mirror of
https://gitlab.com/scemama/QCaml.git
synced 2024-08-31 00:33:40 +02:00
142 lines
3.1 KiB
OCaml
142 lines
3.1 KiB
OCaml
(* 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 *)
|
|
|
|
type task_id = int
|
|
|
|
|
|
(* Server side *)
|
|
let run_parallel_server ~ordered stream =
|
|
|
|
|
|
let fetch_result () : (task_id * 'a) option * int =
|
|
let (message, rank, _tag) : (task_id * 'a) option * int * int =
|
|
Mpi.receive_status Mpi.any_source Mpi.any_tag Mpi.comm_world
|
|
in
|
|
message, rank
|
|
in
|
|
|
|
|
|
let send_task (client_rank : int) : unit =
|
|
let task =
|
|
try
|
|
let task_id = Stream.count stream in
|
|
let element = Stream.next stream in
|
|
Some (task_id, element)
|
|
with Stream.Failure -> None
|
|
in
|
|
Mpi.send task client_rank 0 Mpi.comm_world
|
|
in
|
|
|
|
|
|
(* n_todo is required for clean termination. It is the
|
|
number of tasks to wait for when the input stream is
|
|
empty. *)
|
|
let n_todo = ref (Mpi.comm_size Mpi.comm_world ) in
|
|
|
|
let rec get_result () : (task_id * 'a ) option =
|
|
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 message, rank = fetch_result () in
|
|
send_task rank;
|
|
match message with
|
|
| None -> get_result ()
|
|
| Some (task_id, result) -> Some (task_id, result)
|
|
end
|
|
in
|
|
|
|
let f =
|
|
|
|
if ordered then
|
|
|
|
(* buffer of finished tasks with a task_id greater than the
|
|
current result_id. It allows to put back the results in
|
|
the correct order.
|
|
*)
|
|
let buffer = Hashtbl.create 67 in
|
|
|
|
fun i ->
|
|
begin
|
|
match Hashtbl.find_opt buffer i with
|
|
| Some x ->
|
|
begin
|
|
Hashtbl.remove buffer i;
|
|
Some x
|
|
end
|
|
| None ->
|
|
let rec loop () =
|
|
match get_result () with
|
|
| None -> None
|
|
| Some (task_id, result) ->
|
|
if task_id = i then Some result
|
|
else (Hashtbl.add buffer task_id result; loop () )
|
|
in loop ()
|
|
end
|
|
|
|
else
|
|
|
|
fun _ ->
|
|
match get_result () with
|
|
| Some (_, result) -> Some result
|
|
| None -> None
|
|
|
|
in
|
|
Stream.from f
|
|
|
|
|
|
(** Client side *)
|
|
let run_parallel_client f =
|
|
Mpi.send None 0 0 Mpi.comm_world;
|
|
|
|
let rec run () =
|
|
let message =
|
|
Mpi.receive 0 Mpi.any_tag Mpi.comm_world
|
|
in
|
|
|
|
match message with
|
|
| None -> Mpi.barrier Mpi.comm_world
|
|
| Some (task_id, task) ->
|
|
let result = f task in
|
|
begin
|
|
Mpi.send (Some (task_id, result)) 0 0 Mpi.comm_world;
|
|
run ()
|
|
end
|
|
in
|
|
run ();
|
|
Stream.of_list []
|
|
|
|
|
|
|
|
let run_parallel ~ordered f stream =
|
|
match Mpi.comm_rank Mpi.comm_world with
|
|
| 0 -> run_parallel_server ~ordered stream
|
|
| _ -> run_parallel_client f
|
|
|
|
|
|
let run ?(ordered=true) ~f stream =
|
|
match Mpi.comm_size Mpi.comm_world with
|
|
| 1 -> run_sequential f stream
|
|
| _ -> run_parallel ~ordered f stream
|
|
|
|
|