fix: restore cross-device synchronization in device_to_device copies

When source and destination are on different devices, wait for source
stream writes to complete before scheduling the copy on the destination
stream. Without this, async CUDA/Metal copies could read stale data.

Addresses review feedback on PR #12.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: lukstafi

arrayjit/lib/backend_intf.ml

@@ -36,10 +36,10 @@ end
 
 type merge_buffer_use = No | Copy [@@deriving sexp_of]
 
-type kparam_source =
+type param_source =
   | Log_file_name
   | Merge_buffer
-  | Kparam_ptr of Tnode.t
+  | Param_ptr of Tnode.t
   | Static_idx of Indexing.static_symbol
 [@@deriving sexp_of]
 

arrayjit/lib/backends.ml

@@ -88,6 +88,10 @@ module Add_buffer_retrieval_and_syncing (Backend : No_buffer_retrieval_or_syncin
       match Map.find src.ctx_arrays tn with
       | None -> false
       | Some s_arr -> (
+          (* For cross-device copies, wait for the source stream's writes to complete
+             before the destination stream reads the data. *)
+          if not same_device then
+            Backend.will_wait_for dst (Backend.all_work src.stream);
           match into_merge_buffer with
           | No -> (
               match Map.find dst.ctx_arrays tn with
@@ -122,6 +126,9 @@ module Add_buffer_retrieval_and_syncing (Backend : No_buffer_retrieval_or_syncin
                ("init_from_device: tensor node " ^ Tn.debug_name tn
               ^ " already on same stream, for stream " ^ Backend.get_name src.stream)
         else (
+          (* For cross-device copies, wait for source writes to complete. *)
+          if not same_device then
+            Backend.will_wait_for dst (Backend.all_work src.stream);
           match Map.find dst.ctx_arrays tn with
           | Some _ ->
               raise

arrayjit/lib/cuda_backend.ml

@@ -284,7 +284,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
   type code = {
     traced_store : Low_level.traced_store;
     ptx : Nvrtc.compile_to_ptx_result;
-    kparams : (string * kparam_source) list;
+    params : (string * param_source) list;
     bindings : Indexing.unit_bindings;
     name : string;
   }
@@ -294,7 +294,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
     traced_stores : Low_level.traced_store option array;
     ptx : Nvrtc.compile_to_ptx_result;
     bindings : Indexing.unit_bindings;
-    kparams_and_names : ((string * kparam_source) list * string) option array;
+    params_and_names : ((string * param_source) list * string) option array;
   }
   [@@deriving sexp_of]
 
@@ -839,7 +839,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
       let procs = [| lowered |]
     end)) in
     let idx_params = Indexing.bound_symbols bindings in
-    let kparams, proc_doc = Syntax.compile_proc ~name idx_params lowered in
+    let params, proc_doc = Syntax.compile_proc ~name idx_params lowered in
     let cuda_includes =
       {|#include <cuda_fp16.h>
 #include <cuda_bf16.h>
@@ -861,21 +861,21 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
         ~proc_doc
     in
     let ptx = cuda_to_ptx ~name source in
-    { traced_store; ptx; kparams; bindings; name }
+    { traced_store; ptx; params; bindings; name }
 
   let%diagn2_sexp compile_batch ~names bindings lowereds =
     let module Syntax = C_syntax.C_syntax (Cuda_syntax_config (struct
       let procs = Array.filter_opt lowereds
     end)) in
     let idx_params = Indexing.bound_symbols bindings in
-    let kparams_and_docs =
+    let params_and_docs =
       Array.map2_exn names lowereds
         ~f:
           (Option.map2 ~f:(fun name lowered ->
-               let kparams, doc = Syntax.compile_proc ~name idx_params lowered in
-               ((kparams, name), doc)))
+               let params, doc = Syntax.compile_proc ~name idx_params lowered in
+               ((params, name), doc)))
     in
-    let all_proc_docs = List.filter_map (Array.to_list kparams_and_docs) ~f:(Option.map ~f:snd) in
+    let all_proc_docs = List.filter_map (Array.to_list params_and_docs) ~f:(Option.map ~f:snd) in
     let final_doc = PPrint.(separate hardline all_proc_docs) in
     let cuda_includes =
       {|#include <cuda_fp16.h>
@@ -905,8 +905,8 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
     in
     let ptx = cuda_to_ptx ~name source in
     let traced_stores = Array.map lowereds ~f:(Option.map ~f:(fun l -> l.Low_level.traced_store)) in
-    let kparams_and_names = Array.map kparams_and_docs ~f:(Option.map ~f:fst) in
-    { traced_stores; ptx; kparams_and_names; bindings }
+    let params_and_names = Array.map params_and_docs ~f:(Option.map ~f:fst) in
+    { traced_stores; ptx; params_and_names; bindings }
 
   let get_global_run_id =
     let next_id = ref 0 in
@@ -915,7 +915,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
       if !next_id < 0 then next_id := 0;
       !next_id
 
-  let link_proc ~prior_context ~name ~(kparams : (string * kparam_source) list) ~ctx_arrays
+  let link_proc ~prior_context ~name ~(params : (string * param_source) list) ~ctx_arrays
       lowered_bindings run_module =
     let func = Cu.Module.get_function run_module ~name in
     let stream = prior_context.stream in
@@ -929,13 +929,13 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
         "on",
         stream_name,
         (log_id : int),
-        (kparams : (string * kparam_source) list)];
+        (params : (string * param_source) list)];
       let module S = Cu.Stream in
       let args : S.kernel_param list =
         (* TODO: should we prohibit or warn about local-only tensors that are in
            prior_context.ctx_arrays? *)
-        List.map kparams ~f:(function
-          | _name, Kparam_ptr tn ->
+        List.map params ~f:(function
+          | _name, Param_ptr tn ->
               let arr = Option.value_exn ~here:[%here] @@ Map.find ctx_arrays tn in
               S.Tensor arr
           | _name, Log_file_name -> S.Int log_id
@@ -985,7 +985,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
       List.map idx_params ~f:(fun s -> (s, ref 0))
     in
     let task =
-      link_proc ~prior_context ~name:code.name ~kparams:code.kparams ~ctx_arrays lowered_bindings
+      link_proc ~prior_context ~name:code.name ~params:code.params ~ctx_arrays lowered_bindings
         run_module
     in
     (lowered_bindings, task)
@@ -1000,11 +1000,11 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
     let run_module = Cu.Module.load_data_ex code_batch.ptx (run_options ()) in
     prior_context.stream.device.dev.set_builtins_in run_module;
     let procs =
-      Array.mapi code_batch.kparams_and_names ~f:(fun i pns ->
+      Array.mapi code_batch.params_and_names ~f:(fun i pns ->
           Option.value ~default:None
-          @@ Option.map2 pns ctx_arrays.(i) ~f:(fun (kparams, name) ctx_arrays ->
+          @@ Option.map2 pns ctx_arrays.(i) ~f:(fun (params, name) ctx_arrays ->
               let task =
-                link_proc ~prior_context ~name ~kparams ~ctx_arrays lowered_bindings run_module
+                link_proc ~prior_context ~name ~params ~ctx_arrays lowered_bindings run_module
               in
               Some task))
     in

arrayjit/lib/metal_backend.ml

@@ -387,7 +387,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
     metal_source : string; (* Store source, compile during link if not already compiled *)
     compiled_code : Me.Library.t option array; (* Store compiled code per device *)
     func_name : string;
-    kparams : (string * kparam_source) list;
+    params : (string * param_source) list;
     bindings : Indexing.unit_bindings;
     traced_store : Low_level.traced_store;
   }
@@ -396,7 +396,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
   type code_batch = {
     metal_source : string; (* Store combined source *)
     compiled_code : Me.Library.t option array; (* Store compiled code per device *)
-    funcs : (string * (string * kparam_source) list) option array; (* func_name * kparams *)
+    funcs : (string * (string * param_source) list) option array; (* func_name * params *)
     bindings : Indexing.unit_bindings;
     traced_stores : Low_level.traced_store option array;
   }
@@ -671,7 +671,7 @@ module Fresh : Ir.Backend_impl.Lowered_backend = struct
     end)) in
     let idx_params = Indexing.bound_symbols bindings in
     (* Add Metal address space qualifiers *)
-    let kparams, proc_doc = Syntax.compile_proc ~name idx_params lowered in
+    let params, proc_doc = Syntax.compile_proc ~name idx_params lowered in
     let metal_includes = {|#include <metal_stdlib>
 using namespace metal;|} in
     let source =
@@ -683,7 +683,7 @@ using namespace metal;|} in
       compiled_code = Array.create ~len:num_devs None;
       (* One slot per device *)
       func_name = name;
-      kparams;
+      params;
       bindings;
       traced_store = lowered.traced_store;
     }
@@ -697,8 +697,8 @@ using namespace metal;|} in
       Array.map2_exn names lowereds
         ~f:
           (Option.map2 ~f:(fun name lowered ->
-               let kparams, doc = Syntax.compile_proc ~name idx_params lowered in
-               ((name, kparams), doc)))
+               let params, doc = Syntax.compile_proc ~name idx_params lowered in
+               ((name, params), doc)))
     in
     let all_proc_docs = List.filter_map (Array.to_list funcs_and_docs) ~f:(Option.map ~f:snd) in
     let final_doc = PPrint.(separate hardline all_proc_docs) in
@@ -720,7 +720,7 @@ using namespace metal;|} in
     }
 
   let%debug4_sexp link_proc ~prior_context ~library ~func_name
-      ~(kparams : (string * kparam_source) list) ~lowered_bindings ~(ctx_arrays : buffer_ptr Tn.t_map)
+      ~(params : (string * param_source) list) ~lowered_bindings ~(ctx_arrays : buffer_ptr Tn.t_map)
       : Task.t =
     let stream = prior_context.stream in
     let device = stream.device.dev in
@@ -740,12 +740,12 @@ using namespace metal;|} in
         Me.ComputeCommandEncoder.set_compute_pipeline_state encoder pso;
 
         (* Set arguments *)
-        List.iteri kparams ~f:(fun index (_p_name, p_source) ->
+        List.iteri params ~f:(fun index (_p_name, p_source) ->
             match p_source with
-            | Kparam_ptr tn when Map.mem ctx_arrays tn ->
+            | Param_ptr tn when Map.mem ctx_arrays tn ->
                 let buffer = Map.find_exn ctx_arrays tn in
                 Me.ComputeCommandEncoder.set_buffer encoder ~index buffer
-            | Kparam_ptr tn when Tn.known_constant tn && Tn.is_hosted_force tn 48 ->
+            | Param_ptr tn when Tn.known_constant tn && Tn.is_hosted_force tn 48 ->
                 let buffer =
                   Hashtbl.find_or_add stream.device.device_buffer_cache tn ~default:(fun () ->
                       get_buffer_for_ptr device ~size_in_bytes:(Lazy.force tn.size_in_bytes)
@@ -754,9 +754,9 @@ using namespace metal;|} in
                       @@ Lazy.force tn.array)
                 in
                 Me.ComputeCommandEncoder.set_buffer encoder ~index buffer
-            | Kparam_ptr tn ->
+            | Param_ptr tn ->
                 failwith
-                  [%string "Kparam_ptr %{Tn.debug_name tn} not found in ctx_arrays for %{func_name}"]
+                  [%string "Param_ptr %{Tn.debug_name tn} not found in ctx_arrays for %{func_name}"]
             | Static_idx s ->
                 let value = !(Indexing.find_exn lowered_bindings s) in
                 let size = Ctypes.sizeof Ctypes.int in
@@ -803,7 +803,7 @@ using namespace metal;|} in
       List.map (Indexing.bound_symbols code.bindings) ~f:(fun s -> (s, ref 0))
     in
     let task =
-      link_proc ~prior_context ~library ~func_name:code.func_name ~kparams:code.kparams
+      link_proc ~prior_context ~library ~func_name:code.func_name ~params:code.params
         ~lowered_bindings ~ctx_arrays
     in
     (lowered_bindings, task)
@@ -817,9 +817,9 @@ using namespace metal;|} in
 
     let tasks =
       Array.mapi code_batch.funcs ~f:(fun i func_opt ->
-          Option.bind func_opt ~f:(fun (func_name, kparams) ->
+          Option.bind func_opt ~f:(fun (func_name, params) ->
               Option.map ctx_arrays_opts.(i) ~f:(fun ctx_arrays ->
-                  link_proc ~prior_context ~library ~func_name ~kparams ~lowered_bindings ~ctx_arrays)))
+                  link_proc ~prior_context ~library ~func_name ~params ~lowered_bindings ~ctx_arrays)))
     in
     (lowered_bindings, tasks)
 end