Revert to using our own SPSC unlimited FIFO queue

lukstafi · lukstafi · commit ccf03e05da77 · 2024-09-09T15:25:57.000+02:00
diff --git a/arrayjit.opam b/arrayjit.opam
@@ -24,7 +24,6 @@ depends: [
   "ocannl_npy"
   "stdio"
   "num"
-  "saturn_lockfree"
   "ppxlib"
   "ppx_jane"
   "ppx_expect"
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -166,17 +166,13 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
   type task_list = Tnode.task Utils.mutable_list [@@deriving sexp_of]
 
   module Mut = Stdlib.Mutex
-  module Queue = Saturn_lockfree.Single_prod_single_cons_queue
-
-  type task_queue = Tnode.task Queue.t
-
-  let sexp_of_task_queue q =
-    Sexp.(List [ Atom "task_queue_of_size"; Atom (Int.to_string @@ Queue.size q) ])
 
   type device_state = {
     mutable keep_spinning : bool;
     mutable device_error : exn option;
-    queue : task_queue;
+    mutable host_pos : task_list;
+    mutable dev_pos : task_list;
+    mutable dev_previous_pos : task_list;
     mut : (Mut.t[@sexp.opaque]);
     host_wait_for_idle : (Stdlib.Condition.t[@sexp.opaque]);
     dev_wait_for_work : (Stdlib.Condition.t[@sexp.opaque]);
@@ -205,7 +201,7 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
 
   let expected_merge_node (code : code) = Backend.expected_merge_node code
   let expected_merge_nodes (codes : code_batch) = Backend.expected_merge_nodes codes
-  let is_dev_queue_empty state = Queue.size state.queue = 0
+  let is_dev_queue_empty state = Utils.(is_empty @@ tl_exn state.dev_previous_pos)
   let is_idle device = is_dev_queue_empty device.state && device.state.is_ready
   let name = "multicore " ^ Backend.name
 
@@ -231,9 +227,7 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
     Option.iter d.device_error ~f:(fun e ->
         Exn.reraise e @@ name ^ " device " ^ Int.to_string device.ordinal);
     if not d.keep_spinning then invalid_arg "Multicore_backend: device not available";
-    if not @@ Queue.try_push d.queue task then (
-      await device;
-      Queue.push_exn d.queue task);
+    d.host_pos <- Utils.insert ~next:task d.host_pos;
     if d.is_ready then (
       Mut.lock d.mut;
       Stdlib.Condition.broadcast d.dev_wait_for_work;
@@ -243,11 +237,22 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
 
   let%track3_l_sexp spinup_device ~(ordinal : int) : device =
     Int.incr global_run_no;
+    let init_pos =
+      Utils.Cons
+        {
+          hd =
+            Tnode.Task
+              { context_lifetime = (); description = "root of task queue"; work = (fun () -> ()) };
+          tl = Empty;
+        }
+    in
     let state =
       {
         keep_spinning = true;
         device_error = None;
-        queue = Queue.create ~size_exponent:12;
+        host_pos = init_pos;
+        dev_pos = Empty;
+        dev_previous_pos = init_pos;
         mut = Mut.create ();
         is_ready = false;
         host_wait_for_idle = Stdlib.Condition.create ();
@@ -259,8 +264,8 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
       assert (not @@ Domain.is_main_domain ());
       try
         while state.keep_spinning do
-          match Queue.pop_opt state.queue with
-          | None ->
+          match state.dev_pos with
+          | Empty ->
               Mut.lock state.mut;
               if is_dev_queue_empty state && state.keep_spinning then (
                 state.is_ready <- true;
@@ -269,8 +274,12 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
                   Stdlib.Condition.wait state.dev_wait_for_work state.mut
                 done;
                 state.is_ready <- false);
+              state.dev_pos <- Utils.tl_exn state.dev_previous_pos;
               Mut.unlock state.mut
-          | Some task -> Tnode.run task
+          | Cons { hd; tl } ->
+              Tnode.run hd;
+              state.dev_previous_pos <- state.dev_pos;
+              state.dev_pos <- tl
         done
       with e ->
         state.device_error <- Some e;
@@ -594,12 +603,11 @@ module Pipes_multicore_backend (Backend : No_device_backend) : Backend = struct
           | Empty ->
               let _host_released : bool = host_wait_for_idle.release_if_waiting () in
               let _could_wait : bool = wait_by_dev () in
-              (* not _host_released && not _could_wait: we busy-loop until host processes its release. *)
-              (* [%log "WORK WHILE LOOP: EMPTY AFTER WAIT -- dev pos:", (state.dev_pos : task_list)]; *)
+              (* not _host_released && not _could_wait: we busy-loop until host processes its
+                 release. *)
               state.dev_pos <- Utils.tl_exn state.dev_previous_pos
           | Cons { hd; tl } ->
               Tnode.run hd;
-              (* [%log "WORK WHILE LOOP: AFTER WORK"]; *)
               state.dev_previous_pos <- state.dev_pos;
               state.dev_pos <- tl
         done
diff --git a/arrayjit/lib/dune b/arrayjit/lib/dune
@@ -11,7 +11,6 @@
   str
   ctypes
   ctypes.foreign
-  saturn_lockfree
   (select
    gcc_backend.ml
    from
diff --git a/dune-project b/dune-project
@@ -65,7 +65,6 @@
   ocannl_npy
   stdio
   num
-  saturn_lockfree
   ppxlib
   ppx_jane
   ppx_expect
