Parallel integrals with Parmap and Zmq

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
+  let zmq_addr = Printf.sprintf "ipc://%d" (Unix.getpid ()) in
-        List.iter (fun shell_q ->
+  let () = 
-            let () = 
+    match Unix.fork () with
-              if  Cs.index (Csp.shell_a shell_q) >
+    | 0 ->    begin
-                  Cs.index (Csp.shell_a shell_p) then
+                let zmq = ref None in
-                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
+                Parmap.pariter ~chunksize:1 ~ncores:4
-            | Some cspc -> let cls = class_of_contracted_shell_pair_couple cspc in
+                  ~init:(fun _ ->
-                           store_class ~cutoff eri_array cspc cls
+                    let zmq_context = 
-            | None -> ()
+                      Zmq.Context.create ()
-          ) shell_pairs
+                    in
-      with Exit -> ()
+                    let push_socket = 
-    ) shell_pairs ;
+                      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
 

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
 

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