Backends: Remove now-redundant accessor functions, share get_name

lukstafi · lukstafi · commit 6605dffbd203 · 2024-10-25T21:42:27.000+02:00
diff --git a/arrayjit/lib/backend_types.ml b/arrayjit/lib/backend_types.ml
@@ -123,6 +123,8 @@ module type Device_config = sig
   (** 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. *)
+
+  val name : string
 end
 
 type ('buffer_ptr, 'dev, 'event) device = {
@@ -148,7 +150,9 @@ type ('buffer_ptr, 'dev, 'runner, 'event) stream = {
   merge_buffer : ('buffer_ptr * Tnode.t) option ref;
   stream_id : int;
   mutable allocated_buffer : 'buffer_ptr buffer option;
-  queried_work_for : 'event option Hashtbl.M(Tnode).t; (* The completion event for updating the node via this stream. Only populated after the first time {!} *)
+  queried_work_for : 'event option Hashtbl.M(Tnode).t;
+      (* The completion event for updating the node via this stream. Only populated after the first
+         time {!} *)
 }
 [@@deriving sexp_of]
 
@@ -174,6 +178,7 @@ module type Device = sig
 
   val make_device : dev -> ordinal:int -> device
   val make_stream : device -> runner -> stream_id:int -> stream
+  val get_name : stream -> string
 end
 
 module Device_types (Device_config : Device_config) = struct
@@ -211,6 +216,8 @@ struct
       allocated_buffer = None;
       queried_work_for = Hashtbl.create (module Tnode);
     }
+
+  let get_name stream = [%string "%{name}:%{stream.device.ordinal#Int}:%{stream.stream_id#Int}"]
 end
 
 (** Parts shared by both assignments-level and lowered-level backend interfaces. *)
@@ -224,8 +231,6 @@ module type Backend_any_common = sig
   (** For backends derived via {!No_device_backend}, this is usually the backend name concatenated
       with the device or stream number. For {!Backend}, [init_info = stream]. *)
 
-  val name : string
-
   val initialize : config -> unit
   (** Initializes a backend before first use. Typically does nothing if the backend is already
       initialized, but some backends can do some safe cleanups. *)
@@ -286,6 +291,8 @@ module type No_device_backend = sig
   include Backend_common with type init_info := string and type stream := unit
   include Backend_impl_common with type context := context and type buffer_ptr := buffer_ptr
 
+  val name : string
+
   val link : merge_buffer:(buffer_ptr * Tnode.t) option ref -> context -> code -> context routine
   (** Returns the routine for the code's procedure, in a new context derived from the given context. *)
 
@@ -350,9 +357,6 @@ module type Backend_device_common = sig
 
   val new_stream : device -> stream
   val get_ctx_stream : context -> stream
-  val get_stream_device : stream -> device
-  val to_ordinal : device -> int
-  val get_name : stream -> string
 end
 
 module type With_buffer_retrieval_and_syncing = sig
@@ -427,6 +431,8 @@ module type Lowered_no_device_backend = sig
        and type init_info := string
        and type buffer_ptr := buffer_ptr
 
+  val name : string
+
   type procedure [@@deriving sexp_of]
 
   val compile :
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -47,7 +47,7 @@ module Add_buffer_retrieval_and_syncing (Backend : No_buffer_retrieval_or_syncin
 
   let%diagn2_l_sexp device_to_device (tn : Tn.t) ~into_merge_buffer ~(dst : Backend.context)
       ~(src : Backend.context) =
-    let ordinal_of ctx = Backend.(to_ordinal @@ get_stream_device @@ get_ctx_stream ctx) in
+    let ordinal_of ctx = (Backend.get_ctx_stream ctx).device.ordinal in
     let name_of ctx = Backend.(get_name @@ get_ctx_stream ctx) in
     let same_device = ordinal_of dst = ordinal_of src in
     if same_device && (Tn.known_shared_cross_stream tn || String.equal (name_of src) (name_of dst))
@@ -123,6 +123,8 @@ module Multicore_backend (Backend : No_device_backend) = struct
 
     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 Alloc_buffer = struct
@@ -289,9 +291,7 @@ module Multicore_backend (Backend : No_device_backend) = struct
     Stdlib.Gc.finalise cleanup_stream stream;
     stream
 
-  let get_stream_device stream = stream.device
   let get_ctx_stream { stream; _ } = stream
-  let to_ordinal _ = 0
 
   let from_host ~dst_ptr ~dst hosted =
     let work () = host_to_buffer hosted ~dst:dst_ptr in
@@ -350,6 +350,8 @@ module Sync_backend (Backend : No_device_backend) = struct
     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 Alloc_buffer = struct
@@ -370,7 +372,6 @@ module Sync_backend (Backend : No_device_backend) = struct
     Backend.alloc_buffer ?old_buffer ~size_in_bytes ()
 
   let device : device = make_device CPU ~ordinal:0
-  let to_ordinal device = device.ordinal
 
   let get_device ~ordinal =
     if ordinal <> 0 then
@@ -399,7 +400,6 @@ module Sync_backend (Backend : No_device_backend) = struct
   type context = { stream : stream; ctx : Backend.context } [@@deriving sexp_of]
 
   let get_ctx_stream context = context.stream
-  let get_stream_device stream = stream.device
   let ctx_arrays context = ctx_arrays context.ctx
   let init stream = { stream; ctx = Backend.init name }
   let initialize = Backend.initialize
diff --git a/arrayjit/lib/cc_backend.ml b/arrayjit/lib/cc_backend.ml
@@ -10,6 +10,8 @@ let _get_local_debug_runtime = Utils._get_local_debug_runtime
 include Backend_types.No_device_buffer_and_copying
 open Backend_types
 
+let name = "cc"
+
 let optimization_level () =
   Int.of_string @@ Utils.get_global_arg ~default:"3" ~arg_name:"cc_backend_optimization_level"
 
@@ -257,5 +259,3 @@ let%diagn_sexp link_compiled ~merge_buffer (prior_context : context) (code : pro
         work;
       },
     name )
-
-let name = "cc"
diff --git a/arrayjit/lib/cuda_backend.cudajit.ml b/arrayjit/lib/cuda_backend.cudajit.ml
@@ -36,6 +36,8 @@ module Device_config = struct
   type dev = { dev : Cu.Device.t; primary_context : Cu.Context.t } [@@deriving sexp_of]
   type runner = Cu.Stream.t [@@deriving sexp_of]
   type event = Cu.Delimited_event.t [@@deriving sexp_of]
+
+  let name = "cuda"
 end
 
 module Device_stream = Device_types (Device_config)
@@ -185,10 +187,6 @@ let suggested_num_streams device =
   | Most_parallel_streams -> (cuda_properties device).multiprocessor_count
 
 let get_ctx_stream { stream; _ } = stream
-let get_stream_device { device; _ } = device
-let to_ordinal { ordinal; _ } = ordinal
-let name = "cuda"
-let get_name stream = [%string "%{name}:%{stream.device.ordinal#Int}:%{stream.stream_id#Int}"]
 
 let await stream : unit =
   set_ctx stream.device.dev.primary_context;
diff --git a/bin/compilation_speed.ml b/bin/compilation_speed.ml
@@ -26,7 +26,7 @@ let benchmark_overhead backend () =
   (* Train.every_non_literal_on_host f; *)
   let stream = Backend.(new_stream @@ get_device ~ordinal:0) in
   let ctx = Backend.init stream in
-  let init_mem = Backend.(get_used_memory @@ get_stream_device stream) in
+  let init_mem = Backend.(get_used_memory stream.device) in
   let update_f = Train.grad_update f in
   (* Initialize the context with a mock update of x to ensure that it is not optimized as a
      constant. *)
@@ -63,7 +63,7 @@ let benchmark_overhead backend () =
   in
   let final_time = Time_now.nanoseconds_since_unix_epoch () in
   let time_in_sec = Int63.(to_float @@ (final_time - init_time)) /. 1000_000_000. in
-  let mem_in_bytes = Backend.(get_used_memory @@ get_stream_device stream) - init_mem in
+  let mem_in_bytes = Backend.(get_used_memory stream.device) - init_mem in
   let result =
     PrintBox_utils.Benchmark
       {
