Merge pull request #368 from ahrefs/feature/tensor-rootness-check

lukstafi · web-flow · commit 7fe3406b69de · 2025-08-24T11:40:04.000+02:00
Improve tensor rootness safety checks: precise checks in Tensor.consume_forward_code and Tensor.consume_backprop_code
diff --git a/arrayjit/lib/assignments.ml b/arrayjit/lib/assignments.ml
@@ -113,8 +113,9 @@ 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 =
+(** In the second set, returns the nodes that are not read from after being written to. In the first
+    set, returns the nodes that are ever read from. *)
+let%debug3_sexp collect_nodes_guess_output (asgns : t) : Tn.t_set * Tn.t_set =
   let open Utils.Set_O in
   let empty = Set.empty (module Tn) in
   let one = Set.singleton (module Tn) in
@@ -137,7 +138,7 @@ let%debug3_sexp guess_output_nodes (asgns : t) : Tn.t_set =
     | Set_vec_unop { lhs; rhs; _ } -> (of_node rhs, one lhs)
     | Fetch { array; _ } -> (empty, one array)
   in
-  snd @@ loop asgns
+  loop asgns
 
 let sequential l =
   Option.value ~default:Noop @@ List.reduce l ~f:(fun sts another_st -> Seq (sts, another_st))
diff --git a/lib/tensor.ml b/lib/tensor.ml
@@ -601,18 +601,22 @@ let consume_forward_code t =
     @@ Session_error
          ( "Tensor.consume_forward_code: tensor is not a root for tnode: " ^ Tn.debug_name t.value,
            Some t );
-  (* FIXME(#321): this is too aggressive, instead we should check if the code contains any
-     non-embedded nodes that are embedded nodes of the other roots. *)
-  let unsafe_roots =
-    Map.data session_state.forward_roots
-    |> List.filter ~f:(fun r -> not (List.is_empty r.children || r.id = t.id))
+  (* Check if any non-embedded descendants of t are embedded in other roots *)
+  let all_read = fst @@ Asgns.collect_nodes_guess_output t.forward.asgns in
+  let non_embedded_descendants = Set.diff all_read t.forward.embedded_nodes in
+  let other_roots =
+    Map.data session_state.forward_roots |> List.filter ~f:(fun r -> r.id <> t.id)
   in
-  if not @@ List.is_empty unsafe_roots then
+  let conflicting_roots =
+    List.filter other_roots ~f:(fun root ->
+        not (Set.is_empty (Set.inter non_embedded_descendants root.forward.embedded_nodes)))
+  in
+  if not @@ List.is_empty conflicting_roots then
     raise
     @@ Session_error
          ( [%string
              {|Tensor.consume_forward_code for %{debug_name t}:
-found potentially unsafe roots: %{String.concat ~sep:", " @@ List.map ~f:debug_name unsafe_roots}|}],
+found conflicting roots with shared non-embedded descendants: %{String.concat ~sep:", " @@ List.map ~f:debug_name conflicting_roots}|}],
            Some t );
   remove_fwd_root t;
   t.forward
@@ -630,16 +634,26 @@ let consume_backprop_code t =
     raise
     @@ Session_error
          ("Tensor.consume_backprop_code: tensor is not a root for tnode: " ^ debug_grad t, Some t);
-  let unsafe_roots =
-    Map.data session_state.backprop_roots
-    |> List.filter ~f:(fun r -> not (List.is_empty r.children || r.id = t.id))
-  in
-  if not @@ List.is_empty unsafe_roots then
+  (* Check if any non-embedded grad descendants of t are embedded in other roots *)
+  let all_read = fst @@ Asgns.collect_nodes_guess_output diff.backprop.asgns in
+  let non_embedded_grad_descendants = Set.diff all_read diff.backprop.embedded_nodes in
+  let other_roots =
+    Map.data session_state.backprop_roots |> List.filter ~f:(fun r -> r.id <> t.id)
+  in
+  let conflicting_roots =
+    List.filter other_roots ~f:(fun root ->
+        match root.diff with
+        | Some rdiff ->
+            not
+              (Set.is_empty (Set.inter non_embedded_grad_descendants rdiff.backprop.embedded_nodes))
+        | None -> false)
+  in
+  if not @@ List.is_empty conflicting_roots then
     raise
     @@ Session_error
          ( [%string
              {|Tensor.consume_backprop_code for %{debug_grad t}:
-found potentially unsafe roots: %{String.concat ~sep:", " @@ List.map ~f:debug_name unsafe_roots}|}],
+found conflicting roots with shared non-embedded grad descendants: %{String.concat ~sep:", " @@ List.map ~f:debug_grad conflicting_roots}|}],
            Some t );
   remove_bprop_root t;
   diff.backprop
diff --git a/lib/train.ml b/lib/train.ml
@@ -317,7 +317,7 @@ let to_routine (type buffer_ptr dev runner event optimize_ctx)
        and type event = event
        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;
+  if hosted then Set.iter (snd @@ Asgns.collect_nodes_guess_output comp.Asgns.asgns) ~f:set_hosted;
   Backend.link context @@ Backend.compile context.optimize_ctx ?name bindings comp
 
 (** [init_params] initializes the parameters of [t], via running their forward code or copying from
