Rename unified_memory -> use_host_memory and add more debugging

lukstafi · lukstafi · commit 069322acabad · 2024-11-09T14:20:58.000+01:00
diff --git a/arrayjit/lib/assignments.ml b/arrayjit/lib/assignments.ml
@@ -79,10 +79,10 @@ let is_total ~initialize_neutral ~projections =
 
 (** Returns materialized nodes in the sense of {!Tnode.is_in_context_force}. NOTE: it ideally should
     be called after compilation. *)
-let context_nodes ~unified_memory asgns =
+let context_nodes ~use_host_memory asgns =
   let open Utils.Set_O in
   let empty = Set.empty (module Tn) in
-  let one tn = if Tnode.is_in_context_force ~unified_memory tn 34 then Set.singleton (module Tn) tn else empty in
+  let one tn = if Tnode.is_in_context_force ~use_host_memory tn 34 then Set.singleton (module Tn) tn else empty in
   let of_node = function Node rhs -> one rhs | Merge_buffer _ -> empty in
   let rec loop = function
     | Noop -> empty
diff --git a/arrayjit/lib/backend_impl.ml b/arrayjit/lib/backend_impl.ml
@@ -133,11 +133,10 @@ end
 module type Backend_impl_common = sig
   include Buffer
 
-  val unified_memory : bool
-  (** If true, the node is required to be in the contexts linked with code that uses it.
+  val use_host_memory : bool
+  (** If true, the backend will read from and write to the host memory directly whenever possible.
 
-      Should return false for nodes that are virtual, local, or which the backend prefers to access
-      directly from the host. *)
+      [use_host_memory] can only be true on unified memory devices, like CPU and Apple Metal. *)
 end
 
 (** An interface to adding schedulers for stream-agnostic (typically CPU) backend implementations. *)
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -115,18 +115,18 @@ let lower_batch_assignments ?names ?occupancy bindings asgns_l =
              Some (Assignments.lower ~unoptim_ll_source ~ll_source ~cd_source ~name bound asgns) )
          else (None, None))
 
-let verify_prior_context ~unified_memory ~ctx_arrays ~from_prior_context =
+let verify_prior_context ~use_host_memory ~ctx_arrays ~from_prior_context =
   Set.iter from_prior_context ~f:(fun tn ->
       if
-        Tn.is_in_context_force ~unified_memory tn 342
+        Tn.is_in_context_force ~use_host_memory tn 342
         && not (Option.is_some @@ Map.find ctx_arrays tn)
       then raise @@ Utils.User_error ("The linked context lacks node " ^ Tnode.debug_name tn))
 
-let from_prior_context_batch ~unified_memory comps =
+let from_prior_context_batch ~use_host_memory comps =
   Array.filter_map comps ~f:(fun comp ->
       Option.map comp ~f:(fun comp ->
           Set.diff
-            (Assignments.context_nodes ~unified_memory comp.Assignments.asgns)
+            (Assignments.context_nodes ~use_host_memory comp.Assignments.asgns)
             comp.embedded_nodes))
   |> Array.fold ~init:(Set.empty (module Tnode)) ~f:Set.union
 
@@ -270,7 +270,7 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
     let name, lowered = lower_assignments ?name bindings comp.Assignments.asgns in
     let code = compile ?shared ~name bindings lowered in
     let from_prior_context =
-      Set.diff (Assignments.context_nodes ~unified_memory comp.asgns) comp.embedded_nodes
+      Set.diff (Assignments.context_nodes ~use_host_memory comp.asgns) comp.embedded_nodes
     in
     { from_prior_context; name; lowered; code; expected_merge_node = lowered.Low_level.merge_node }
 
@@ -281,7 +281,7 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
     in
     let code_batch = compile_batch ?shared ~names bindings lowereds in
     let from_prior_context =
-      from_prior_context_batch ~unified_memory
+      from_prior_context_batch ~use_host_memory
       @@ Array.mapi lowereds ~f:(fun i -> Option.map ~f:(fun _ -> comps.(i)))
     in
     {
@@ -295,7 +295,7 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
     }
 
   let%track3_sexp alloc_if_needed (stream : stream) ~key ~data:node ctx_arrays =
-    if Tnode.is_in_context_force ~unified_memory key 345 && not (Map.mem ctx_arrays key) then (
+    if Tnode.is_in_context_force ~use_host_memory key 345 && not (Map.mem ctx_arrays key) then (
       [%log2 Tn.debug_name key];
       [%log3 (key : Tnode.t)];
       let default () =
@@ -307,7 +307,7 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
         if Tn.known_non_cross_stream key then add_new ()
         else (
           if Hashtbl.mem device.cross_stream_candidates key then
-            Tn.update_memory_sharing key Tn.Shared_cross_stream 40;
+            Tn.update_memory_sharing key Tn.Shared_cross_stream 39;
           let data = Hashtbl.find_or_add device.cross_stream_candidates key ~default in
           Map.add_exn ctx_arrays ~key ~data)
       else if Tn.known_shared_cross_stream key then (
@@ -326,8 +326,8 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
         add_new ()))
     else ctx_arrays
 
-  let link context (code : code) =
-    verify_prior_context ~unified_memory ~ctx_arrays:context.ctx_arrays
+  let%debug3_sexp link context (code : code) =
+    verify_prior_context ~use_host_memory ~ctx_arrays:context.ctx_arrays
       ~from_prior_context:code.from_prior_context;
     let inputs, outputs = Low_level.input_and_output_nodes code.lowered in
     let ctx_arrays =
@@ -344,8 +344,8 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
     in
     { context; schedule; bindings; name = code.name; inputs; outputs }
 
-  let link_batch context code_batch =
-    verify_prior_context ~unified_memory ~ctx_arrays:context.ctx_arrays
+  let%debug3_sexp link_batch context code_batch =
+    verify_prior_context ~use_host_memory ~ctx_arrays:context.ctx_arrays
       ~from_prior_context:code_batch.from_prior_context;
     let ctx_arrays =
       Array.map code_batch.lowereds
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -18,8 +18,7 @@ module C_syntax (B : sig
       shared and already known. *)
 
   val hardcoded_context_ptr : (buffer_ptr -> Ops.prec -> string) option
-  val unified_memory : bool
-  val host_ptrs_for_readonly : bool
+  val use_host_memory : bool
   val logs_to_stdout : bool
   val main_kernel_prefix : string
   val kernel_prep_line : string
@@ -33,7 +32,7 @@ struct
   let get_ident =
     Low_level.get_ident_within_code ~no_dots:true @@ Array.map B.procs ~f:(fun (l, _) -> l.llc)
 
-  let in_ctx tn = B.(Tn.is_in_context_force ~unified_memory tn 341)
+  let in_ctx tn = B.(Tn.is_in_context_force ~use_host_memory tn 341)
 
   let pp_zero_out ppf tn =
     Stdlib.Format.fprintf ppf "@[<2>memset(%s, 0, %d);@]@ " (get_ident tn) @@ Tn.size_in_bytes tn
@@ -73,18 +72,15 @@ struct
             if not @@ Hash_set.mem is_global tn then
               let ctx_ptr = B.hardcoded_context_ptr in
               let mem : (Tn.memory_mode * int) option = tn.memory_mode in
-              match
-                (in_ctx tn, ctx_ptr, ctx_arrays, B.host_ptrs_for_readonly, mem, node.read_only)
-              with
-              | true, Some get_ptr, Some ctx_arrays, _, _, _ ->
+              match (in_ctx tn, ctx_ptr, ctx_arrays, B.use_host_memory, mem) with
+              | true, Some get_ptr, Some ctx_arrays, _, _ ->
                   let ident = get_ident tn in
                   let ctx_array =
                     Option.value_exn ~here:[%here] ~message:ident @@ Map.find ctx_arrays tn
                   in
                   fprintf ppf "#define %s (%s)@," ident @@ get_ptr ctx_array (Lazy.force tn.prec);
                   Hash_set.add is_global tn
-              | false, _, _, true, Some (Hosted _, _), true ->
-                  (* In-context nodes to read directly from host would be error prone. *)
+              | false, _, _, true, Some (Hosted _, _) ->
                   let nd = Option.value_exn ~here:[%here] @@ Lazy.force tn.array in
                   fprintf ppf "#define %s (%s)@," (get_ident tn) (Ndarray.c_ptr_to_string nd);
                   Hash_set.add is_global tn
@@ -288,7 +284,7 @@ struct
     in
     pp_ll ppf llc
 
-  let%diagn_sexp compile_proc ~name ppf idx_params ~is_global
+  let%track3_sexp compile_proc ~name ppf idx_params ~is_global
       Low_level.{ traced_store; llc; merge_node } =
     let open Stdlib.Format in
     let params : (string * param_source) list =
diff --git a/arrayjit/lib/cc_backend.ml b/arrayjit/lib/cc_backend.ml
@@ -33,7 +33,7 @@ type procedure = {
 }
 [@@deriving sexp_of]
 
-let unified_memory = true
+let use_host_memory = true
 
 let get_global_run_id =
   let next_id = ref 0 in
@@ -79,8 +79,7 @@ struct
 
   let procs = Input.procs
   let hardcoded_context_ptr = c_ptr_to_string
-  let unified_memory = unified_memory
-  let host_ptrs_for_readonly = true
+  let use_host_memory = use_host_memory
   let logs_to_stdout = false
   let main_kernel_prefix = ""
   let kernel_prep_line = ""
@@ -136,7 +135,7 @@ let%diagn_sexp compile_batch ~names ~opt_ctx_arrays bindings
       Option.map names.(i) ~f:(fun name ->
           { result; params = Option.value_exn ~here:[%here] params; bindings; name }))
 
-let%diagn_sexp link_compiled ~merge_buffer ~runner_label ctx_arrays (code : procedure) =
+let%track3_sexp link_compiled ~merge_buffer ~runner_label ctx_arrays (code : procedure) =
   let name : string = code.name in
   List.iter code.params ~f:(function
     | _, Param_ptr tn ->
diff --git a/arrayjit/lib/cuda_backend.cudajit.ml b/arrayjit/lib/cuda_backend.cudajit.ml
@@ -30,7 +30,7 @@ module Backend_buffer = struct
   end)
 end
 
-let unified_memory = false
+let use_host_memory = false
 
 module Device_config = struct
   include Backend_buffer
@@ -271,9 +271,7 @@ struct
 
   let procs = Input.procs
   let hardcoded_context_ptr = None
-  let unified_memory = unified_memory
-  let host_ptrs_for_readonly = false
-  (* GPUs cannot access host memory pointers directly. *)
+  let use_host_memory = use_host_memory
 
   let logs_to_stdout = true
   let main_kernel_prefix = "extern \"C\" __global__"
diff --git a/arrayjit/lib/cuda_backend.missing.ml b/arrayjit/lib/cuda_backend.missing.ml
@@ -19,7 +19,7 @@ let compile ?shared:_ ~name:_ bindings _optimized = bindings
 let compile_batch ?shared:_ ~names:_ (bindings : Indexing.unit_bindings) optimized : code_batch =
   Array.map optimized ~f:(fun _ -> bindings)
 
-let unified_memory = false
+let use_host_memory = false
 let ctx_arrays Unimplemented_ctx = Map.empty (module Tnode)
 
 let link (Unimplemented_ctx : context) (code : code) =
diff --git a/arrayjit/lib/gcc_backend.gccjit.ml b/arrayjit/lib/gcc_backend.gccjit.ml
@@ -66,7 +66,7 @@ type procedure = {
 }
 [@@deriving sexp_of]
 
-let unified_memory = true
+let use_host_memory = true
 
 let gcc_typ_of_prec =
   let open Gccjit in
@@ -113,7 +113,7 @@ let prepare_node ~debug_log_zero_out ~get_ident ctx traced_store ~opt_ctx_arrays
   let num_typ = Type.(get ctx c_typ) in
   let ptr_typ = Type.pointer num_typ in
   let ident = get_ident tn in
-  let in_ctx = Tn.is_in_context_force ~unified_memory tn 343 in
+  let in_ctx = Tn.is_in_context_force ~use_host_memory tn 343 in
   let ptr =
     match (in_ctx, opt_ctx_arrays, Tn.is_hosted_force tn 344) with
     | true, Some ctx_arrays, _ ->
diff --git a/arrayjit/lib/gcc_backend.missing.ml b/arrayjit/lib/gcc_backend.missing.ml
@@ -5,7 +5,7 @@ include Backend_impl.No_device_buffer_and_copying ()
 let expected_merge_node Unimplemented_proc =
   failwith "gcc backend missing: install the optional dependency gccjit"
 
-let unified_memory = true
+let use_host_memory = true
 
 let to_buffer _tn ~dst:_ ~src:_ =
   failwith "gcc backend missing: install the optional dependency gccjit"
diff --git a/arrayjit/lib/tnode.ml b/arrayjit/lib/tnode.ml
@@ -183,11 +183,11 @@ let is_materialized_force tn provenance =
   | Some ((On_device _ | Hosted _ | Materialized), _) -> true
   | Some ((Never_virtual | Device_only | Effectively_constant), _) -> assert false
 
-let is_in_context_force ~unified_memory tn provenance =
+let is_in_context_force ~use_host_memory tn provenance =
   default_to_most_local tn provenance;
   match tn.memory_mode with
   | Some (Hosted (Constant | Volatile | Changed_on_devices Shared_cross_stream), _)
-    when unified_memory ->
+    when use_host_memory ->
       false
   | Some ((Virtual | Local), _) -> false
   | _ -> true
diff --git a/bin/moons_demo_parallel.ml b/bin/moons_demo_parallel.ml
@@ -9,9 +9,7 @@ module Utils = Arrayjit.Utils
 module Rand = Arrayjit.Rand.Lib
 
 let experiment ~seed ~backend_name ~config () =
-  (* Utils.set_log_level 1; *)
-  (* Utils.set_log_level 3; *)
-  (* Utils.settings.output_debug_files_in_build_directory <- true; *)
+  Utils.settings.output_debug_files_in_build_directory <- true;
   (* Utils.settings.debug_log_from_routines <- true; *)
   let hid_dim = 16 in
   (* let hid_dim = 4 in *)
