Add a benchmark using work-stealing deque as a SPMC queue

polytypic · polytypic · commit 228dc654199b · 2024-03-13T09:45:35.000+02:00
This is a technically possible use case of a work-stealing deque and also a
possible, although non-ideal, pattern for some programs running on a
work-stealing scheduler.
diff --git a/bench/bench_ws_deque.ml b/bench/bench_ws_deque.ml
@@ -1,7 +1,7 @@
 open Multicore_bench
 module Ws_deque = Saturn_lockfree.Work_stealing_deque.M
 
-let run_one ~budgetf ?(n_domains = 1) () =
+let run_as_scheduler ~budgetf ?(n_domains = 1) () =
   let spawns =
     Array.init n_domains @@ fun _ -> ref 0 |> Multicore_magic.copy_as_padded
   in
@@ -85,6 +85,83 @@ let run_one ~budgetf ?(n_domains = 1) () =
   let n = Array.fold_left (fun n c -> n + !c) 0 spawns in
   Times.to_thruput_metrics ~n ~singular:"spawn" ~config times
 
+let run_as_one_domain ~budgetf ?(n_msgs = 150 * Util.iter_factor) order =
+  let t = Ws_deque.create () in
+
+  let op_lifo push =
+    if push then Ws_deque.push t 101
+    else match Ws_deque.pop t with _ -> () | exception Exit -> ()
+  and op_fifo push =
+    if push then Ws_deque.push t 101
+    else match Ws_deque.steal t with _ -> () | exception Exit -> ()
+  in
+
+  let init _ =
+    assert (match Ws_deque.steal t with _ -> false | exception Exit -> true);
+    Util.generate_push_and_pop_sequence n_msgs
+  in
+  let work _ bits =
+    Util.Bits.iter (match order with `FIFO -> op_fifo | `LIFO -> op_lifo) bits
+  in
+
+  let config =
+    let label = match order with `FIFO -> "FIFO" | `LIFO -> "LIFO" in
+    Printf.sprintf "one domain (%s)" label
+  in
+  Times.record ~budgetf ~n_domains:1 ~init ~work ()
+  |> Times.to_thruput_metrics ~n:n_msgs ~singular:"message" ~config
+
+let run_as_spmc ~budgetf ~n_thiefs () =
+  let n_domains = n_thiefs + 1 in
+
+  let n_msgs = 70 * Util.iter_factor in
+
+  let t = Ws_deque.create () in
+
+  let n_msgs_to_steal = Atomic.make 0 |> Multicore_magic.copy_as_padded in
+
+  let init _ =
+    assert (match Ws_deque.steal t with _ -> false | exception Exit -> true);
+    Atomic.set n_msgs_to_steal n_msgs
+  in
+  let work i () =
+    if i < n_thiefs then
+      let rec work () =
+        let n = Util.alloc n_msgs_to_steal in
+        if 0 < n then
+          let rec loop n =
+            if 0 < n then
+              match Ws_deque.steal t with
+              | exception Exit ->
+                  Domain.cpu_relax ();
+                  loop n
+              | _ -> loop (n - 1)
+            else work ()
+          in
+          loop n
+      in
+      work ()
+    else
+      for i = 1 to n_msgs do
+        Ws_deque.push t i
+      done
+  in
+
+  let config =
+    Printf.sprintf "1 adder, %d taker%s" n_thiefs
+      (if n_thiefs = 1 then "" else "s")
+  in
+  Times.record ~budgetf ~n_domains ~init ~work ()
+  |> Times.to_thruput_metrics ~n:n_msgs ~singular:"message" ~config
+
 let run_suite ~budgetf =
-  [ 1; 2; 4; 8 ]
-  |> List.concat_map @@ fun n_domains -> run_one ~budgetf ~n_domains ()
+  List.concat
+    [
+      [ 1; 2; 4; 8 ]
+      |> List.concat_map (fun n_domains ->
+             run_as_scheduler ~budgetf ~n_domains ());
+      [ 1; 2; 4 ]
+      |> List.concat_map (fun n_thiefs -> run_as_spmc ~budgetf ~n_thiefs ());
+      run_as_one_domain ~budgetf `FIFO;
+      run_as_one_domain ~budgetf `LIFO;
+    ]
