QCaml/Parallel/Farm.ml

(* Single process function *)
let run_sequential f stream =

  let rec next _ =
    try
      let task = Stream.next stream in
      Some (f task) 
    with Stream.Failure -> None in
    Stream.from next




(* Multi-process functions *)

(* Server side *)
let run_parallel_server stream =


  let fetch_result () : 'a option * int =
    let  result, rank, _tag  =
      Mpi.receive_status  Mpi.any_source  Mpi.any_tag  Mpi.comm_world
    in
    result, rank 
  in


  let send_task (client_rank : int) : unit =
    let task =
      try Some (Stream.next stream)
      with Stream.Failure -> None
    in
    Mpi.send  task  client_rank  0  Mpi.comm_world
  in


  let n_todo = ref (Mpi.comm_size Mpi.comm_world ) in

  let f i = 
    let rec get_result () =
      begin
        match Stream.peek stream with
        | None -> decr n_todo
        | _    -> ()
      end;
      match !n_todo with
      | 0 -> 
        begin
          Mpi.barrier  Mpi.comm_world;
          None
        end
      | _ ->
        begin
          let result, client  = fetch_result () in
          send_task  client;
          match result with
          | None -> get_result ()
          | _    -> result
        end
    in
    get_result ()
  in
  Stream.from f


(** Client side *)
let run_parallel_client f =
  Mpi.send  None  0  0  Mpi.comm_world;

  let rec run () =
    let task = 
      Mpi.receive  0  Mpi.any_tag  Mpi.comm_world
    in

    match task with
    | None -> Mpi.barrier  Mpi.comm_world
    | Some task ->
      let result = f task in
      begin
        Mpi.send  (Some result)  0  0  Mpi.comm_world;
        run ()
      end
  in
  run ();
  Stream.of_list []



let run_parallel f stream =
  match Mpi.comm_rank  Mpi.comm_world with
  | 0 -> run_parallel_server stream
  | _ -> run_parallel_client f


let run f stream =
  match Mpi.comm_size Mpi.comm_world with
  | 1 -> run_sequential f stream
  | _ -> run_parallel   f stream