Make in-place ops DCE-able after update_aliases

thunder/__init__.py

@@ -530,6 +530,7 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
                 prologue_trc,
                 executors_list=(pythonex,),
                 use_del_last_used=False,
+                alias_tensor_indices=alias_tensor_indices,
             )
             prologue_trc = prologue_traces[-1]
             pro = prologue_trc.python_callable(include_decorators=False)
@@ -553,7 +554,7 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
             if requires_grad:
                 from thunder.transforms.autodiff import grad_transform_on_trace
 
-                computation_trc = grad_transform_on_trace(computation_trc)
+                computation_trc = grad_transform_on_trace(computation_trc, alias_tensor_indices)
                 computation_traces.append(computation_trc)
 
             from thunder.executors.passes import _transform_for_operator_executor_execution
@@ -569,6 +570,7 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
                 computation_trc,
                 executors_list=cd.executors_list,
                 use_del_last_used=False,
+                alias_tensor_indices=alias_tensor_indices,
             )
             computation_trc = extraces[-1]
 

thunder/common.py

@@ -619,7 +619,7 @@ def wait_for_future(f: FutureTensorProxy) -> TensorProxy:
             # TODO Stop calling this here and make it a separate trace in the sequence
             #   of traces
             if use_dce:
-                trace = dce(trace)
+                trace = dce(trace, keep_inplace_ops=True)
 
         finally:
             # Resets contexts
@@ -644,6 +644,7 @@ def transform_for_execution(
     *,
     only_execute_prims=False,
     use_del_last_used=True,
+    alias_tensor_indices: list[list[int]] | None = None,
 ) -> list[TraceCtx]:
     traces: list[TraceCtx] = []
 
@@ -652,7 +653,7 @@ def transform_for_execution(
     # cse_trace = cse(dce_trace)
     # traces.append(cse_trace)
 
-    extrace = executors.passes.transform_for_execution(trace, executors_list)
+    extrace = executors.passes.transform_for_execution(trace, executors_list, alias_tensor_indices)
 
     traces.append(extrace)
 

thunder/core/jit_ext.py

@@ -835,6 +835,10 @@ def core_of_forward(*args, **kwargs):
 
     from thunder.core.update_aliases import insert_alias_updates
 
+    # Copy attributes needed for TensorProxy name construction
+    trace_of_augmented_fwd.name_ctr = get_jit_ctx().computation_trace.name_ctr
+    trace_of_augmented_fwd.names = set(get_jit_ctx().computation_trace.names)
+
     alias_tensor_indices = [[i] for i in range(len(trace_of_augmented_fwd.args))]
     aliased_trace_of_augmented_fwd = insert_alias_updates(trace_of_augmented_fwd, alias_tensor_indices)
 
@@ -869,6 +873,10 @@ def core_of_forward(*args, **kwargs):
     )
     bwd_trace_impl.args = tuple(ctx_proxy.saved_consts + ctx_proxy.saved_tensors + grads)
 
+    # Copy attributes needed for TensorProxy name construction
+    bwd_trace_impl.name_ctr = get_jit_ctx().computation_trace.name_ctr
+    bwd_trace_impl.names = set(get_jit_ctx().computation_trace.names)
+
     alias_tensor_indices = [[i] for i in range(len(bwd_trace_impl.args))]
     aliased_bwd_trace_impl = insert_alias_updates(bwd_trace_impl, alias_tensor_indices)
 
@@ -951,6 +959,10 @@ def _generate_random_str_id() -> str:
 
     from thunder.core.update_aliases import insert_alias_updates
 
+    # Copy attributes needed for TensorProxy name construction
+    aug_fwd_trace.name_ctr = get_jit_ctx().computation_trace.name_ctr
+    aug_fwd_trace.names = set(get_jit_ctx().computation_trace.names)
+
     alias_tensor_indices = [[i] for i in range(len(aug_fwd_trace.args))]
     aliased_aug_fwd_trace = insert_alias_updates(aug_fwd_trace, alias_tensor_indices)
 
@@ -988,7 +1000,9 @@ def forward(*args, **kwargs):
     ]
     bwd_trace.bound_symbols = bwd_unpack_bsyms + bwd_trace.bound_symbols
 
-    from thunder.core.update_aliases import insert_alias_updates
+    # Copy attributes needed for TensorProxy name construction
+    bwd_trace.name_ctr = get_jit_ctx().computation_trace.name_ctr
+    bwd_trace.names = set(get_jit_ctx().computation_trace.names)
 
     alias_tensor_indices = [[i] for i in range(len(bwd_trace.args))]
     aliased_bwd_trace = insert_alias_updates(bwd_trace, alias_tensor_indices)

thunder/core/prims.py

@@ -4333,7 +4333,7 @@ def copy__meta(
     return TensorProxy(like=copy_to)
 
 
-copy_ = make_prim(PrimIDs.COPY_, "copy_", meta=copy__meta, tags=(OpTags.DONT_DCE,))
+copy_ = make_prim(PrimIDs.COPY_, "copy_", meta=copy__meta, tags=(OpTags.IN_PLACE,))
 
 
 def bitcast_meta(

thunder/core/transform_common.py

@@ -142,7 +142,7 @@ def keep_or_swap(p):
 #   that only produce non-proxy objects
 # NOTE needed_proxies is an in/out argument, it takes an initial set of Variables you want to keep, and return
 #   all the needed proxies of the input trace
-def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace:
+def dce(trace: Trace, needed_proxies: None | set[Variable] = None, keep_inplace_ops: bool = False) -> Trace:
     start_time_ns = time.perf_counter_ns()
 
     producer_map: ProxyDict = producers(trace)
@@ -159,6 +159,8 @@ def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace:
         # Preserves symbols that should never be collected
         if has_tags(bsym, {prims.OpTags.DONT_DCE}):
             needed = True
+        elif keep_inplace_ops and has_tags(bsym, {prims.OpTags.IN_PLACE}):
+            needed = True
         else:
             needed = False
 

thunder/core/update_aliases.py

@@ -39,6 +39,7 @@ def _get_new_aliases(aliases, trace):
 
 
 def _is_inplace_op(bsym):
+    # TODO: Handle higher order bsyms containing inplace ops
     return (bsym.sym.tags and prims.OpTags.IN_PLACE in bsym.sym.tags) or (
         bsym.subsymbols and bsym.subsymbols[-1].sym.id == prims.PrimIDs.COPY_
     )
@@ -67,7 +68,7 @@ def _can_be_reshaped(arg, arg_to_replace):
 
 def replace_args_with_alias_map(
     computation_trace: Trace,
-    alias_tensor_indices: list[list[int]],
+    alias_tensor_indices: list[list[int]] | None = None,
 ) -> tuple[Trace, list[set[VariableInterface]]]:
     if not alias_tensor_indices:
         return computation_trace, []
@@ -140,7 +141,7 @@ def _helper(alias):
     return list(map(_helper, aliases))
 
 
-def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[list[int]]) -> Trace:
+def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[list[int]] | None = None) -> Trace:
     if not any(_is_inplace_op(bsym) for bsym in computation_trace.bound_symbols):
         return computation_trace
 
@@ -157,7 +158,8 @@ def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[li
     for bsym in computation_trace.bound_symbols:
         if _is_inplace_op(bsym) or _is_view_creation_op(bsym):
             # only interested in the input which is modified by the inplace op
-            in_tensor = variableify(bsym.flat_proxy_args[0])
+            mutated_or_aliased_index = 1 if bsym.sym.id == prims.PrimIDs.COPY_ else 0
+            in_tensor = variableify(bsym.flat_proxy_args[mutated_or_aliased_index])
             out_tensors = set(map(variableify, filter(lambda p: isinstance(p, TensorProxy), bsym.flat_proxy_outs)))
             if _is_inplace_op(bsym):
                 inplace_inputs.add(in_tensor)
@@ -180,10 +182,11 @@ def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[li
         if (
             _is_inplace_op(bsym)
             or _is_view_creation_op(bsym)
-            or (bsym.sym.id != prims.PrimIDs.RETURN and _involves_viewed_args(set(unswapped_in_tensors), viewed))
+            or _involves_viewed_args(set(unswapped_in_tensors), viewed)
         ):
             if _is_inplace_op(bsym) and in_tensors:
-                in_tensors = {in_tensors[0]}
+                mutated_index = 1 if bsym.sym.id == prims.PrimIDs.COPY_ else 0
+                in_tensors = {in_tensors[mutated_index]}
                 unswapped_in_tensors = {unswapped_in_tensors[0]}
             else:
                 in_tensors = set(in_tensors)

thunder/executors/nvfuserex_impl.py

@@ -646,7 +646,7 @@ def __init__(self):
         super().__init__("nvfuser", version=nvfuser.version())
 
         # TODO: Replace this with a query to a compile option
-        self._use_rematerialization = True
+        self._use_rematerialization = False
 
         fuel_str = os.getenv("NVFUSER_OPTIMIZATION_FUEL")
         if fuel_str:

thunder/executors/passes.py

@@ -10,6 +10,7 @@
 from thunder.core.trace import from_trace, TraceProvenance
 from thunder.core.trace_interpreter import TraceSubstitutionProcessor
 from thunder.core.transform_common import dce
+from thunder.core.update_aliases import insert_alias_updates
 from thunder.core.utils import ProxyDict
 from thunder.executors.pythonex import clear_mutable_collection
 from thunder.extend import Executor, get_always_executors, OperatorExecutor, FusionExecutor
@@ -104,7 +105,9 @@ def process_bsym(self, bsym: BoundSymbol) -> None:
     return extrace
 
 
-def transform_for_execution(trace: TraceCtx, executors_list: Sequence[Executor]) -> TraceCtx:
+def transform_for_execution(
+    trace: TraceCtx, executors_list: Sequence[Executor], alias_tensor_indices: list[list[int]] | None = None
+) -> TraceCtx:
     import torch
 
     start_time_ns = time.perf_counter_ns()
@@ -122,7 +125,11 @@ def transform_for_execution(trace: TraceCtx, executors_list: Sequence[Executor])
     # Step 1 Performs execution transforms
     #
     extrace = _transform_for_operator_executor_execution(trace, executors_list)
+    # Insert alias updates before DCE for bsyms exposed by decomposition
+    # Inserted prims.update_aliases will be handled in Step 3
+    extrace = insert_alias_updates(extrace, alias_tensor_indices)
     extrace = dce(extrace)
+
     #
     # Step 2 Fusion executors can transform the trace
     #

thunder/executors/torchex.py

@@ -2371,7 +2371,7 @@ def _copy__impl(copy_from, copy_to, grad_enabled):
 
 
 copy_ = ex.register_operator(
-    "copy_", meta=prims.copy_, tags=(prims.OpTags.DONT_DCE,), fn=_copy__impl, module=torch.Tensor
+    "copy_", meta=prims.copy_, tags=(prims.OpTags.IN_PLACE,), fn=_copy__impl, module=torch.Tensor
 )
 _register_implementation(prims.copy_, copy_, checker=_always_executable)
 

thunder/tests/test_inplace_copy.py

@@ -42,16 +42,16 @@ def test_prim_inplace_copy_bwd(executor, device, dtype):
     def torch_foo(x, y):
         z = x * y
         z = z * x
-        x.copy_(z)
+        o = x.copy_(z)
         p = y * y
-        return p
+        return p, o
 
     def foo(x, y):
         z = x * y
         z = z * x
-        thunder.core.prims.copy_(z, x, grad_enabled=True)
+        o = thunder.core.prims.copy_(z, x, grad_enabled=True)
         p = y * y
-        return p
+        return p, o
 
     traced_nvfuser_foo = executor.make_callable(foo)
 
@@ -72,11 +72,11 @@ def foo(x, y):
     )
     custom_comparator(a, a1)
 
-    g = torch.ones_like(thunder_result)
-    thunder_result.backward(g)
+    g = torch.ones_like(thunder_result[0])
+    thunder_result[0].backward(g)
 
-    g1 = torch.ones_like(torch_result)
-    torch_result.backward(g1)
+    g1 = torch.ones_like(torch_result[0])
+    torch_result[0].backward(g1)
     assert_close(g, g1)
     assert_close(b.grad, b1.grad)
 
@@ -131,7 +131,7 @@ def func2(x, y):
         return y, o1, o2
 
     for foo in (func1, func2):
-        traced_foo = executor.make_callable(foo)
+        traced_foo = executor.make_callable(foo, skip_inplace_alias_updates=True)
 
         tdtype = ttorch.to_torch_dtype(dtype)
         a = make_tensor((4, 4), device=device, dtype=tdtype)