Backends: expose Types.stream from the signatures, implement work_for

lukstafi · lukstafi · commit e0515ef4f6de · 2024-10-22T16:56:38.000+02:00
diff --git a/arrayjit/lib/backend_types.ml b/arrayjit/lib/backend_types.ml
@@ -170,15 +170,24 @@ end
 (** Parts shared by both assignments-level and lowered-level backend interfaces providing streams
     and devices. *)
 module type Backend_device_common = sig
-  type stream [@@deriving sexp_of]
-
-  include Backend_any_common with type init_info := stream and type stream := stream
+  type buffer_ptr
+  type device
 
   type event
   (** An event tracks if a stream finished computing past a particular point in its schedue. These
       values are used internally for scheduling across streams of the backend, and can be used for
       explicit scheduling. *)
 
+  type stream_state [@@deriving sexp_of]
+  type runner [@@deriving sexp_of]
+  type stream = (buffer_ptr, event, device, stream_state, runner) Types.stream [@@deriving sexp_of]
+
+  include
+    Backend_any_common
+      with type buffer_ptr := buffer_ptr
+       and type init_info := stream
+       and type stream := stream
+
   val sync : event -> unit
   (** Blocks till the event completes, if it's not done already. *)
 
@@ -192,8 +201,6 @@ module type Backend_device_common = sig
       called internally when necessary. But there is one exception, see {!device_to_device} when
       [into_merge_buffer=Streaming]. *)
 
-  type device
-
   val get_used_memory : device -> int
   (** Returns (an upper bound of) the memory used for arrays, in bytes. *)
 
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -24,8 +24,14 @@ struct
   type context = Backend.context
   type event = Backend.event
 
-  (* FIXME: *)
-  let work_for _context _tn = failwith "NOT IMPLEMENTED YET"
+  let work_for context tn =
+    let stream = Backend.get_ctx_stream context in
+    let default () = Some (Backend.all_work stream) in
+    if not @@ Map.mem (Backend.ctx_arrays context) tn then None
+    else
+      Hashtbl.update_and_return stream.requested_work_for tn ~f:(function
+        | None | Some None -> default ()
+        | Some (Some _ as event) -> event)
 
   let%diagn2_l_sexp from_host (ctx : Backend.context) tn =
     match (tn, Map.find (Backend.ctx_arrays ctx) tn) with
@@ -109,13 +115,13 @@ module Multicore_backend (Backend : Backend_types.No_device_backend) = struct
   }
   [@@deriving sexp_of]
 
-  type domain = unit Domain.t
+  type runner = unit Domain.t
 
-  let sexp_of_domain (d : domain) = Sexp.Atom ("domain-" ^ Int.to_string (Domain.get_id d :> int))
+  let sexp_of_runner (d : runner) = Sexp.Atom ("domain-" ^ Int.to_string (Domain.get_id d :> int))
 
   type device = CPU [@@deriving sexp_of]
   type event = Not_implemented_yet [@@deriving sexp_of]
-  type nonrec stream = (buffer_ptr, event, device, stream_state, domain) stream [@@deriving sexp_of]
+  type nonrec stream = (buffer_ptr, event, device, stream_state, runner) stream [@@deriving sexp_of]
 
   (** TODO: Blocks till the event completes, if it's not done already. *)
   let sync Not_implemented_yet = ()
@@ -359,7 +365,9 @@ module Sync_backend (Backend : Backend_types.No_device_backend) = struct
   let get_used_memory CPU = Backend.get_used_memory ()
   let next_stream = ref 0
 
-  type nonrec stream = (buffer_ptr, event, device, unit, unit) stream [@@deriving sexp_of]
+  type stream_state = unit [@@deriving sexp_of]
+  type runner = unit [@@deriving sexp_of]
+  type nonrec stream = (buffer_ptr, event, device, stream_state, runner) stream [@@deriving sexp_of]
 
   let new_stream CPU : stream =
     Int.incr next_stream;
diff --git a/arrayjit/lib/cuda_backend.cudajit.ml b/arrayjit/lib/cuda_backend.cudajit.ml
@@ -42,7 +42,9 @@ type device = {
 }
 [@@deriving sexp_of]
 
-type nonrec stream = (buffer_ptr, event, device, unit, Cu.Stream.t) stream [@@deriving sexp_of]
+type stream_state = unit [@@deriving sexp_of]
+type runner = Cu.Stream.t [@@deriving sexp_of]
+type nonrec stream = (buffer_ptr, event, device, stream_state, runner) stream [@@deriving sexp_of]
 
 type context = {
   label : string;

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,9 @@ type device = {`
`42`	`42`	`}`
`43`	`43`	`[@@deriving sexp_of]`
`44`	`44`
`45`		`-type nonrec stream = (buffer_ptr, event, device, unit, Cu.Stream.t) stream [@@deriving sexp_of]`
	`45`	`+type stream_state = unit [@@deriving sexp_of]`
	`46`	`+type runner = Cu.Stream.t [@@deriving sexp_of]`
	`47`	`+type nonrec stream = (buffer_ptr, event, device, stream_state, runner) stream [@@deriving sexp_of]`
`46`	`48`
`47`	`49`	`type context = {`
`48`	`50`	`label : string;`