Host memory wrapper needs to know the size

lukstafi · lukstafi · commit e586db290e23 · 2025-04-29T22:27:18.000+02:00
diff --git a/arrayjit/lib/backend_impl.ml b/arrayjit/lib/backend_impl.ml
@@ -15,7 +15,7 @@ open Backend_intf
 module type No_device_buffer_and_copying = sig
   include Alloc_buffer with type stream := unit
 
-  val use_host_memory : (unit Ctypes.ptr -> buffer_ptr) option
+  val use_host_memory : (size_in_bytes:int -> unit Ctypes.ptr -> buffer_ptr) option
 
   val get_used_memory : unit -> int
   (** Returns (an upper bound of) the memory used for arrays, in bytes. *)
@@ -29,7 +29,7 @@ module No_device_buffer_and_copying () :
   No_device_buffer_and_copying with type buffer_ptr = unit Ctypes.ptr = struct
   type buffer_ptr = unit Ctypes.ptr
 
-  let use_host_memory = Some Fn.id
+  let use_host_memory = Some (fun ~size_in_bytes:_ ptr -> ptr)
   let sexp_of_buffer_ptr = Ops.sexp_of_voidptr
 
   include Buffer_types (struct
@@ -145,9 +145,9 @@ end
 module type Backend_impl_common = sig
   include Backend_intf.Buffer
 
-  val use_host_memory : (unit Ctypes.ptr -> buffer_ptr) option
+  val use_host_memory : (size_in_bytes:int -> unit Ctypes.ptr -> buffer_ptr) option
   (** If not [None], the backend will read from and write to the host memory directly whenever
-      reasonable.
+      reasonable. [size_in_bytes] is the size of the memory allocated on the host.
 
       [use_host_memory] can only be [Some] on unified memory devices, like CPU and Apple Metal. *)
 end
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -413,9 +413,11 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
                   && Tn.known_shared_cross_streams key && Tn.is_hosted_force key 44
                 then
                   Hashtbl.update_and_return device.cross_stream_candidates key ~f:(fun _ ->
-                      get_buffer_ptr @@ Ndarray.get_voidptr_not_managed
-                      @@ Option.value_exn ~here:[%here]
-                      @@ Lazy.force key.array)
+                      get_buffer_ptr
+                        ~size_in_bytes:(Lazy.force key.size_in_bytes)
+                        @@ Ndarray.get_voidptr_not_managed
+                        @@ Option.value_exn ~here:[%here]
+                        @@ Lazy.force key.array)
                 else Hashtbl.find_or_add device.cross_stream_candidates key ~default
           in
           if Hashtbl.mem device.cross_stream_candidates key then
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -16,7 +16,7 @@ module C_syntax (B : sig
 
   type buffer_ptr
 
-  val use_host_memory : (unit Ctypes.ptr -> buffer_ptr) option
+  val use_host_memory : (size_in_bytes:int -> unit Ctypes.ptr -> buffer_ptr) option
   val logs_to_stdout : bool
   val main_kernel_prefix : string
   val kernel_prep_line : string
diff --git a/arrayjit/lib/dune b/arrayjit/lib/dune
@@ -72,6 +72,8 @@
 
 (library
  (name metal_backend)
+ ; Temporary disabled until we have a working Metal backend
+ (enabled_if false)
  (optional)
  (modules metal_backend)
  (libraries base metal utils ir)
