Make in-place ops DCE-able after update_aliases

thunder/__init__.py

@@ -37,7 +37,6 @@
     SHARP_EDGES_OPTIONS,
 )
 from thunder.core.proxies import TensorProxy
-from thunder.core.pytree import tree_flatten
 from thunder.core.recipe import Recipe, Plugin
 from thunder.core.symbol import has_tags
 from thunder.core.trace import (
@@ -56,6 +55,7 @@
     wrap_return_value_together_with_arguments,
 )
 from thunder.core.update_aliases import insert_alias_updates
+from thunder.core.utils import encode_alias_tensor_indices
 from thunder.executors.torch_autograd import connect_to_autograd
 import thunder.extend as extend
 from thunder.extend import Executor, add_default_executor
@@ -405,37 +405,6 @@ def jit(
     cs = CompileStats()
     weakref_cs = weakref.ref(cs)
 
-    def _alias_tensor_of_args_kwargs_dict(*args, **kwargs) -> dict[int, list[int]]:
-        flat_args, _ = tree_flatten((args, kwargs))
-        data_ptr_to_tensor_group_index = {}
-        tensor_group_index_to_tensor_indices = defaultdict(list)
-        for idx, t in enumerate(flat_args):
-            # Using type(t) is pytorch.Tensor as TensorSubclasses don't support calling
-            # data_ptr().
-            # Eg. RuntimeError: Attempted to access the data pointer on an invalid python storage. (data_ptr access on TensorSubclass)
-            #
-            # isinstance(t, pytorch.Tensor) or pytorch.is_tensor(t) will match all Tensor objects including
-            # subclasses.
-            if type(t) is pytorch.Tensor and t.layout is pytorch.strided:
-                data_ptr = t.untyped_storage().data_ptr()
-                if data_ptr not in data_ptr_to_tensor_group_index:
-                    data_ptr_to_tensor_group_index[data_ptr] = len(data_ptr_to_tensor_group_index)
-                tgi = data_ptr_to_tensor_group_index[data_ptr]
-                tensor_group_index_to_tensor_indices[tgi].append(idx)
-        return tensor_group_index_to_tensor_indices
-
-    def _alias_tensor_of_args_kwargs(*args, **kwargs) -> str:
-        """If no aliases found, empty string, otherwise, aliases are comma separated, groups are hyphen separated."""
-
-        alias_indices = []
-        for k, v in _alias_tensor_of_args_kwargs_dict(*args, **kwargs).items():
-            if len(v) > 1:
-                s = ",".join(f"{i}" for i in v)
-                alias_indices.append(s)
-        if not alias_indices:
-            return ""
-        return "-".join(alias_indices)
-
     def acquire_initial_trace(fn, args, kwargs, cd, cs, ad_hoc_executor):
         with compile_data_and_stats(cd, cs):
             # Acquires the trace OR inlines the trace into an existing trace and
@@ -488,16 +457,10 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
             computation_trc = remove_context_manager_prims_from_trace(computation_trc)
             computation_traces.append(computation_trc)
 
-            alias_tensor_indices_str = cache_info.get("alias_tensor_indices", "")
-            alias_tensor_indices: list[list[int]] = [
-                [int(i) for i in s.split(",")] for s in alias_tensor_indices_str.split("-") if s != ""
-            ]
-
-            if not compile_options.get("skip_inplace_alias_updates", False):
-                aliased_trace = insert_alias_updates(computation_trc, alias_tensor_indices)
-                if aliased_trace is not computation_trc:
-                    computation_traces.append(aliased_trace)
-                    computation_trc = computation_traces[-1]
+            aliased_trace = insert_alias_updates(computation_trc)
+            if aliased_trace is not computation_trc:
+                computation_traces.append(aliased_trace)
+                computation_trc = computation_traces[-1]
 
             cs.last_trace_tracing_stop = time.perf_counter_ns()
 
@@ -674,7 +637,7 @@ def populate_cache_info(cache_info, *args, **kwargs):
         # It however would require the computation trace to interact with `cache_info`,
         # which seems to break the consistency of cache_info, leading to a failure in cache_info check.
         if not compile_options.get("skip_inplace_alias_updates", False):
-            cache_info["alias_tensor_indices"] = _alias_tensor_of_args_kwargs(*args, **kwargs)
+            cache_info["alias_tensor_indices"] = encode_alias_tensor_indices(*args, **kwargs)
 
         # Store the `is_grad_enabled` state of PyTorch. This is used by vjp transform
         # to treat certain Symbols as constant.

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

thunder/core/jit_ext.py

@@ -835,8 +835,11 @@ def core_of_forward(*args, **kwargs):
 
     from thunder.core.update_aliases import insert_alias_updates
 
-    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)
+    # 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)
+
+    aliased_trace_of_augmented_fwd = insert_alias_updates(trace_of_augmented_fwd)
 
     # Backward definition
     custom_backward = custom_autograd_function_cls.backward
@@ -869,8 +872,11 @@ def core_of_forward(*args, **kwargs):
     )
     bwd_trace_impl.args = tuple(ctx_proxy.saved_consts + ctx_proxy.saved_tensors + grads)
 
-    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)
+    # 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)
+
+    aliased_bwd_trace_impl = insert_alias_updates(bwd_trace_impl)
 
     @wraps(bwd_trace_impl.python_callable())
     def bwd_impl_callable(*args, **kwargs):
@@ -951,8 +957,11 @@ def _generate_random_str_id() -> str:
 
     from thunder.core.update_aliases import insert_alias_updates
 
-    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)
+    # 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)
+
+    aliased_aug_fwd_trace = insert_alias_updates(aug_fwd_trace)
 
     trace_of_forward = from_trace(aliased_aug_fwd_trace)
     for bsym in aug_fwd_trace.bound_symbols:
@@ -988,10 +997,11 @@ 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)
+    aliased_bwd_trace = insert_alias_updates(bwd_trace)
 
     @wraps(forward)
     def grad_transform(*args, **kwargs):

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

@@ -1,11 +1,13 @@
 from functools import reduce, partial
 
-from thunder.core.compile_data import using_symbolic_values
+import thunder
+from thunder.core.compile_data import get_compile_data, using_symbolic_values
 import thunder.core.prims as prims
 from thunder.core.proxies import TensorProxy, variableify, unvariableify
 from thunder.core.pytree import tree_flatten
 from thunder.core.symbol import BoundSymbol, BoundSymbolTag, has_tags
 from thunder.core.trace import from_trace, tracectx, TraceCtx as Trace, TraceProvenance, VariableInterface
+from thunder.core.utils import parse_alias_tensor_indices
 
 
 def _update_swap_map(swap_map, old_alias, new_alias):
@@ -39,6 +41,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_
     )
@@ -51,8 +54,6 @@ def _is_view_creation_op(bsym):
 
 
 def _involves_viewed_args(bsym, viewed):
-    if bsym.sym.id == prims.PrimIDs.RETURN:
-        return False
     return any(isinstance(p, TensorProxy) and variableify(p) in viewed for p in bsym.flat_proxy_args)
 
 
@@ -131,10 +132,17 @@ def replace_args_with_alias_map(
     return no_implicit_alias_trace, view_groups
 
 
-def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[list[int]]) -> Trace:
+def insert_alias_updates(computation_trace: Trace) -> Trace:
+    cd = get_compile_data()
+    if cd is not None and cd.compile_options.get("skip_inplace_alias_updates", False):
+        return computation_trace
+
     if not any(_is_inplace_op(bsym) for bsym in computation_trace.bound_symbols):
         return computation_trace
 
+    alias_tensor_indices_str = thunder._get_cache_info().get("alias_tensor_indices", "")
+    alias_tensor_indices = parse_alias_tensor_indices(alias_tensor_indices_str)
+
     swap_map = dict()
     bsyms = []
 
@@ -148,7 +156,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)
@@ -169,7 +178,8 @@ def insert_alias_updates(computation_trace: Trace, alias_tensor_indices: list[li
         if _is_inplace_op(bsym) or _is_view_creation_op(bsym) or _involves_viewed_args(bsym, viewed):
             in_tensors = list(map(variableify, filter(lambda p: isinstance(p, TensorProxy), bsym.flat_proxy_args)))
             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]}
             else:
                 in_tensors = set(in_tensors)
             out_tensors = set(map(variableify, filter(lambda p: isinstance(p, TensorProxy), bsym.flat_proxy_outs)))

thunder/core/utils.py

@@ -1262,3 +1262,27 @@ def create_python_callable_from_bsym(bsym: BoundSymbolInterface) -> str:
         prims.python_return(bsym.output)
 
     return trace.python(include_decorators=False)
+
+
+def parse_alias_tensor_indices(alias_tensor_indices_str: str) -> list[list[int]]:
+    return [[int(i) for i in s.split(",")] for s in alias_tensor_indices_str.split("-") if s != ""]
+
+
+def encode_alias_tensor_indices(*args, **kwargs) -> str:
+    flat_args, _ = tree_flatten((args, kwargs))
+    data_ptr_to_tensor_indices = defaultdict(list)
+
+    for idx, t in enumerate(flat_args):
+        # Using type(t) is torch.Tensor as TensorSubclasses don't support calling data_ptr().
+        # Eg. RuntimeError: Attempted to access the data pointer on an invalid python storage. (data_ptr access on TensorSubclass)
+        #
+        # isinstance(t, torch.Tensor) or torch.is_tensor(t) will match all Tensor objects including subclasses.
+        if type(t) is torch.Tensor and t.layout is torch.strided:
+            data_ptr = t.untyped_storage().data_ptr()
+            data_ptr_to_tensor_indices[data_ptr].append(idx)
+
+    alias_indices = []
+    for indices in data_ptr_to_tensor_indices.values():
+        if len(indices) > 1:
+            alias_indices.append(",".join(str(idx) for idx in indices))
+    return "-".join(alias_indices)

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
@@ -122,7 +123,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)
     extrace = dce(extrace)
+
     #
     # Step 2 Fusion executors can transform the trace
     #

thunder/executors/torchex.py

@@ -2367,7 +2367,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)