Backends: Unused per-context used_memory

lukstafi · lukstafi · commit a156c25b82d4 · 2024-10-21T20:04:17.000+02:00
diff --git a/arrayjit/lib/backend_types.ml b/arrayjit/lib/backend_types.ml
@@ -9,11 +9,9 @@ let _get_local_debug_runtime = Utils._get_local_debug_runtime
 
 module No_device_types = struct
   type ctx_array = Ndarray.t [@@deriving sexp_of]
+  type ctx_arrays = { ctx_arrays : ctx_array Map.M(Tnode).t } [@@deriving sexp_of]
 
-  type ctx_arrays = { used_memory : Utils.atomic_int; ctx_arrays : ctx_array Map.M(Tnode).t }
-  [@@deriving sexp_of]
-
-  let empty_ctx_arrays = { used_memory = Atomic.make 0; ctx_arrays = Map.empty (module Tnode) }
+  let empty_ctx_arrays = { ctx_arrays = Map.empty (module Tnode) }
   let get_array arrays = Map.find arrays.ctx_arrays
 end
 
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -97,25 +97,6 @@ module Multicore_backend (Backend : Backend_types.No_device_backend) = struct
   let sexp_of_task_queue q =
     Sexp.(List [ Atom "task_queue_of_size"; Atom (Int.to_string @@ Queue.size q) ])
 
-  type event = Not_implemented_yet  (** TODO: NOT IMPLEMENTED YET *)
-
-  (** 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: If the tensor node is in the context, returns the event indicating if currently running
-      or scheduled computations modifying that node on the context's stream have completed.
-
-      NOTE: [work_for ctx tn], if work tracking was not registered for [tn], will register work
-      tracking for [tn] and return the event tracking all currently scheduled computations on
-      [ctx]'s stream. *)
-  let work_for _ctx _tn = Some Not_implemented_yet
-
-  (** TODO: Schedules waiting for the given event on the context's stream. *)
-  let will_wait_for _ctx Not_implemented_yet = ()
-
   type stream_state = {
     mutable keep_spinning : bool;
     mutable stream_error : exn option;
@@ -136,6 +117,25 @@ module Multicore_backend (Backend : Backend_types.No_device_backend) = struct
   }
   [@@deriving sexp_of]
 
+  type event = Not_implemented_yet
+
+  (** 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: If the tensor node is in the context, returns the event indicating if currently running
+      or scheduled computations modifying that node on the context's stream have completed.
+
+      NOTE: [work_for ctx tn], if work tracking was not registered for [tn], will register work
+      tracking for [tn] and return the event tracking all currently scheduled computations on
+      [ctx]'s stream. *)
+  let work_for _ctx _tn = Some Not_implemented_yet
+
+  (** TODO: Schedules waiting for the given event on the context's stream. *)
+  let will_wait_for _ctx Not_implemented_yet = ()
+
   let alloc_buffer ?old_buffer ~size_in_bytes _stream = alloc_buffer ?old_buffer ~size_in_bytes ()
   let get_used_memory _device = get_used_memory ()
 
diff --git a/arrayjit/lib/cc_backend.ml b/arrayjit/lib/cc_backend.ml
@@ -137,7 +137,7 @@ let%diagn_sexp compile ~(name : string) ~opt_ctx_arrays bindings (lowered : Low_
                     Ndarray.create_array ~debug (Lazy.force tn.Tn.prec) ~dims:(Lazy.force tn.dims)
                     @@ Constant_fill { values = [| 0. |]; strict = false }
                   in
-                  { ctx_arrays with ctx_arrays = Map.add_exn ctx_arrays.ctx_arrays ~key:tn ~data }
+                  { ctx_arrays = Map.add_exn ctx_arrays.ctx_arrays ~key:tn ~data }
                 else ctx_arrays
             | Some _ -> ctx_arrays))
   in
@@ -176,7 +176,7 @@ let%diagn_sexp compile_batch ~names ~opt_ctx_arrays bindings
                       else ctx_arrays
                   | Some _ -> ctx_arrays))
         in
-        { arrays with ctx_arrays })
+        { ctx_arrays })
   in
   let module Syntax = C_syntax.C_syntax (C_syntax_config (struct
     let for_lowereds = for_lowereds
@@ -231,7 +231,7 @@ let%diagn_sexp link_compiled ~merge_buffer (prior_context : context) (code : pro
                     Ndarray.create_array ~debug (Lazy.force tn.Tn.prec) ~dims:(Lazy.force tn.dims)
                     @@ Constant_fill { values = [| 0. |]; strict = false }
               in
-              { ctx_arrays with ctx_arrays = Map.update ctx_arrays.ctx_arrays tn ~f }
+              { ctx_arrays = Map.update ctx_arrays.ctx_arrays tn ~f }
           | _ -> ctx_arrays)
   in
   let context = { label; arrays } in
diff --git a/arrayjit/lib/gcc_backend.gccjit.ml b/arrayjit/lib/gcc_backend.gccjit.ml
@@ -697,7 +697,7 @@ let%diagn_sexp compile_proc ~name ~opt_ctx_arrays ctx bindings ~get_ident
   let opt_ctx_arrays =
     match (opt_ctx_arrays, ctx_nodes) with
     | None, _ | _, Param_ptrs _ -> None
-    | Some arrays, Ctx_arrays { contents } -> Some { arrays with ctx_arrays = contents }
+    | Some _arrays, Ctx_arrays { contents } -> Some { ctx_arrays = contents }
   in
   (ctx_info, opt_ctx_arrays, params)
 
@@ -786,7 +786,7 @@ let%diagn_sexp link_compiled ~merge_buffer (prior_context : context) (code : pro
           | _ -> arrays)
   in
   let context =
-    { label; arrays = { prior_context.arrays with ctx_arrays }; result = Some code.result }
+    { label; arrays = { ctx_arrays }; result = Some code.result }
   in
   let log_file_name = Utils.diagn_log_file [%string "debug-%{label}-%{code.name}.log"] in
   let run_variadic =
