Parallelization in f12

MOBasis/HF12.ml

@@ -5,8 +5,6 @@ open Lacaml.D
 module Fis = FourIdxStorage
 
 
-type q = (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Genarray.t 
-
 type t = {
   simulation : Simulation.t ;
   aux_basis  : MOBasis.t ;
@@ -19,54 +17,99 @@ type t = {
 
 let sum l f = List.fold_left (fun accu i -> accu +. f i) 0. l
 
-let array_3_init d1 d2 d3 f =
+let array_3_init d1 d2 d3 fx =
+  let f k = 
+    let x = 
+      Bigarray.(Array2.create Float64 fortran_layout d1 d2)
+    in
+    for j=1 to d2 do
+      for i=1 to d1 do
+        x.{i,j} <-  fx i j k
+      done
+    done;
+    (k,x)
+  in
   let result =
       SharedMemory.create Bigarray.Float64 [| d1 ; d2 ; d3 |]
       |> Bigarray.array3_of_genarray
   in
-  if Parallel.master then 
-    for k=1 to d3 do
+  Util.list_range 1 d3
+  |> Stream.of_list 
+  |> Farm.run ~f
+  |> Stream.iter (fun (k,x) ->
         for j=1 to d2 do
             for i=1 to d1 do
-              result.{i,j,k} <-  f i j k
+              result.{i,j,k} <-  x.{i,j}
             done
-        done
-    done;
-  result
+        done)
+  ;
+  Parallel.broadcast (lazy result)
 
-let array_4_init d1 d2 d3 d4 f =
+let array_4_init d1 d2 d3 d4 fx =
+  let f l = 
+    let x = 
+      Bigarray.(Array3.create Float64 fortran_layout d1 d2 d3)
+    in
+    for k=1 to d3 do
+      for j=1 to d2 do
+        for i=1 to d1 do
+          x.{i,j,k} <-  fx i j k l
+        done
+      done
+    done;
+    (l,x)
+  in
   let result =
       SharedMemory.create Bigarray.Float64 [| d1;d2;d3;d4 |] 
   in
-  if Parallel.master then 
+  Util.list_range 1 d4
+  |> Stream.of_list 
+  |> Farm.run ~f
+  |> Stream.iter (fun (l,x) ->
+        for k=1 to d3 do
+          for j=1 to d2 do
+            for i=1 to d1 do
+              result.{i,j,k,l} <-  x.{i,j,k}
+            done
+          done
+        done)
+  ;
+  Parallel.broadcast (lazy result)
+
+let array_5_init d1 d2 d3 d4 d5 fx =
+  let f m =
+    let x = 
+      Bigarray.(Genarray.create Float64 fortran_layout [| d1; d2; d3; d4 |])
+    in
     for l=1 to d4 do
       for k=1 to d3 do
         for j=1 to d2 do
           for i=1 to d1 do
-            result.{i,j,k,l} <- f i j k l
+            x.{i,j,k,l} <- fx i j k l m
           done
         done
       done
     done;
-  result
-
-let array_5_init d1 d2 d3 d4 d5 f =
+    (m,x)
+  in
   let result =
       SharedMemory.create Bigarray.Float64 [| d1;d2;d3;d4;d5 |]
   in
-  if Parallel.master then 
-    for m=1 to d5 do
+  Util.list_range 1 d5
+  |> Stream.of_list 
+  |> Farm.run ~f
+  |> Stream.iter (fun (m,x) ->
       for l=1 to d4 do
         for k=1 to d3 do
           for j=1 to d2 do
             for i=1 to d1 do
-              result.{i,j,k,l,m} <- f i j k l m
+              result.{i,j,k,l,m} <-  x.{i,j,k,l}
             done
           done
         done
-      done
-    done;
-  result
+      done)
+  ;
+  Parallel.broadcast (lazy result)
 
 
 let make ~simulation ~mo_basis ~aux_basis_filename () =