Broadcast for large arrays

Parallel_mpi/Parallel.ml

@@ -23,12 +23,54 @@ let barrier () =
   Mpi.barrier Mpi.comm_world
 
 
+let broadcast_bytes data root comm =
+  let length = Bytes.length data in
+  let lmax = Int32.(to_int max_int) in
+  if length < lmax then
+    Mpi.broadcast data root comm
+    |> ignore
+  else
+    let rec aux start =
+      if length - start < lmax then (
+        let message = Bytes.sub data start (length - start) in
+        Mpi.broadcast message root comm
+        |> ignore;
+        Bytes.blit message 0 data start lmax;
+      ) else (
+        let message = Bytes.sub data start lmax in
+        Mpi.broadcast message root comm
+        |> ignore;
+        Bytes.blit message 0 data start lmax;
+        aux (start + lmax)
+      )
+    in
+    aux 0 
+
+
+let broadcast v root comm =
+  if rank = root then
+    begin
+      let data = Marshal.to_bytes v [Marshal.Closures] in
+      ignore(Mpi.broadcast_int (Bytes.length data) root comm);
+      broadcast_bytes data root comm;
+      v
+    end
+  else
+    begin
+      (* Other processes receive length, allocate buffer, receive data,
+             and unmarshal it. *)
+      let len = Mpi.broadcast_int 0 root comm in
+      let data = Bytes.create len in
+      broadcast_bytes data root comm;
+      Marshal.from_bytes data 0
+    end
+
 let broadcast x =
   let x = 
     if master then Some (Lazy.force x)
     else None
   in
-  match Mpi.broadcast x 0 Mpi.comm_world with
+  match broadcast x 0 Mpi.comm_world with
   | Some x -> x
   | None -> assert false