Fix names conflict when lifting (pytorch#129817)

## Bug description
When pending args that are potentially to be lift [here](https://github.com/pytorch/pytorch/blob/58f346c874a8a982679b4d4f3876602cc05d66d4/torch/_dynamo/output_graph.py#L1866) having same base name, like `contiguous` and `contiguous_1`, the call into [create_graph_input](https://github.com/pytorch/pytorch/blob/58f346c874a8a982679b4d4f3876602cc05d66d4/torch/_dynamo/output_graph.py#L2081) can finally create a name ([here](https://github.com/pytorch/pytorch/blob/58f346c874a8a982679b4d4f3876602cc05d66d4/torch/fx/graph.py#L1008)) that overwrite args to lift. And thus causing a wrong output of graph.

## Reproducing
Below is an reproduceable example,
```python
import logging
from typing import List

import torch
from functorch.compile import aot_module_simplified, make_boxed_func

@torch.library.custom_op("mylib::somefunc_forward", mutates_args=())
def somefunc_forward(
    input_: torch.Tensor,
    weight: torch.Tensor,
    shape: List[int],
) -> torch.Tensor:
    return torch.ones_like(input_)

@somefunc_forward.register_fake
def _(input_, shape, weight):
    return torch.empty_like(input_)

@torch.library.custom_op("mylib::somefunc_backward", mutates_args=())
def somefunc_backward(
    grad_output: torch.Tensor,
    input_: torch.Tensor,
    weight: torch.Tensor,
    shape: List[int],
) -> torch.Tensor:
    print(f"backward.{grad_output.shape=}")
    print(f"backward.{input_.shape=}")
    print(f"backward.{weight.shape=}")
    print(f"backward.{shape=}")
    assert list(weight.shape) == shape
    return torch.ones_like(weight)

@somefunc_backward.register_fake
def _(grad_output, input_, weight, shape):
    return torch.empty_like(weight)

def a_func(grad_output, input_, weight_, shape):
    return torch.ones_like(input_.sum() * weight_)

class SomeFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, normalized_shape):
        ctx.normalized_shape = normalized_shape
        input_ = input.contiguous()
        weight_ = weight.contiguous()
        output = somefunc_forward(input_, weight_, ctx.normalized_shape)
        ctx.save_for_backward(input_, weight_)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input_, weight_ = ctx.saved_tensors
        # grad_weight = a_func(grad_output, input_, weight_, ctx.normalized_shape)
        grad_weight = somefunc_backward(
            grad_output.contiguous(),
            input_,
            weight_,
            ctx.normalized_shape,
        )
        return None, grad_weight, None

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(7))

    def forward(self, x):
        return SomeFunc.apply(x, self.weight, [7])

model = MyModel()
torch._logging.set_logs(dynamo=logging.DEBUG, aot=logging.DEBUG, graph_code=True)

def aot_print_backend(gm, sample_inputs):
    # Forward compiler capture
    def fw(gm, sample_inputs):
        print(f"----- fw")
        gm.print_readable()
        return make_boxed_func(gm.forward)

    # Backward compiler capture
    def bw(gm, sample_inputs):
        print(f"----- bw")
        gm.print_readable()
        return make_boxed_func(gm.forward)

    # Call AOTAutograd
    gm_forward = aot_module_simplified(
        gm, sample_inputs, fw_compiler=fw, bw_compiler=bw
    )
    return gm_forward

model = torch.compile(
    model,
    backend=aot_print_backend,
    dynamic=False,
)
out = model(torch.rand((128, 4, 7)))
out.mean().backward()
```

I can see log that showing calling into create_graph_input like
```log
V0629 02:08:46.839914 8200981504 torch/_dynamo/output_graph.py:2042] [0/0] create_graph_input contiguous (none)
V0629 02:08:46.839998 8200981504 torch/_dynamo/output_graph.py:2042] [0/0] create_graph_input contiguous_1 (none)
```

And the backward graph generate will be like
```log
class GraphModule(torch.nn.Module):
    def forward(self, function_ctx, somefunc_forward_default: "f32[128, 4, 7]", contiguous: "f32[128, 4, 7]", contiguous_1: "f32[7]"):
        contiguous_1 = contiguous
        contiguous_2 = contiguous_1

        # No stacktrace found for following nodes
        _set_grad_enabled = torch._C._set_grad_enabled(False)

         # File: /Users/bytedance/testtorch/test_custom_op_bug.py:61 in backward, code: grad_output.contiguous(),
        contiguous: "f32[128, 4, 7]" = somefunc_forward_default.contiguous();  somefunc_forward_default = None

         # File: /opt/tiger/pytorch/torch/_library/custom_ops.py:506 in __call__, code: return self._opoverload(*args, **kwargs)
        somefunc_backward_default: "f32[7]" = torch.ops.mylib.somefunc_backward.default(contiguous, contiguous_1, contiguous_2, [7]);  contiguous = contiguous_1 = contiguous_2 = None

        # No stacktrace found for following nodes
        _set_grad_enabled_1 = torch._C._set_grad_enabled(True)
        return (None, somefunc_backward_default)
```

The original code of `somefunc_backward` takes a input list of `grad_output`, `input_`, `weight` and `shape`, where `weight` should be shape of `torch.Size([7])`. However, in the graph, `contiguous1` and `contiguous_2` are assigned with `contiguous`, this leads to assertion failure I added in `somefunc_backward`.

## Environment
```log
Collecting environment information...
PyTorch version: 2.5.0a0+git0b7e8df
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.5 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.9.19 (main, May  6 2024, 14:39:30)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.5-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M3 Pro

Versions of relevant libraries:
[pip3] numpy==2.0.0
[pip3] optree==0.11.0
[pip3] torch==2.5.0a0+git0b7e8df
[pip3] torchgraph==0.0.1
[conda] numpy                     2.0.0                    pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] torch                     2.5.0a0+git0b7e8df           dev_0    <develop>
[conda] torchgraph                0.0.1                     dev_0    <develop>
```

## How to fix?

I put a naive fix that add the potential args to lift into the used_names. This visits private variables, will fix that if this issue makes sense to you.

@zou3519 @oulgen

Co-authored-by: rzou <[email protected]>
Pull Request resolved: pytorch#129817
Approved by: https://github.com/zou3519

Author: 2 people

test/dynamo/test_autograd_function.py

@@ -4,12 +4,14 @@
 import math
 
 from dataclasses import dataclass
+from typing import List
 
 import torch
 
 import torch._dynamo.test_case
 import torch._dynamo.testing
 import torch._dynamo.utils
+from functorch.compile import aot_module_simplified
 from torch.testing._internal.triton_utils import HAS_CUDA, requires_cuda
 
 if HAS_CUDA:
@@ -223,6 +225,69 @@ def forward(self, x):
         return self.f(x)
 
 
+@torch.library.custom_op("_torch_testing::custom_op_forward", mutates_args=())
+def custom_op_forward(
+    foo: torch.Tensor,
+    bar: torch.Tensor,
+    shape: List[int],
+) -> torch.Tensor:
+    return torch.ones_like(foo)
+
+
+@custom_op_forward.register_fake
+def _(foo, bar, weight):
+    return torch.empty_like(foo)
+
+
+@torch.library.custom_op("_torch_testing::custom_op_backward", mutates_args=())
+def custom_op_backward(
+    grad_output: torch.Tensor,
+    foo: torch.Tensor,
+    bar: torch.Tensor,
+    shape: List[int],
+) -> torch.Tensor:
+    assert list(bar.shape) == shape
+    return torch.ones_like(bar)
+
+
+@custom_op_backward.register_fake
+def _(grad_output, foo, bar, shape):
+    return torch.empty_like(bar)
+
+
+class CustomOpFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weight, normalized_shape):
+        ctx.normalized_shape = normalized_shape
+        input_ = input.contiguous()
+        weight_ = weight.contiguous()
+        output = custom_op_forward(input_, weight_, ctx.normalized_shape)
+        ctx.save_for_backward(input_, weight_)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input_, weight_ = ctx.saved_tensors
+        # grad_weight = a_func(grad_output, input_, weight_, ctx.normalized_shape)
+        grad_weight = custom_op_backward(
+            grad_output.contiguous(),
+            input_,
+            weight_,
+            ctx.normalized_shape,
+        )
+        return None, grad_weight, None
+
+
+class CustomOpModule(torch.nn.Module):
+    def __init__(self, shape):
+        super().__init__()
+        self.shape = shape
+        self.weight = torch.nn.Parameter(torch.ones(self.shape))
+
+    def forward(self, x):
+        return CustomOpFunc.apply(x, self.weight, self.shape)
+
+
 class AutogradFunctionTests(torch._dynamo.test_case.TestCase):
     # Sound behaviors, tested for working capture
     def test_autograd_function_equivalence(self):
@@ -527,18 +592,29 @@ def forward(self, L_x_: "f32[]", L_z_: "f32[]", L_weird_b: "f32[]", L_weird_c: "
 
     class GraphModule(torch.nn.Module):
         def forward(self, ctx, x: "f32[]", z: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):
-            mul: "f32[]" = l_weird_b * l_weird_c
-            clone: "f32[]" = x.clone();  x = None
+            ctx_1 = ctx
+            x_1 = x
+            z_1 = z
+            l_weird_b_1 = l_weird_b
+            l_weird_c_1 = l_weird_c
+
+            mul: "f32[]" = l_weird_b_1 * l_weird_c_1
+            clone: "f32[]" = x_1.clone();  x_1 = None
             mul_1: "f32[]" = mul * clone;  mul = clone = None
-            return (mul_1, [l_weird_b, l_weird_c])
+            return (mul_1, [l_weird_b_1, l_weird_c_1])
 
     class GraphModule(torch.nn.Module):
         def forward(self, ctx, grad: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):
+            ctx_1 = ctx
+            grad_1 = grad
+            l_weird_b_1 = l_weird_b
+            l_weird_c_1 = l_weird_c
+
             _set_grad_enabled = torch._C._set_grad_enabled(False)
 
-            mul: "f32[]" = grad * l_weird_b;  l_weird_b = None
-            mul_1: "f32[]" = mul * l_weird_c;  mul = l_weird_c = None
-            mul_2: "f32[]" = grad * 2;  grad = None
+            mul: "f32[]" = grad_1 * l_weird_b_1;  l_weird_b_1 = None
+            mul_1: "f32[]" = mul * l_weird_c_1;  mul = l_weird_c_1 = None
+            mul_2: "f32[]" = grad_1 * 2;  grad_1 = None
 
             _set_grad_enabled_1 = torch._C._set_grad_enabled(True)
             return (mul_1, mul_2)
@@ -1103,6 +1179,22 @@ def fn():
         self.assertEqual(cnt.frame_count, 1)
         self.assertEqual(cnt.op_count, 2)
 
+    def test_custom_op(self):
+        shape = [7]
+        x = torch.rand(128, shape[0])
+        model = CustomOpModule(shape)
+        out = model(x)
+
+        def backend(gm, example_inputs):
+            return aot_module_simplified(
+                gm, example_inputs, fw_compiler=lambda gm, _: gm
+            )
+
+        opt_model = torch.compile(model, backend=backend)
+        opt_out = opt_model(x)
+        opt_out.mean().backward()
+        self.assertEqual(out, opt_out)
+
     @requires_cuda
     def test_triton_kernel_basic(self):
         class Add(torch.autograd.Function):

test/dynamo/test_export.py

@@ -1894,16 +1894,16 @@ def forward(self, x):
                 out_graph.cond_true_0.code.strip(),
                 """\
 def forward(self, l_x_):
-    l_x__1 = l_x_
-    add = l_x__1 + l_x__1;  l_x__1 = None
+    l_x__2 = l_x_
+    add = l_x__2 + l_x__2;  l_x__2 = None
     return (add,)""",
             )
             self.assertExpectedInline(
                 out_graph.cond_false_0.code.strip(),
                 """\
 def forward(self, l_x_):
-    l_x__1 = l_x_
-    getitem = l_x__1[slice(None, 2, None)];  l_x__1 = None
+    l_x__2 = l_x_
+    getitem = l_x__2[slice(None, 2, None)];  l_x__2 = None
     return (getitem,)""",
             )
             with self.assertRaisesRegex(
@@ -3947,13 +3947,13 @@ def forward(self, pred, x):
             out_graph.cond_true_0.code.strip(),
             """\
 def forward(self, a, b, l_x_, d_true_branch, c_false_branch):
-    a_1 = a
-    b_1 = b
-    l_x__1 = l_x_
-    add = l_x__1 + l_x__1;  l_x__1 = None
-    cos = a_1.cos();  a_1 = None
+    a_2 = a
+    b_2 = b
+    l_x__2 = l_x_
+    add = l_x__2 + l_x__2;  l_x__2 = None
+    cos = a_2.cos();  a_2 = None
     add_1 = add + cos;  add = cos = None
-    cos_1 = b_1.cos();  b_1 = None
+    cos_1 = b_2.cos();  b_2 = None
     add_2 = add_1 + cos_1;  add_1 = cos_1 = None
     cos_2 = d_true_branch.cos();  d_true_branch = None
     add_3 = add_2 + cos_2;  add_2 = cos_2 = None
@@ -3964,13 +3964,13 @@ def forward(self, a, b, l_x_, d_true_branch, c_false_branch):
             out_graph.cond_false_0.code.strip(),
             """\
 def forward(self, a, b, l_x_, d_true_branch, c_false_branch):
-    a_1 = a
-    b_1 = b
-    l_x__1 = l_x_
-    mul = l_x__1 * l_x__1;  l_x__1 = None
-    sin = a_1.sin();  a_1 = None
+    a_2 = a
+    b_2 = b
+    l_x__2 = l_x_
+    mul = l_x__2 * l_x__2;  l_x__2 = None
+    sin = a_2.sin();  a_2 = None
     add = mul + sin;  mul = sin = None
-    sin_1 = b_1.sin();  b_1 = None
+    sin_1 = b_2.sin();  b_2 = None
     add_1 = add + sin_1;  add = sin_1 = None
     sin_2 = c_false_branch.sin();  c_false_branch = None
     add_2 = add_1 + sin_2;  add_1 = sin_2 = None

test/dynamo/test_functions.py

@@ -1964,6 +1964,7 @@ def forward(self, L_lambda0_keywords_y_: "f32[2, 2]"):
                 """\
 class GraphModule(torch.nn.Module):
     def forward(self, s0: "Sym(s0)", L_lambda0_keywords_y_: "f32[s0, s0]"):
+        s0_1 = s0
         l_lambda0_keywords_y_ = L_lambda0_keywords_y_
 
         mul: "f32[s0, s0]" = l_lambda0_keywords_y_ * l_lambda0_keywords_y_
@@ -2012,6 +2013,7 @@ def forward(self, L_lambda0_keywords_y_: "f32[2, 2]"):
                 """\
 class GraphModule(torch.nn.Module):
     def forward(self, s0: "Sym(s0)", L_lambda0_keywords_y_: "f32[s0, s0]"):
+        s0_1 = s0
         l_lambda0_keywords_y_ = L_lambda0_keywords_y_
 
         mul: "f32[s0, s0]" = l_lambda0_keywords_y_ * l_lambda0_keywords_y_
@@ -2063,6 +2065,7 @@ def forward(self, L_lambda0_keywords_y_: "f32[2, 2]"):
                 """\
 class GraphModule(torch.nn.Module):
     def forward(self, s0: "Sym(s0)", L_lambda0_keywords_y_: "f32[s0, s0]"):
+        s0_1 = s0
         l_lambda0_keywords_y_ = L_lambda0_keywords_y_
 
         mul: "f32[s0, s0]" = l_lambda0_keywords_y_ * l_lambda0_keywords_y_
@@ -2111,6 +2114,7 @@ def forward(self, L_x_: "f32[2, 2]"):
                 """\
 class GraphModule(torch.nn.Module):
     def forward(self, s0: "Sym(s0)", L_x_: "f32[s0, s0]"):
+        s0_1 = s0
         l_x_ = L_x_
 
         mul: "f32[s0, s0]" = l_x_ * 4