Parallel integrals with Parmap and Zmq

This commit is contained in:
Anthony Scemama 2018-06-01 14:34:28 +02:00
parent 63f0b379de
commit 7c7e79285a
3 changed files with 113 additions and 33 deletions

View File

@ -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

View File

@ -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

2
_tags
View File

@ -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)
<odoc-ltxhtml>: not_hygienic