Yay, fix the scheduler bug: the old d.is_idle check in await ignored the queue

lukstafi · lukstafi · commit 418af7c1617c · 2024-09-09T15:00:02.000+02:00
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -180,7 +180,7 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
     mut : (Mut.t[@sexp.opaque]);
     host_wait_for_idle : (Stdlib.Condition.t[@sexp.opaque]);
     dev_wait_for_work : (Stdlib.Condition.t[@sexp.opaque]);
-    mutable is_idle : bool;
+    mutable is_ready : bool;
     mutable host_is_waiting : bool;  (** The host is waiting for this specific device. *)
   }
   [@@deriving sexp_of]
@@ -206,17 +206,17 @@ 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_idle device = is_dev_queue_empty device.state && device.state.is_idle
+  let is_idle device = is_dev_queue_empty device.state && device.state.is_ready
   let name = "multicore " ^ Backend.name
 
   let%track3_l_sexp await device =
     assert (Domain.is_main_domain ());
     let d = device.state in
-    if (not d.is_idle) && d.keep_spinning then (
+    if (not @@ is_idle device) && d.keep_spinning then (
       Mut.lock d.mut;
-      if (not d.is_idle) && d.keep_spinning then (
+      if (not @@ is_idle device) && d.keep_spinning then (
         d.host_is_waiting <- true;
-        while (not d.is_idle) && d.keep_spinning do
+        while (not @@ is_idle device) && d.keep_spinning do
           Stdlib.Condition.wait d.host_wait_for_idle d.mut
         done;
         d.host_is_waiting <- false);
@@ -234,7 +234,7 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
     if not @@ Queue.try_push d.queue task then (
       await device;
       Queue.push_exn d.queue task);
-    if d.is_idle then (
+    if d.is_ready then (
       Mut.lock d.mut;
       Stdlib.Condition.broadcast d.dev_wait_for_work;
       Mut.unlock d.mut)
@@ -249,25 +249,26 @@ module Multicore_backend (Backend : No_device_backend) : Backend = struct
         device_error = None;
         queue = Queue.create ~size_exponent:12;
         mut = Mut.create ();
-        is_idle = false;
+        is_ready = false;
         host_wait_for_idle = Stdlib.Condition.create ();
         dev_wait_for_work = Stdlib.Condition.create ();
         host_is_waiting = false;
       }
     in
     let%track3_l_sexp worker (() : unit) : unit =
+      assert (not @@ Domain.is_main_domain ());
       try
         while state.keep_spinning do
           match Queue.pop_opt state.queue with
           | None ->
               Mut.lock state.mut;
               if is_dev_queue_empty state && state.keep_spinning then (
-                state.is_idle <- true;
+                state.is_ready <- true;
                 while is_dev_queue_empty state && state.keep_spinning do
                   if state.host_is_waiting then Stdlib.Condition.broadcast state.host_wait_for_idle;
                   Stdlib.Condition.wait state.dev_wait_for_work state.mut
                 done;
-                state.is_idle <- false);
+                state.is_ready <- false);
               Mut.unlock state.mut
           | Some task -> Tnode.run task
         done
diff --git a/bin/moons_benchmark.ml b/bin/moons_benchmark.ml
@@ -40,9 +40,9 @@ let classify_moons ~seed ~on_device ~inlining_cutoff ~num_devices ~batch_size ~b
   (* let hid_dim = 4 in *)
   let len = batch_size * 20 in
   (* let epochs = 100 in *)
-  (* let epochs = 20 in *)
-  (* let epochs = 5 in *)
-  let epochs = 1 in
+  let epochs = 20 in
+  (* let epochs = 10 in *)
+  (* let epochs = 1 in *)
   let init_lr = 0.1 in
   let noise () = Rand.float_range (-0.1) 0.1 in
   let moons_flat =
@@ -177,7 +177,7 @@ let _cpu_benchmarks =
 
 let cuda_benchmarks =
   List.concat_map [ 0; (* 1; 2; *) 3 ] ~f:(fun inlining_cutoff ->
-      List.concat_map [ 1; 2; 5; 8; 10; 16; 20; 30 (* *; 32; 40; 64 *) ] ~f:(fun num_devices ->
+      List.concat_map [ 1; 2; 5; 8; 10 (* ; 16; 20; 30; 32; 40; 64 *) ] ~f:(fun num_devices ->
           List.concat_map [ 120; 160 (* ; 320; 640; 1280 *) ] ~f:(fun batch_size ->
               List.concat_map [ 0; 1 (* ; 2; 3; 4 *) ] ~f:(fun seed ->
                   List.concat_map [ (* "gccjit" ; *) "cc" (* ; "cuda" *) ]
diff --git a/lib/train.ml b/lib/train.ml
@@ -508,7 +508,7 @@ let example_train_loop ?(disable_rootness_check = false) ~seed ~batch_size ~init
     assert (Backend.from_host routine.context infer.value);
     run routine;
     assert (Backend.to_host routine.context model_result.value);
-    Backend.(await @@ get_ctx_device prior_contexts.(0));
+    Backend.(await @@ get_ctx_device routine.context);
     Tensor.get_values model_result
   in
   (* Note: infer_callback is significantly less efficient than using the model via arrayjit. *)
