Split schedulers.ml out of backends.ml

lukstafi · lukstafi · commit e34f941340d0 · 2024-11-01T12:01:53.000+01:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,4 +1,4 @@
-## [0.4.2] -- current
+## [0.5.0] -- current
 
 ### Added
 
@@ -11,7 +11,7 @@
 - Migrated to cudajit 0.5.
 - Verifying that code is linked with the right contexts, by tracking `embedded_nodes` with assignments.
 - Renaming: (virtual) `device` -> `stream`, `physical_device` -> `device`.
-- New files: split out `backend_types.ml` from `backends.ml`; moved `Tnode.task` to `task.ml`; renamed `backend_utils.ml` to `c_syntax.ml`.
+- New files: split out `backend_intf.ml`, `backend_impl.ml`, `schedulers.ml` from `backends.ml`; moved `Tnode.task` to `task.ml`; renamed `backend_utils.ml` to `c_syntax.ml`.
 - Removed half-static verification of merge buffer nodes inside `device_to_device`.
 - Fixed #286: cross-stream-sharing incorporated into `Tnode.memory_mode`.
 - Moved the multicore backend from a `device = stream` model to a single device model.
diff --git a/arrayjit/lib/backend_impl.ml b/arrayjit/lib/backend_impl.ml
@@ -291,3 +291,17 @@ module type Lowered_backend = sig
       batch. The returned [ctx_arrays] will be part of a context downstream of all the tasks and the
       tasks' contexts are not independent (typically, they are cumulative). *)
 end
+
+module Alloc_buffer_ignore_stream
+    (Device_types : Device_types)
+    (Backend : Alloc_buffer with type buffer_ptr = Device_types.buffer_ptr and type stream := unit) :
+  Alloc_buffer with type buffer_ptr = Backend.buffer_ptr and type stream = Device_types.stream =
+struct
+  include Device_types
+
+  let alloc_buffer ?old_buffer ~size_in_bytes _stream =
+    Backend.alloc_buffer ?old_buffer ~size_in_bytes ()
+
+  let alloc_zero_init_array prec ~dims _stream = Backend.alloc_zero_init_array prec ~dims ()
+  let free_buffer = Option.map Backend.free_buffer ~f:(fun memfree _stream ptr -> memfree () ptr)
+end
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -87,249 +87,6 @@ module Add_buffer_retrieval_and_syncing (Backend : No_buffer_retrieval_or_syncin
               true)
 end
 
-module Alloc_buffer_ignore_stream
-    (Device_types : Device_types)
-    (Backend : Alloc_buffer with type buffer_ptr = Device_types.buffer_ptr and type stream := unit) :
-  Alloc_buffer with type buffer_ptr = Backend.buffer_ptr and type stream = Device_types.stream =
-struct
-  include Device_types
-
-  let alloc_buffer ?old_buffer ~size_in_bytes _stream =
-    Backend.alloc_buffer ?old_buffer ~size_in_bytes ()
-
-  let alloc_zero_init_array prec ~dims _stream = Backend.alloc_zero_init_array prec ~dims ()
-  let free_buffer = Option.map Backend.free_buffer ~f:(fun memfree _stream ptr -> memfree () ptr)
-end
-
-module Multicore_scheduler (Backend : For_add_scheduler) :
-  With_scheduler with type buffer_ptr = Backend.buffer_ptr = struct
-  include Backend
-  module Domain = Domain [@warning "-3"]
-
-  let global_config = ref Only_devices_parallel
-
-  let initialize config =
-    global_config := config;
-    initialize config
-
-  let is_initialized = is_initialized
-
-  type task_list = Task.t Utils.mutable_list [@@deriving sexp_of]
-
-  module Mut = Stdlib.Mutex
-  module Queue = Saturn_lockfree.Single_prod_single_cons_queue
-
-  type task_queue = Task.t Queue.t
-
-  let sexp_of_task_queue q =
-    Sexp.(List [ Atom "task_queue_of_size"; Atom (Int.to_string @@ Queue.size q) ])
-
-  module Device_config = struct
-    include (
-      Backend : Buffer with type buffer_ptr = Backend.buffer_ptr and type buffer = Backend.buffer)
-
-    type dev = CPU [@@deriving sexp_of]
-
-    type stream_state = {
-      mutable keep_spinning : bool;
-      mutable stream_error : exn option;
-      queue : task_queue;
-      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_ready : bool;
-    }
-    [@@deriving sexp_of]
-
-    type domain = unit Domain.t
-
-    let sexp_of_domain (d : domain) = Sexp.Atom ("domain-" ^ Int.to_string (Domain.get_id d :> int))
-
-    type runner = { state : stream_state; domain : domain } [@@deriving sexp_of]
-    type event = Not_implemented_yet [@@deriving sexp_of]
-
-    let name = "multicore_" ^ Backend.name
-  end
-
-  module Device_types = Device_types (Device_config)
-  include Device (Device_types) (Alloc_buffer_ignore_stream (Device_types) (Backend))
-  open Device_config
-
-  (** TODO: Blocks till the event completes, if it's not done already. *)
-  let sync Not_implemented_yet = ()
-
-  (** TODO: Whether the event completed. *)
-  let is_done Not_implemented_yet = true
-
-  (** TODO: Schedules waiting for the given event on the context's stream. *)
-  let will_wait_for _ctx Not_implemented_yet = ()
-
-  let get_used_memory _device = get_used_memory ()
-  let is_dev_queue_empty state = Queue.size state.queue = 0
-  let is_idle stream = is_dev_queue_empty stream.runner.state && stream.runner.state.is_ready
-  let name = "multicore_" ^ name
-  let get_name stream = [%string "%{name}:0:%{stream.stream_id#Int}"]
-
-  let%track3_l_sexp await stream =
-    assert (Domain.is_main_domain ());
-    let d = stream.runner.state in
-    if (not @@ is_idle stream) && d.keep_spinning then (
-      Mut.lock d.mut;
-      while (not @@ is_idle stream) && d.keep_spinning do
-        (* If the stream "is ready", it needs to be woken up first to finish the work. *)
-        if d.is_ready then Stdlib.Condition.broadcast d.dev_wait_for_work;
-        Stdlib.Condition.wait d.host_wait_for_idle d.mut
-      done;
-      Mut.unlock d.mut;
-      Option.iter d.stream_error ~f:(fun e -> Exn.reraise e @@ get_name stream))
-
-  (** TODO: Returns the event indicating if any currently running or scheduled computations on the
-      stream have completed. *)
-  let all_work _stream = Not_implemented_yet
-
-  let%track3_l_sexp schedule_task stream task =
-    assert (Domain.is_main_domain ());
-    [%log_result "schedule_task", Task.describe task, get_name stream];
-    let d = stream.runner.state in
-    Option.iter d.stream_error ~f:(fun e -> Exn.reraise e @@ get_name stream);
-    if not d.keep_spinning then invalid_arg "Multicore_scheduler: stream not available";
-    if not @@ Queue.try_push d.queue task then (
-      await stream;
-      Queue.push_exn d.queue task);
-    if d.is_ready then (
-      Mut.lock d.mut;
-      Stdlib.Condition.broadcast d.dev_wait_for_work;
-      Mut.unlock d.mut)
-
-  let global_run_no = ref 0
-  let device : device = make_device CPU ~ordinal:0
-
-  let%track3_l_sexp spinup_stream ~stream_id : stream =
-    Int.incr global_run_no;
-    let state =
-      {
-        keep_spinning = true;
-        stream_error = None;
-        queue = Queue.create ~size_exponent:12;
-        mut = Mut.create ();
-        is_ready = false;
-        host_wait_for_idle = Stdlib.Condition.create ();
-        dev_wait_for_work = Stdlib.Condition.create ();
-      }
-    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;
-              state.is_ready <- true;
-              Stdlib.Condition.broadcast state.host_wait_for_idle;
-              while is_dev_queue_empty state && state.keep_spinning do
-                Stdlib.Condition.wait state.dev_wait_for_work state.mut
-              done;
-              state.is_ready <- false;
-              Mut.unlock state.mut
-          | Some task -> Task.run task
-        done
-      with e ->
-        state.stream_error <- Some e;
-        state.keep_spinning <- false;
-        [%log1 "stream", (stream_id : int), "exception", Exn.to_string e];
-        (* TODO: we risk raising this error multiple times because await and schedule_task raise
-           stream_error. But this is fine if we assume all exceptions are fatal. *)
-        raise e
-    in
-    make_stream device { state; domain = Domain.spawn worker } ~stream_id
-
-  module Dynarr = Stdlib.Dynarray
-
-  let num_devices () = 1
-  let suggested_num_streams _device = Domain.recommended_domain_count () - 1
-
-  let cleanup_stream stream =
-    assert (Domain.is_main_domain ());
-    await stream;
-    let r = stream.runner in
-    r.state.keep_spinning <- false;
-    Stdlib.Condition.broadcast r.state.dev_wait_for_work;
-    Domain.join r.domain
-
-  let get_device ~ordinal =
-    if ordinal <> 0 then
-      invalid_arg [%string "Multicore_scheduler.get_device %{ordinal#Int}: only device 0 exists"];
-    device
-
-  let latest_stream_id = ref (-1)
-
-  let new_stream _device =
-    assert (Domain.is_main_domain ());
-    Int.incr latest_stream_id;
-    let stream = spinup_stream ~stream_id:!latest_stream_id in
-    Stdlib.Gc.finalise cleanup_stream stream;
-    stream
-end
-
-(** For debugging, allow [Sync_scheduler(...).suggested_num_streams] calls to return >1 numbers. *)
-let sync_suggested_num_streams = ref 1
-
-(** A minimalisitc wrapper creating backends where all calls run synchronously on the main thread.
-    There is only one device, but an arbitrary number of streams. *)
-module Sync_scheduler (Backend : For_add_scheduler) = struct
-  include Backend
-
-  module Device_config = struct
-    include (
-      Backend : Buffer with type buffer_ptr = Backend.buffer_ptr and type buffer = Backend.buffer)
-
-    type dev = CPU [@@deriving sexp_of]
-    type runner = unit [@@deriving sexp_of]
-    type event = unit [@@deriving sexp_of]
-
-    let name = "sync_" ^ Backend.name
-  end
-
-  module Device_types = Device_types (Device_config)
-  include Device (Device_types) (Alloc_buffer_ignore_stream (Device_types) (Backend))
-  open Device_config
-
-  let sync () = ()
-  let is_done () = true
-  let will_wait_for _context () = ()
-
-  let alloc_buffer ?old_buffer ~size_in_bytes _stream =
-    Backend.alloc_buffer ?old_buffer ~size_in_bytes ()
-
-  let device : device = make_device CPU ~ordinal:0
-
-  let get_device ~ordinal =
-    if ordinal <> 0 then
-      invalid_arg @@ "Sync_scheduler.get_device: there is only one device, but ordinal="
-      ^ Int.to_string ordinal;
-    device
-
-  let num_devices () = 1
-  let suggested_num_streams _ = !sync_suggested_num_streams
-  let get_used_memory _ = Backend.get_used_memory ()
-  let latest_stram_id = ref (-1)
-
-  let new_stream device =
-    Int.incr latest_stram_id;
-    make_stream device () ~stream_id:!latest_stram_id
-
-  let all_work _stream = ()
-  let is_idle _stream = true
-  let name = "sync_" ^ Backend.name
-  let await _stream = ()
-  (* let global_run_no = ref 0 *)
-
-  let initialize = Backend.initialize
-  let is_initialized = Backend.is_initialized
-  let get_name stream = [%string "%{name}:0:%{stream.stream_id#Int}"]
-  let schedule_task _stream task = Task.run task
-end
-
 let lower_assignments ?name bindings asgns =
   let name = Option.value_or_thunk name ~default:(fun () -> Assignments.get_name_exn asgns) in
   let unoptim_ll_source = Utils.get_debug_formatter ~fname:(name ^ "-unoptimized.ll") in
@@ -587,10 +344,10 @@ struct
   include Backend_device
 end
 
-module Cc_multicore = Make_device_backend_from_lowered (Multicore_scheduler) (Cc_backend)
-module Gcc_multicore = Make_device_backend_from_lowered (Multicore_scheduler) (Gcc_backend)
-module Cc_sync = Make_device_backend_from_lowered (Sync_scheduler) (Cc_backend)
-module Gcc_sync = Make_device_backend_from_lowered (Sync_scheduler) (Gcc_backend)
+module Cc_multicore = Make_device_backend_from_lowered (Schedulers.Multicore) (Cc_backend)
+module Gcc_multicore = Make_device_backend_from_lowered (Schedulers.Multicore) (Gcc_backend)
+module Cc_sync = Make_device_backend_from_lowered (Schedulers.Sync) (Cc_backend)
+module Gcc_sync = Make_device_backend_from_lowered (Schedulers.Sync) (Gcc_backend)
 
 let reinitialize (module Backend : Backend) config =
   if not @@ Backend.is_initialized () then Backend.initialize config
diff --git a/arrayjit/lib/backends.mli b/arrayjit/lib/backends.mli
@@ -2,8 +2,6 @@
 
 open Base
 
-val sync_suggested_num_streams : int ref
-
 val reinitialize : (module Backend_intf.Backend) -> Backend_intf.config -> unit
 (** Initializes the backend, and if it was already initialized, performs garbage collection. *)
 
diff --git a/arrayjit/lib/dune b/arrayjit/lib/dune
@@ -48,6 +48,7 @@
   cc_backend
   gcc_backend
   cuda_backend
+  schedulers
   backends)
  (modes byte native))
 
diff --git a/arrayjit/lib/schedulers.ml b/arrayjit/lib/schedulers.ml
diff --git a/bin/moons_benchmark.ml b/bin/moons_benchmark.ml
