From a156c25b82d47f12501285f086beea06fe9c8b21 Mon Sep 17 00:00:00 2001 From: Lukasz Stafiniak Date: Mon, 21 Oct 2024 20:04:17 +0200 Subject: [PATCH] Backends: Unused per-context `used_memory` --- arrayjit/lib/backend_types.ml | 6 ++--- arrayjit/lib/backends.ml | 38 +++++++++++++++--------------- arrayjit/lib/cc_backend.ml | 6 ++--- arrayjit/lib/gcc_backend.gccjit.ml | 4 ++-- 4 files changed, 26 insertions(+), 28 deletions(-) diff --git a/arrayjit/lib/backend_types.ml b/arrayjit/lib/backend_types.ml index 98207671..51db91a0 100644 --- 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 index 4d4725ab..4131245c 100644 --- 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 index 5f1c13fc..21d9d015 100644 --- 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 index 89e819d7..eeec34ed 100644 --- 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 =