Auto-set hosted for Train.to_routine; fix ambiguous mem mode provenances

lukstafi · lukstafi · commit acb41d23296a · 2025-06-29T11:20:57.000+02:00
diff --git a/arrayjit/lib/assignments.ml b/arrayjit/lib/assignments.ml
@@ -117,6 +117,27 @@ let%debug3_sexp context_nodes ~(use_host_memory : 'a option) (asgns : t) : Tn.t_
   in
   loop asgns
 
+(** Returns the nodes that are not read from after being written to. *)
+let%debug3_sexp guess_output_nodes (asgns : t) : Tn.t_set =
+  let open Utils.Set_O in
+  let empty = Set.empty (module Tn) in
+  let one = Set.singleton (module Tn) in
+  let of_node = function Node rhs -> one rhs | Merge_buffer _ -> empty in
+  let rec loop = function
+    | Noop -> (empty, empty)
+    | Seq (t1, t2) ->
+        let i1, o1 = loop t1 in
+        let i2, o2 = loop t2 in
+        (i1 + i2, o1 + o2 - (i1 + i2))
+    | Block_comment (_, t) -> loop t
+    | Accum_unop { lhs; rhs; _ } -> (of_node rhs, one lhs)
+    | Accum_binop { lhs; rhs1; rhs2; _ } -> (of_node rhs1 + of_node rhs2, one lhs)
+    | Accum_ternop { lhs; rhs1; rhs2; rhs3; _ } ->
+        (of_node rhs1 + of_node rhs2 + of_node rhs3, one lhs)
+    | Fetch { array; _ } -> (empty, one array)
+  in
+  snd @@ loop asgns
+
 let sequential l =
   Option.value ~default:Noop @@ List.reduce l ~f:(fun sts another_st -> Seq (sts, another_st))
 
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -455,7 +455,7 @@ module Raise_backend (Device : Lowered_backend) : Backend = struct
         let data = Hashtbl.find_exn device.cross_stream_candidates key in
         Map.add_exn ctx_arrays ~key ~data)
       else (
-        Tn.update_memory_sharing key Tn.Per_stream 41;
+        Tn.update_memory_sharing key Tn.Per_stream 410;
         Hashtbl.remove device.cross_stream_candidates key;
         add_new ()))
     else ctx_arrays
diff --git a/lib/train.ml b/lib/train.ml
@@ -69,8 +69,8 @@ let set_on_host ?(from_device = true) (a : Tn.t) =
 let set_materialized (a : Tn.t) = Tn.update_memory_mode a Materialized 28
 
 let set_hosted (a : Tn.t) =
-  if Tn.known_constant a then Tn.update_memory_mode a (Hosted Constant) 41
-  else Tn.update_memory_mode a (Hosted (Changed_on_devices Unset)) 41
+  if Tn.known_constant a then Tn.update_memory_mode a (Hosted Constant) 411
+  else Tn.update_memory_mode a (Hosted (Changed_on_devices Unset)) 412
 
 (** Sets the tensor's value as "fully on host", returns the tensor's forward code with a
     label-derived comment. *)
@@ -346,7 +346,8 @@ let to_routine (type buffer_ptr dev runner event optimize_ctx)
        and type dev = dev
        and type runner = runner
        and type event = event
-       and type optimize_ctx = optimize_ctx) (context : Backend.context) ?name bindings comp =
+       and type optimize_ctx = optimize_ctx) (context : Backend.context) ?(hosted=true) ?name bindings comp =
+  if hosted then Set.iter (Asgns.guess_output_nodes comp.Asgns.asgns) ~f:set_hosted;
   Backend.link context @@ Backend.compile context.optimize_ctx ?name bindings comp
 
 type example_train_result = {
