From 7c7e79285a513a6af5c20a6d5f413751637b6b98 Mon Sep 17 00:00:00 2001 From: Anthony Scemama Date: Fri, 1 Jun 2018 14:34:28 +0200 Subject: [PATCH] Parallel integrals with Parmap and Zmq --- Basis/ERI.ml | 136 +++++++++++++++++++++++++++++++++++++++------------ INSTALL.md | 8 +++ _tags | 2 +- 3 files changed, 113 insertions(+), 33 deletions(-) diff --git a/Basis/ERI.ml b/Basis/ERI.ml index 2282246..13ccff8 100644 --- a/Basis/ERI.ml +++ b/Basis/ERI.ml @@ -98,7 +98,7 @@ let filter_contracted_shell_pair_couples ?(cutoff=integrals_cutoff) shell_pair_c *) -let store_class ?(cutoff=integrals_cutoff) data contracted_shell_pair_couple cls = +let store_class ?(cutoff=integrals_cutoff) push_socket contracted_shell_pair_couple cls = let to_powers x = let open Zkey in match to_powers x with @@ -110,6 +110,7 @@ let store_class ?(cutoff=integrals_cutoff) data contracted_shell_pair_couple cls and shell_q = Cspc.shell_pair_q contracted_shell_pair_couple in + let msg = ref [] in Array.iteri (fun i_c powers_i -> let i_c = Cs.index (Csp.shell_a shell_p) + i_c + 1 in let xi = to_powers powers_i in @@ -124,11 +125,12 @@ let store_class ?(cutoff=integrals_cutoff) data contracted_shell_pair_couple cls let xl = to_powers powers_l in let key = Zkey.of_powers_twelve xi xj xk xl in let value = Zmap.find cls key in - set_chem data i_c j_c k_c l_c value; + msg := (i_c,j_c,k_c,l_c,value) :: !msg; ) (Cs.zkey_array (Csp.shell_b shell_q)) ) (Cs.zkey_array (Csp.shell_a shell_q)) ) (Cs.zkey_array (Csp.shell_b shell_p)) - ) (Cs.zkey_array (Csp.shell_a shell_p)) + ) (Cs.zkey_array (Csp.shell_a shell_p)); + Zmq.Socket.send_all push_socket ["0" ; Bytes.to_string (Marshal.to_bytes !msg []) ] @@ -162,38 +164,108 @@ let of_basis basis = let t0 = Unix.gettimeofday () in let ishell = ref 0 in - List.iter (fun shell_p -> - let () = - if (Cs.index (Csp.shell_a shell_p) > !ishell) then - (ishell := Cs.index (Csp.shell_a shell_p) ; print_int !ishell ; print_newline ()) - in - let sp = - Csp.shell_pairs shell_p - in - try - List.iter (fun shell_q -> - let () = - if Cs.index (Csp.shell_a shell_q) > - Cs.index (Csp.shell_a shell_p) then - raise Exit - in - let sq = Csp.shell_pairs shell_q in - let cspc = - if Array.length sp < Array.length sq then - Cspc.make ~cutoff shell_p shell_q - else - Cspc.make ~cutoff shell_q shell_p - in + let zmq_addr = Printf.sprintf "ipc://%d" (Unix.getpid ()) in + let () = + match Unix.fork () with + | 0 -> begin + let zmq = ref None in - match cspc with - | Some cspc -> let cls = class_of_contracted_shell_pair_couple cspc in - store_class ~cutoff eri_array cspc cls - | None -> () - ) shell_pairs - with Exit -> () - ) shell_pairs ; + Parmap.pariter ~chunksize:1 ~ncores:4 + ~init:(fun _ -> + let zmq_context = + Zmq.Context.create () + in + let push_socket = + Zmq.Socket.create zmq_context Zmq.Socket.push + in + Zmq.Socket.connect push_socket zmq_addr; + zmq := Some (zmq_context, push_socket) + ) + + (fun shell_p -> + let push_socket = + match !zmq with + | Some (_, push_socket) -> push_socket + | None -> failwith "ZMQ" + in + let () = + if (Cs.index (Csp.shell_a shell_p) > !ishell) then + (ishell := Cs.index (Csp.shell_a shell_p) ; print_int !ishell ; print_newline ()) + in + + let sp = + Csp.shell_pairs shell_p + in + + try + List.iter (fun shell_q -> + let () = + if Cs.index (Csp.shell_a shell_q) > + Cs.index (Csp.shell_a shell_p) then + raise Exit + in + let sq = Csp.shell_pairs shell_q in + let cspc = + if Array.length sp < Array.length sq then + Cspc.make ~cutoff shell_p shell_q + else + Cspc.make ~cutoff shell_q shell_p + in + + match cspc with + | Some cspc -> let cls = class_of_contracted_shell_pair_couple cspc in + store_class ~cutoff push_socket cspc cls + | None -> () + ) shell_pairs + with Exit -> () + ) (Parmap.L shell_pairs) + ~finalize:(fun _ -> + let zmq_context, push_socket = + match !zmq with + | Some (zmq_context, push_socket) -> zmq_context, push_socket + | None -> failwith "ZMQ" + in + Zmq.Socket.close push_socket; + Zmq.Context.terminate zmq_context + ); + let zmq_context = + Zmq.Context.create () + in + let push_socket = Zmq.Socket.create zmq_context Zmq.Socket.push in + Zmq.Socket.connect push_socket zmq_addr; + Zmq.Socket.send_all push_socket [ "1" ; ""]; + Zmq.Socket.close push_socket; + Zmq.Context.terminate zmq_context; + ignore @@ exit 0 + end + + | pid -> begin + let zmq_context = + Zmq.Context.create () + in + let pull_socket = + Zmq.Socket.create zmq_context Zmq.Socket.pull + in + Zmq.Socket.bind pull_socket zmq_addr; + + try + while true do + match Zmq.Socket.recv_all pull_socket with + | "0" :: rest :: [] -> + List.iter (fun (i,j,k,l,value) -> + set_chem eri_array i j k l value) (Marshal.from_bytes (Bytes.of_string rest) 0) + | "1" :: _ -> raise Exit + | _ -> invalid_arg "ERI" + done + with Exit -> (); + + Zmq.Socket.close pull_socket; + Zmq.Context.terminate zmq_context; + ignore (Unix.wait ()) + end + in Printf.printf "Computed ERIs in %f seconds\n%!" (Unix.gettimeofday () -. t0); eri_array diff --git a/INSTALL.md b/INSTALL.md index 014a5f8..351ef28 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -21,6 +21,14 @@ export LACAML_LIBS="-L${MKLROOT}/lib/intel64 -Wl,--no-as-needed -lmkl_rt -lpthr opam install lacaml ``` +# Parmap + +Multicore library. + +```bash +opam install parmap +``` + # odoc-ltxhtml diff --git a/_tags b/_tags index dca8492..9568e5b 100644 --- a/_tags +++ b/_tags @@ -1,4 +1,4 @@ -true: package(str,unix,bigarray,lacaml) +true: package(str,unix,bigarray,lacaml,parmap,zmq) <*.byte> : linkdep(Utils/math_functions.o), custom <*.native>: linkdep(Utils/math_functions.o) : not_hygienic