Fixed MPI Broadcast in FCI-f12

MOBasis/HF12.ml

@@ -47,31 +47,31 @@ let array_3_init d1 d2 d3 fx =
   Parallel.Node.broadcast (lazy result)
 
 let array_4_init d1 d2 d3 d4 fx =
-  let f l = 
+  let f (k,l) = 
     let x = 
-      Bigarray.(Array3.create Float64 fortran_layout d1 d2 d3)
+      Bigarray.(Array2.create Float64 fortran_layout d1 d2)
     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
+    for j=1 to d2 do
+      for i=1 to d1 do
+        x.{i,j} <-  fx i j k l
       done
     done;
-    (l,x)
+    (k,l,x)
   in
   let result =
       SharedMemory.create Bigarray.Float64 [| d1;d2;d3;d4 |] 
   in
   Util.list_range 1 d4
+  |> List.map (fun l -> 
+    Util.list_range 1 d3
+    |> List.map (fun k -> (k,l)) )
+  |> List.concat
   |> 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
+  |> Stream.iter (fun (k,l,x) ->
+        for j=1 to d2 do
+          for i=1 to d1 do
+            result.{i,j,k,l} <-  x.{i,j}
           done
         done)
   ;
@@ -79,40 +79,77 @@ let array_4_init d1 d2 d3 d4 fx =
   Parallel.Node.broadcast (lazy result)
 
 let array_5_init d1 d2 d3 d4 d5 fx =
-  let f m =
+  let f (l,m) =
     let x = 
-      Bigarray.(Genarray.create Float64 fortran_layout [| d1; d2; d3; d4 |])
+      Bigarray.(Array3.create Float64 fortran_layout d1 d2 d3)
     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
-            x.{i,j,k,l} <- fx i j k l m
-          done
+    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 m
         done
       done
     done;
-    (m,x)
+    (l,m,x)
   in
   let result =
       SharedMemory.create Bigarray.Float64 [| d1;d2;d3;d4;d5 |]
   in
   Util.list_range 1 d5
+  |> List.map (fun m -> 
+    Util.list_range 1 d4
+    |> List.map (fun l -> (l,m)) )
+  |> List.concat
   |> 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} <-  x.{i,j,k,l}
-            done
+  |> Stream.iter (fun (l,m,x) ->
+      for k=1 to d3 do
+        for j=1 to d2 do
+          for i=1 to d1 do
+            result.{i,j,k,l,m} <-  x.{i,j,k}
           done
         done
       done)
   ;
   if Parallel.master then Printf.printf "Broadcast d5\n" ;
-  Parallel.Node.broadcast (lazy result)
+  try
+    Parallel.Node.broadcast (lazy result)
+  with Invalid_argument _ ->
+    begin
+      Printf.eprintf "Array too large... splitting.\n%!";
+      let x = 
+        Bigarray.(Genarray.create Float64 fortran_layout [| d1; d2; d3; d4 |])
+      in
+      for m=1 to d5 do
+        if Parallel.master then
+          begin
+            for l=1 to d4 do
+              for k=1 to d3 do
+                for j=1 to d2 do
+                  for i=1 to d1 do
+                    x.{i,j,k,l} <- result.{i,j,k,l,m} 
+                  done
+                done
+              done
+            done;
+            ignore @@ Parallel.Node.broadcast (lazy x)
+          end
+        else 
+          begin
+            ignore @@ Parallel.Node.broadcast (lazy 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} <- x.{i,j,k,l} 
+                  done
+                done
+              done
+            done
+          end
+      done;
+      result
+    end
 
 
 let make ~simulation ~mo_basis ~aux_basis_filename () =
@@ -938,8 +975,7 @@ let make ~simulation ~mo_basis ~aux_basis_filename () =
       f_0 ; f_1 ; f_2 ; f_3 
     } 
   in
-  if Parallel.master then Printf.printf "Broadcast f3\n" ;
-  Parallel.broadcast (lazy result)
+  result