We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
hf_Bart `assert dst.data_ptr() == src.data_ptr()` error
@Chillee is this heuristic for static inputs still valid with your latest changes?
```
    def is_not_gradout(x):
        return "tangents" not in x.name
```


I am seeing the error below, which I believe means one of the "tangents" is changing when we didn't expect it to.


```
cuda train hf_Bart                            WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING:root:using triton random, expect difference from eager
WARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices
Similarity score=0.766252875328064
Accuracy failed for key name logits
INCORRECT
Traceback (most recent call last):
  File "./benchmarks/torchbench.py", line 390, in <module>
    main(TorchBenchmarkRunner(), original_dir)
  File "/home/jansel/torchdynamo/benchmarks/common.py", line 1547, in main
    runner.run_one_model(
  File "/home/jansel/torchdynamo/benchmarks/common.py", line 900, in run_one_model
    model_iter_fn(model, example_inputs)
  File "./benchmarks/torchbench.py", line 374, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "./benchmarks/torchbench.py", line 374, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "./benchmarks/torchbench.py", line 380, in forward_and_backward_pass
    self.grad_scaler.scale(loss).backward()
  File "/home/jansel/pytorch/functorch/functorch/_src/monkey_patching.py", line 77, in _backward
    return _old_backward(*args, **kwargs)
  File "/home/jansel/pytorch/torch/_tensor.py", line 484, in backward
    torch.autograd.backward(
  File "/home/jansel/pytorch/torch/autograd/__init__.py", line 191, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/jansel/pytorch/torch/autograd/function.py", line 265, in apply
    return user_fn(self, *args)
  File "/home/jansel/torchdynamo/torchdynamo/eval_frame.py", line 94, in _fn
    return fn(*args, **kwargs)
  File "/home/jansel/pytorch/functorch/functorch/_src/aot_autograd.py", line 271, in backward
    out = normalize_as_list(compiled_bw(*ctx.saved_tensors, *contiguous_args))
  File "/home/jansel/torchdynamo/torchdynamo/eval_frame.py", line 94, in _fn
    return fn(*args, **kwargs)
  File "/home/jansel/torchdynamo/torchinductor/compile_fx.py", line 205, in run
    assert dst.data_ptr() == src.data_ptr()
AssertionError

```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchinductor/compile_fx.py`
Content:
```
1 import functools
2 import itertools
3 import logging
4 import operator
5 import os
6 import textwrap
7 from typing import List
8 
9 import torch.fx
10 from functorch.compile import min_cut_rematerialization_partition
11 
12 import torchdynamo.config
13 from torchdynamo.optimizations.backends import aot_autograd
14 from torchdynamo.optimizations.normalize import normalize_ir
15 from torchdynamo.optimizations.python_key import python_key_normalize
16 from torchdynamo.testing import same
17 from torchdynamo.utils import identity
18 from torchdynamo.utils import init_logging
19 
20 from . import config
21 from . import overrides
22 from .decomposition import decompositions
23 from .graph import GraphLowering
24 from .virtualized import V
25 
26 log = logging.getLogger(__name__)
27 
28 
29 class CheckEachNode(torch.fx.Interpreter):
30     def call_function(self, target, args, kwargs):
31         expected = target(*args, **kwargs)
32         if target in (operator.getitem,):
33             return expected
34 
35         g = torch.fx.Graph()
36         g_args = []
37         a_args = []
38         for n, arg in enumerate(args):
39             if isinstance(arg, torch.Tensor):
40                 g_args.append(g.placeholder(f"arg{n}"))
41                 a_args.append(arg)
42             else:
43                 g_args.append(arg)
44         assert all(not isinstance(x, torch.Tensor) for x in kwargs.values())
45         node = g.call_function(target, tuple(g_args), kwargs)
46         if isinstance(expected, torch.Tensor):
47             node = (node,)
48         g.output(node)
49 
50         gm = torch.fx.GraphModule({}, g)
51         graph = GraphLowering(gm)
52         with V.set_graph_handler(graph):
53             graph.run(*args, **kwargs)
54             actual = graph.compile_to_fn()(*a_args)
55 
56         if isinstance(expected, torch.Tensor):
57             actual = actual[0]
58 
59         print(target, same(expected, actual))
60         assert same(expected, actual)
61 
62         return expected
63 
64 
65 def dump_to_repro(gm, *args):
66     with open(os.path.join(torchdynamo.config.base_dir, "repro.py"), "w") as fd:
67         fd.write(
68             textwrap.dedent(
69                 """
70                 import torch
71                 import torchdynamo
72                 from torchdynamo.testing import rand_strided, same
73                 """
74             )
75         )
76         fd.write("class Repro:\n")
77         for i in itertools.count():
78             try:
79                 val = getattr(gm, f"_tensor_constant{i}")
80             except AttributeError:
81                 break
82             fd.write(f"    _tensor_constant{i} = {val.item()!r}\n")
83         fd.write(textwrap.indent(gm.code, "    "))
84         fd.write("\n")
85 
86         fd.write("args = (\n")
87         for arg in args:
88             fd.write(
89                 f"  rand_strided({tuple(arg.size())!r}, {tuple(arg.stride())!r},"
90                 f" {arg.dtype!r}, {arg.device.type!r}),"
91             )
92             fd.write("\n")
93         fd.write(")\n")
94         fd.write(
95             textwrap.dedent(
96                 """
97                 expected = Repro().forward(*args)
98                 with torchdynamo.optimize("inductor", nopython=True):
99                     actual = Repro().forward(*args)
100                 assert same(actual, expected)
101                 """
102             )
103         )
104         print("wrote repro.py")
105 
106 
107 def compile_fx_python_key(
108     model: torch.fx.GraphModule, example_inputs: List[torch.Tensor], cudagraphs=None
109 ):
110     """Alternate version for inference only"""
111     assert isinstance(model, torch.fx.GraphModule)
112     assert all(isinstance(x, torch.Tensor) for x in example_inputs)
113 
114     with overrides.patch_functions():
115         model = overrides.replace_fx(model)
116         gm, wrap = python_key_normalize(
117             model, example_inputs, decompositions=decompositions
118         )
119 
120     if config.dce:
121         gm.graph.eliminate_dead_code()
122     if config.debug:
123         gm.graph.print_tabular()
124 
125     if os.environ.get("TORCHINDUCTOR_CHECK_OPS") == "1":
126         wrap(CheckEachNode(gm).run)(*example_inputs)
127 
128     return compile_fx_inner(gm, example_inputs, wrap=wrap, cudagraphs=cudagraphs)
129 
130 
131 def compile_fx_inner(
132     gm: torch.fx.GraphModule,
133     example_inputs: List[torch.Tensor],
134     wrap=identity,
135     cudagraphs=None,
136     num_fixed=0,
137 ):
138     init_logging()
139 
140     if cudagraphs is None:
141         cudagraphs = config.triton.cudagraphs
142 
143     try:
144         graph = GraphLowering(gm, num_dynamic_inputs=len(example_inputs))
145         with V.set_graph_handler(graph):
146             wrap(graph.run)(*example_inputs)
147             compiled_fn = wrap(graph.compile_to_fn())
148 
149         # make sure it works, causes issues for mutation
150         # compiled_fn(*example_inputs)
151 
152         if (
153             cudagraphs
154             and set(graph.device_types) == {"cuda"}
155             and not graph.mutated_inputs
156         ):
157             return cudagraphify(
158                 compiled_fn, example_inputs, static_input_idxs=range(num_fixed)
159             )
160         elif cudagraphs and set(graph.device_types) == {"cuda"}:
161             log.warning("skipping cudagraphs due to input mutation")
162         elif cudagraphs and len(graph.device_types) > 1:
163             log.warning("skipping cudagraphs due to multple devices")
164         return compiled_fn
165     except Exception:
166         if os.environ.get("TORCHINDUCTOR_DUMP_REPRO") == "1":
167             wrap(functools.partial(dump_to_repro, gm))(*example_inputs)
168 
169         raise
170 
171 
172 def cudagraphify(model, inputs, static_input_idxs=()):
173     """
174     Assumes inputs[static_input_idxs[i]] are always the same memory address
175     """
176     assert isinstance(inputs, (list, tuple))
177     static_inputs = [
178         torch.zeros_like(x) if idx not in static_input_idxs else inputs[idx]
179         for idx, x in enumerate(inputs)
180     ]
181 
182     # warmup
183     torch.cuda.synchronize()
184     stream = torch.cuda.Stream()
185     stream.wait_stream(torch.cuda.current_stream())
186     with torch.cuda.stream(stream):
187         model(*inputs)
188     stream.synchronize()
189     torch.cuda.current_stream().wait_stream(stream)
190     torch.cuda.synchronize()
191 
192     # record
193     graph = torch.cuda.CUDAGraph()
194     with torch.cuda.graph(graph, stream=stream):
195         static_outputs = model(*static_inputs)
196     if not isinstance(static_outputs, (list, tuple)):
197         static_outputs = (static_outputs,)
198 
199     if config.size_asserts:
200 
201         def run(*new_inputs):
202             assert len(static_inputs) == len(new_inputs)
203             for idx, (dst, src) in enumerate(zip(static_inputs, new_inputs)):
204                 if idx in static_input_idxs:
205                     assert dst.data_ptr() == src.data_ptr()
206                 else:
207                     dst.copy_(src)
208             graph.replay()
209             return static_outputs
210 
211     else:
212         copy_indices = [
213             idx for idx in range(len(static_inputs)) if idx not in static_input_idxs
214         ]
215 
216         def run(*new_inputs):
217             for idx in copy_indices:
218                 static_inputs[idx].copy_(new_inputs[idx])
219             graph.replay()
220             return static_outputs
221 
222     return run
223 
224 
225 def count_tangents(fx_g: torch.fx.GraphModule):
226     """
227     Infers which inputs are static for a backwards graph
228     """
229 
230     def is_not_gradout(x):
231         return "tangents" not in x.name
232 
233     arg_count = 0
234     static_arg_idxs = []
235     for n in fx_g.graph.nodes:
236         if n.op == "placeholder":
237             if is_not_gradout(n):
238                 static_arg_idxs.append(arg_count)
239             arg_count += 1
240 
241     assert static_arg_idxs == list(range(len(static_arg_idxs)))
242     return len(static_arg_idxs)
243 
244 
245 def compile_fx_aot(model_: torch.fx.GraphModule, example_inputs_: List[torch.Tensor]):
246     """Main entrypoint to a compile given FX graph"""
247     with overrides.patch_functions():
248         model_ = normalize_ir(model_, example_inputs_)
249         model_ = overrides.replace_fx(model_)
250     num_example_inputs = len(example_inputs_)
251 
252     def fw_compiler(model: torch.fx.GraphModule, example_inputs):
253         if config.debug:
254             print("FORWARD GRAPH:")
255             model.graph.print_tabular()
256         fixed = len(example_inputs) - num_example_inputs
257         return compile_fx_inner(model, example_inputs, num_fixed=fixed)
258 
259     def bw_compiler(model: torch.fx.GraphModule, example_inputs):
260         if config.debug:
261             print("BACKWARD GRAPH:")
262             model.graph.print_tabular()
263         fixed = count_tangents(model)
264         return compile_fx_inner(model, example_inputs, num_fixed=fixed)
265 
266     with overrides.patch_functions():
267         return aot_autograd(
268             model_,
269             example_inputs_,
270             fw_compiler=fw_compiler,
271             bw_compiler=bw_compiler,
272             decompositions=decompositions,
273             partition_fn=functools.partial(
274                 min_cut_rematerialization_partition, compiler="inductor"
275             ),
276         )
277 
278 
279 def compile_fx(model_: torch.fx.GraphModule, example_inputs_: List[torch.Tensor]):
280     """Main entrypoint to a compile given FX graph"""
281     logging.getLogger("torchinductor").setLevel(
282         logging.DEBUG if config.debug else logging.WARNING
283     )
284     if config.aot_autograd:
285         return compile_fx_aot(model_, example_inputs_)
286     else:
287         return compile_fx_python_key(model_, example_inputs_)
288
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torchinductor/compile_fx.py b/torchinductor/compile_fx.py
--- a/torchinductor/compile_fx.py
+++ b/torchinductor/compile_fx.py
@@ -1,3 +1,4 @@
+import dataclasses
 import functools
 import itertools
 import logging
@@ -26,6 +27,14 @@
 log = logging.getLogger(__name__)
 
 
+@dataclasses.dataclass
+class BoxedBool:
+    value: bool
+
+    def __bool__(self):
+        return self.value
+
+
 class CheckEachNode(torch.fx.Interpreter):
     def call_function(self, target, args, kwargs):
         expected = target(*args, **kwargs)
@@ -157,10 +166,15 @@
             return cudagraphify(
                 compiled_fn, example_inputs, static_input_idxs=range(num_fixed)
             )
-        elif cudagraphs and set(graph.device_types) == {"cuda"}:
-            log.warning("skipping cudagraphs due to input mutation")
-        elif cudagraphs and len(graph.device_types) > 1:
-            log.warning("skipping cudagraphs due to multple devices")
+        elif cudagraphs:
+            if isinstance(cudagraphs, BoxedBool):
+                # Disable cudagraphs in the backwards pass too:
+                cudagraphs.value = False
+
+            if set(graph.device_types) == {"cuda"}:
+                log.warning("skipping cudagraphs due to input mutation")
+            elif set(graph.device_types) != {"cpu"}:
+                log.warning("skipping cudagraphs due to multiple devices")
         return compiled_fn
     except Exception:
         if os.environ.get("TORCHINDUCTOR_DUMP_REPRO") == "1":
@@ -248,20 +262,25 @@
         model_ = normalize_ir(model_, example_inputs_)
         model_ = overrides.replace_fx(model_)
     num_example_inputs = len(example_inputs_)
+    cudagraphs = BoxedBool(config.triton.cudagraphs)
 
     def fw_compiler(model: torch.fx.GraphModule, example_inputs):
         if config.debug:
             print("FORWARD GRAPH:")
             model.graph.print_tabular()
         fixed = len(example_inputs) - num_example_inputs
-        return compile_fx_inner(model, example_inputs, num_fixed=fixed)
+        return compile_fx_inner(
+            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs
+        )
 
     def bw_compiler(model: torch.fx.GraphModule, example_inputs):
         if config.debug:
             print("BACKWARD GRAPH:")
             model.graph.print_tabular()
         fixed = count_tangents(model)
-        return compile_fx_inner(model, example_inputs, num_fixed=fixed)
+        return compile_fx_inner(
+            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs
+        )
 
     with overrides.patch_functions():
         return aot_autograd(

{"golden_diff": "diff --git a/torchinductor/compile_fx.py b/torchinductor/compile_fx.py\n--- a/torchinductor/compile_fx.py\n+++ b/torchinductor/compile_fx.py\n@@ -1,3 +1,4 @@\n+import dataclasses\n import functools\n import itertools\n import logging\n@@ -26,6 +27,14 @@\n log = logging.getLogger(__name__)\n \n \n+@dataclasses.dataclass\n+class BoxedBool:\n+    value: bool\n+\n+    def __bool__(self):\n+        return self.value\n+\n+\n class CheckEachNode(torch.fx.Interpreter):\n     def call_function(self, target, args, kwargs):\n         expected = target(*args, **kwargs)\n@@ -157,10 +166,15 @@\n             return cudagraphify(\n                 compiled_fn, example_inputs, static_input_idxs=range(num_fixed)\n             )\n-        elif cudagraphs and set(graph.device_types) == {\"cuda\"}:\n-            log.warning(\"skipping cudagraphs due to input mutation\")\n-        elif cudagraphs and len(graph.device_types) > 1:\n-            log.warning(\"skipping cudagraphs due to multple devices\")\n+        elif cudagraphs:\n+            if isinstance(cudagraphs, BoxedBool):\n+                # Disable cudagraphs in the backwards pass too:\n+                cudagraphs.value = False\n+\n+            if set(graph.device_types) == {\"cuda\"}:\n+                log.warning(\"skipping cudagraphs due to input mutation\")\n+            elif set(graph.device_types) != {\"cpu\"}:\n+                log.warning(\"skipping cudagraphs due to multiple devices\")\n         return compiled_fn\n     except Exception:\n         if os.environ.get(\"TORCHINDUCTOR_DUMP_REPRO\") == \"1\":\n@@ -248,20 +262,25 @@\n         model_ = normalize_ir(model_, example_inputs_)\n         model_ = overrides.replace_fx(model_)\n     num_example_inputs = len(example_inputs_)\n+    cudagraphs = BoxedBool(config.triton.cudagraphs)\n \n     def fw_compiler(model: torch.fx.GraphModule, example_inputs):\n         if config.debug:\n             print(\"FORWARD GRAPH:\")\n             model.graph.print_tabular()\n         fixed = len(example_inputs) - num_example_inputs\n-        return compile_fx_inner(model, example_inputs, num_fixed=fixed)\n+        return compile_fx_inner(\n+            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs\n+        )\n \n     def bw_compiler(model: torch.fx.GraphModule, example_inputs):\n         if config.debug:\n             print(\"BACKWARD GRAPH:\")\n             model.graph.print_tabular()\n         fixed = count_tangents(model)\n-        return compile_fx_inner(model, example_inputs, num_fixed=fixed)\n+        return compile_fx_inner(\n+            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs\n+        )\n \n     with overrides.patch_functions():\n         return aot_autograd(\n", "issue": "hf_Bart `assert dst.data_ptr() == src.data_ptr()` error\n@Chillee is this heuristic for static inputs still valid with your latest changes?\r\n```\r\n    def is_not_gradout(x):\r\n        return \"tangents\" not in x.name\r\n```\r\n\r\n\r\nI am seeing the error below, which I believe means one of the \"tangents\" is changing when we didn't expect it to.\r\n\r\n\r\n```\r\ncuda train hf_Bart                            WARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING:root:using triton random, expect difference from eager\r\nWARNING torchinductor.compile_fx: skipping cudagraphs due to multple devices\r\nSimilarity score=0.766252875328064\r\nAccuracy failed for key name logits\r\nINCORRECT\r\nTraceback (most recent call last):\r\n  File \"./benchmarks/torchbench.py\", line 390, in <module>\r\n    main(TorchBenchmarkRunner(), original_dir)\r\n  File \"/home/jansel/torchdynamo/benchmarks/common.py\", line 1547, in main\r\n    runner.run_one_model(\r\n  File \"/home/jansel/torchdynamo/benchmarks/common.py\", line 900, in run_one_model\r\n    model_iter_fn(model, example_inputs)\r\n  File \"./benchmarks/torchbench.py\", line 374, in forward_and_backward_pass\r\n    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):\r\n  File \"./benchmarks/torchbench.py\", line 374, in forward_and_backward_pass\r\n    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):\r\n  File \"./benchmarks/torchbench.py\", line 380, in forward_and_backward_pass\r\n    self.grad_scaler.scale(loss).backward()\r\n  File \"/home/jansel/pytorch/functorch/functorch/_src/monkey_patching.py\", line 77, in _backward\r\n    return _old_backward(*args, **kwargs)\r\n  File \"/home/jansel/pytorch/torch/_tensor.py\", line 484, in backward\r\n    torch.autograd.backward(\r\n  File \"/home/jansel/pytorch/torch/autograd/__init__.py\", line 191, in backward\r\n    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\r\n  File \"/home/jansel/pytorch/torch/autograd/function.py\", line 265, in apply\r\n    return user_fn(self, *args)\r\n  File \"/home/jansel/torchdynamo/torchdynamo/eval_frame.py\", line 94, in _fn\r\n    return fn(*args, **kwargs)\r\n  File \"/home/jansel/pytorch/functorch/functorch/_src/aot_autograd.py\", line 271, in backward\r\n    out = normalize_as_list(compiled_bw(*ctx.saved_tensors, *contiguous_args))\r\n  File \"/home/jansel/torchdynamo/torchdynamo/eval_frame.py\", line 94, in _fn\r\n    return fn(*args, **kwargs)\r\n  File \"/home/jansel/torchdynamo/torchinductor/compile_fx.py\", line 205, in run\r\n    assert dst.data_ptr() == src.data_ptr()\r\nAssertionError\r\n\r\n```\n", "before_files": [{"content": "import functools\nimport itertools\nimport logging\nimport operator\nimport os\nimport textwrap\nfrom typing import List\n\nimport torch.fx\nfrom functorch.compile import min_cut_rematerialization_partition\n\nimport torchdynamo.config\nfrom torchdynamo.optimizations.backends import aot_autograd\nfrom torchdynamo.optimizations.normalize import normalize_ir\nfrom torchdynamo.optimizations.python_key import python_key_normalize\nfrom torchdynamo.testing import same\nfrom torchdynamo.utils import identity\nfrom torchdynamo.utils import init_logging\n\nfrom . import config\nfrom . import overrides\nfrom .decomposition import decompositions\nfrom .graph import GraphLowering\nfrom .virtualized import V\n\nlog = logging.getLogger(__name__)\n\n\nclass CheckEachNode(torch.fx.Interpreter):\n    def call_function(self, target, args, kwargs):\n        expected = target(*args, **kwargs)\n        if target in (operator.getitem,):\n            return expected\n\n        g = torch.fx.Graph()\n        g_args = []\n        a_args = []\n        for n, arg in enumerate(args):\n            if isinstance(arg, torch.Tensor):\n                g_args.append(g.placeholder(f\"arg{n}\"))\n                a_args.append(arg)\n            else:\n                g_args.append(arg)\n        assert all(not isinstance(x, torch.Tensor) for x in kwargs.values())\n        node = g.call_function(target, tuple(g_args), kwargs)\n        if isinstance(expected, torch.Tensor):\n            node = (node,)\n        g.output(node)\n\n        gm = torch.fx.GraphModule({}, g)\n        graph = GraphLowering(gm)\n        with V.set_graph_handler(graph):\n            graph.run(*args, **kwargs)\n            actual = graph.compile_to_fn()(*a_args)\n\n        if isinstance(expected, torch.Tensor):\n            actual = actual[0]\n\n        print(target, same(expected, actual))\n        assert same(expected, actual)\n\n        return expected\n\n\ndef dump_to_repro(gm, *args):\n    with open(os.path.join(torchdynamo.config.base_dir, \"repro.py\"), \"w\") as fd:\n        fd.write(\n            textwrap.dedent(\n                \"\"\"\n                import torch\n                import torchdynamo\n                from torchdynamo.testing import rand_strided, same\n                \"\"\"\n            )\n        )\n        fd.write(\"class Repro:\\n\")\n        for i in itertools.count():\n            try:\n                val = getattr(gm, f\"_tensor_constant{i}\")\n            except AttributeError:\n                break\n            fd.write(f\"    _tensor_constant{i} = {val.item()!r}\\n\")\n        fd.write(textwrap.indent(gm.code, \"    \"))\n        fd.write(\"\\n\")\n\n        fd.write(\"args = (\\n\")\n        for arg in args:\n            fd.write(\n                f\"  rand_strided({tuple(arg.size())!r}, {tuple(arg.stride())!r},\"\n                f\" {arg.dtype!r}, {arg.device.type!r}),\"\n            )\n            fd.write(\"\\n\")\n        fd.write(\")\\n\")\n        fd.write(\n            textwrap.dedent(\n                \"\"\"\n                expected = Repro().forward(*args)\n                with torchdynamo.optimize(\"inductor\", nopython=True):\n                    actual = Repro().forward(*args)\n                assert same(actual, expected)\n                \"\"\"\n            )\n        )\n        print(\"wrote repro.py\")\n\n\ndef compile_fx_python_key(\n    model: torch.fx.GraphModule, example_inputs: List[torch.Tensor], cudagraphs=None\n):\n    \"\"\"Alternate version for inference only\"\"\"\n    assert isinstance(model, torch.fx.GraphModule)\n    assert all(isinstance(x, torch.Tensor) for x in example_inputs)\n\n    with overrides.patch_functions():\n        model = overrides.replace_fx(model)\n        gm, wrap = python_key_normalize(\n            model, example_inputs, decompositions=decompositions\n        )\n\n    if config.dce:\n        gm.graph.eliminate_dead_code()\n    if config.debug:\n        gm.graph.print_tabular()\n\n    if os.environ.get(\"TORCHINDUCTOR_CHECK_OPS\") == \"1\":\n        wrap(CheckEachNode(gm).run)(*example_inputs)\n\n    return compile_fx_inner(gm, example_inputs, wrap=wrap, cudagraphs=cudagraphs)\n\n\ndef compile_fx_inner(\n    gm: torch.fx.GraphModule,\n    example_inputs: List[torch.Tensor],\n    wrap=identity,\n    cudagraphs=None,\n    num_fixed=0,\n):\n    init_logging()\n\n    if cudagraphs is None:\n        cudagraphs = config.triton.cudagraphs\n\n    try:\n        graph = GraphLowering(gm, num_dynamic_inputs=len(example_inputs))\n        with V.set_graph_handler(graph):\n            wrap(graph.run)(*example_inputs)\n            compiled_fn = wrap(graph.compile_to_fn())\n\n        # make sure it works, causes issues for mutation\n        # compiled_fn(*example_inputs)\n\n        if (\n            cudagraphs\n            and set(graph.device_types) == {\"cuda\"}\n            and not graph.mutated_inputs\n        ):\n            return cudagraphify(\n                compiled_fn, example_inputs, static_input_idxs=range(num_fixed)\n            )\n        elif cudagraphs and set(graph.device_types) == {\"cuda\"}:\n            log.warning(\"skipping cudagraphs due to input mutation\")\n        elif cudagraphs and len(graph.device_types) > 1:\n            log.warning(\"skipping cudagraphs due to multple devices\")\n        return compiled_fn\n    except Exception:\n        if os.environ.get(\"TORCHINDUCTOR_DUMP_REPRO\") == \"1\":\n            wrap(functools.partial(dump_to_repro, gm))(*example_inputs)\n\n        raise\n\n\ndef cudagraphify(model, inputs, static_input_idxs=()):\n    \"\"\"\n    Assumes inputs[static_input_idxs[i]] are always the same memory address\n    \"\"\"\n    assert isinstance(inputs, (list, tuple))\n    static_inputs = [\n        torch.zeros_like(x) if idx not in static_input_idxs else inputs[idx]\n        for idx, x in enumerate(inputs)\n    ]\n\n    # warmup\n    torch.cuda.synchronize()\n    stream = torch.cuda.Stream()\n    stream.wait_stream(torch.cuda.current_stream())\n    with torch.cuda.stream(stream):\n        model(*inputs)\n    stream.synchronize()\n    torch.cuda.current_stream().wait_stream(stream)\n    torch.cuda.synchronize()\n\n    # record\n    graph = torch.cuda.CUDAGraph()\n    with torch.cuda.graph(graph, stream=stream):\n        static_outputs = model(*static_inputs)\n    if not isinstance(static_outputs, (list, tuple)):\n        static_outputs = (static_outputs,)\n\n    if config.size_asserts:\n\n        def run(*new_inputs):\n            assert len(static_inputs) == len(new_inputs)\n            for idx, (dst, src) in enumerate(zip(static_inputs, new_inputs)):\n                if idx in static_input_idxs:\n                    assert dst.data_ptr() == src.data_ptr()\n                else:\n                    dst.copy_(src)\n            graph.replay()\n            return static_outputs\n\n    else:\n        copy_indices = [\n            idx for idx in range(len(static_inputs)) if idx not in static_input_idxs\n        ]\n\n        def run(*new_inputs):\n            for idx in copy_indices:\n                static_inputs[idx].copy_(new_inputs[idx])\n            graph.replay()\n            return static_outputs\n\n    return run\n\n\ndef count_tangents(fx_g: torch.fx.GraphModule):\n    \"\"\"\n    Infers which inputs are static for a backwards graph\n    \"\"\"\n\n    def is_not_gradout(x):\n        return \"tangents\" not in x.name\n\n    arg_count = 0\n    static_arg_idxs = []\n    for n in fx_g.graph.nodes:\n        if n.op == \"placeholder\":\n            if is_not_gradout(n):\n                static_arg_idxs.append(arg_count)\n            arg_count += 1\n\n    assert static_arg_idxs == list(range(len(static_arg_idxs)))\n    return len(static_arg_idxs)\n\n\ndef compile_fx_aot(model_: torch.fx.GraphModule, example_inputs_: List[torch.Tensor]):\n    \"\"\"Main entrypoint to a compile given FX graph\"\"\"\n    with overrides.patch_functions():\n        model_ = normalize_ir(model_, example_inputs_)\n        model_ = overrides.replace_fx(model_)\n    num_example_inputs = len(example_inputs_)\n\n    def fw_compiler(model: torch.fx.GraphModule, example_inputs):\n        if config.debug:\n            print(\"FORWARD GRAPH:\")\n            model.graph.print_tabular()\n        fixed = len(example_inputs) - num_example_inputs\n        return compile_fx_inner(model, example_inputs, num_fixed=fixed)\n\n    def bw_compiler(model: torch.fx.GraphModule, example_inputs):\n        if config.debug:\n            print(\"BACKWARD GRAPH:\")\n            model.graph.print_tabular()\n        fixed = count_tangents(model)\n        return compile_fx_inner(model, example_inputs, num_fixed=fixed)\n\n    with overrides.patch_functions():\n        return aot_autograd(\n            model_,\n            example_inputs_,\n            fw_compiler=fw_compiler,\n            bw_compiler=bw_compiler,\n            decompositions=decompositions,\n            partition_fn=functools.partial(\n                min_cut_rematerialization_partition, compiler=\"inductor\"\n            ),\n        )\n\n\ndef compile_fx(model_: torch.fx.GraphModule, example_inputs_: List[torch.Tensor]):\n    \"\"\"Main entrypoint to a compile given FX graph\"\"\"\n    logging.getLogger(\"torchinductor\").setLevel(\n        logging.DEBUG if config.debug else logging.WARNING\n    )\n    if config.aot_autograd:\n        return compile_fx_aot(model_, example_inputs_)\n    else:\n        return compile_fx_python_key(model_, example_inputs_)\n", "path": "torchinductor/compile_fx.py"}], "after_files": [{"content": "import dataclasses\nimport functools\nimport itertools\nimport logging\nimport operator\nimport os\nimport textwrap\nfrom typing import List\n\nimport torch.fx\nfrom functorch.compile import min_cut_rematerialization_partition\n\nimport torchdynamo.config\nfrom torchdynamo.optimizations.backends import aot_autograd\nfrom torchdynamo.optimizations.normalize import normalize_ir\nfrom torchdynamo.optimizations.python_key import python_key_normalize\nfrom torchdynamo.testing import same\nfrom torchdynamo.utils import identity\nfrom torchdynamo.utils import init_logging\n\nfrom . import config\nfrom . import overrides\nfrom .decomposition import decompositions\nfrom .graph import GraphLowering\nfrom .virtualized import V\n\nlog = logging.getLogger(__name__)\n\n\n@dataclasses.dataclass\nclass BoxedBool:\n    value: bool\n\n    def __bool__(self):\n        return self.value\n\n\nclass CheckEachNode(torch.fx.Interpreter):\n    def call_function(self, target, args, kwargs):\n        expected = target(*args, **kwargs)\n        if target in (operator.getitem,):\n            return expected\n\n        g = torch.fx.Graph()\n        g_args = []\n        a_args = []\n        for n, arg in enumerate(args):\n            if isinstance(arg, torch.Tensor):\n                g_args.append(g.placeholder(f\"arg{n}\"))\n                a_args.append(arg)\n            else:\n                g_args.append(arg)\n        assert all(not isinstance(x, torch.Tensor) for x in kwargs.values())\n        node = g.call_function(target, tuple(g_args), kwargs)\n        if isinstance(expected, torch.Tensor):\n            node = (node,)\n        g.output(node)\n\n        gm = torch.fx.GraphModule({}, g)\n        graph = GraphLowering(gm)\n        with V.set_graph_handler(graph):\n            graph.run(*args, **kwargs)\n            actual = graph.compile_to_fn()(*a_args)\n\n        if isinstance(expected, torch.Tensor):\n            actual = actual[0]\n\n        print(target, same(expected, actual))\n        assert same(expected, actual)\n\n        return expected\n\n\ndef dump_to_repro(gm, *args):\n    with open(os.path.join(torchdynamo.config.base_dir, \"repro.py\"), \"w\") as fd:\n        fd.write(\n            textwrap.dedent(\n                \"\"\"\n                import torch\n                import torchdynamo\n                from torchdynamo.testing import rand_strided, same\n                \"\"\"\n            )\n        )\n        fd.write(\"class Repro:\\n\")\n        for i in itertools.count():\n            try:\n                val = getattr(gm, f\"_tensor_constant{i}\")\n            except AttributeError:\n                break\n            fd.write(f\"    _tensor_constant{i} = {val.item()!r}\\n\")\n        fd.write(textwrap.indent(gm.code, \"    \"))\n        fd.write(\"\\n\")\n\n        fd.write(\"args = (\\n\")\n        for arg in args:\n            fd.write(\n                f\"  rand_strided({tuple(arg.size())!r}, {tuple(arg.stride())!r},\"\n                f\" {arg.dtype!r}, {arg.device.type!r}),\"\n            )\n            fd.write(\"\\n\")\n        fd.write(\")\\n\")\n        fd.write(\n            textwrap.dedent(\n                \"\"\"\n                expected = Repro().forward(*args)\n                with torchdynamo.optimize(\"inductor\", nopython=True):\n                    actual = Repro().forward(*args)\n                assert same(actual, expected)\n                \"\"\"\n            )\n        )\n        print(\"wrote repro.py\")\n\n\ndef compile_fx_python_key(\n    model: torch.fx.GraphModule, example_inputs: List[torch.Tensor], cudagraphs=None\n):\n    \"\"\"Alternate version for inference only\"\"\"\n    assert isinstance(model, torch.fx.GraphModule)\n    assert all(isinstance(x, torch.Tensor) for x in example_inputs)\n\n    with overrides.patch_functions():\n        model = overrides.replace_fx(model)\n        gm, wrap = python_key_normalize(\n            model, example_inputs, decompositions=decompositions\n        )\n\n    if config.dce:\n        gm.graph.eliminate_dead_code()\n    if config.debug:\n        gm.graph.print_tabular()\n\n    if os.environ.get(\"TORCHINDUCTOR_CHECK_OPS\") == \"1\":\n        wrap(CheckEachNode(gm).run)(*example_inputs)\n\n    return compile_fx_inner(gm, example_inputs, wrap=wrap, cudagraphs=cudagraphs)\n\n\ndef compile_fx_inner(\n    gm: torch.fx.GraphModule,\n    example_inputs: List[torch.Tensor],\n    wrap=identity,\n    cudagraphs=None,\n    num_fixed=0,\n):\n    init_logging()\n\n    if cudagraphs is None:\n        cudagraphs = config.triton.cudagraphs\n\n    try:\n        graph = GraphLowering(gm, num_dynamic_inputs=len(example_inputs))\n        with V.set_graph_handler(graph):\n            wrap(graph.run)(*example_inputs)\n            compiled_fn = wrap(graph.compile_to_fn())\n\n        # make sure it works, causes issues for mutation\n        # compiled_fn(*example_inputs)\n\n        if (\n            cudagraphs\n            and set(graph.device_types) == {\"cuda\"}\n            and not graph.mutated_inputs\n        ):\n            return cudagraphify(\n                compiled_fn, example_inputs, static_input_idxs=range(num_fixed)\n            )\n        elif cudagraphs:\n            if isinstance(cudagraphs, BoxedBool):\n                # Disable cudagraphs in the backwards pass too:\n                cudagraphs.value = False\n\n            if set(graph.device_types) == {\"cuda\"}:\n                log.warning(\"skipping cudagraphs due to input mutation\")\n            elif set(graph.device_types) != {\"cpu\"}:\n                log.warning(\"skipping cudagraphs due to multiple devices\")\n        return compiled_fn\n    except Exception:\n        if os.environ.get(\"TORCHINDUCTOR_DUMP_REPRO\") == \"1\":\n            wrap(functools.partial(dump_to_repro, gm))(*example_inputs)\n\n        raise\n\n\ndef cudagraphify(model, inputs, static_input_idxs=()):\n    \"\"\"\n    Assumes inputs[static_input_idxs[i]] are always the same memory address\n    \"\"\"\n    assert isinstance(inputs, (list, tuple))\n    static_inputs = [\n        torch.zeros_like(x) if idx not in static_input_idxs else inputs[idx]\n        for idx, x in enumerate(inputs)\n    ]\n\n    # warmup\n    torch.cuda.synchronize()\n    stream = torch.cuda.Stream()\n    stream.wait_stream(torch.cuda.current_stream())\n    with torch.cuda.stream(stream):\n        model(*inputs)\n    stream.synchronize()\n    torch.cuda.current_stream().wait_stream(stream)\n    torch.cuda.synchronize()\n\n    # record\n    graph = torch.cuda.CUDAGraph()\n    with torch.cuda.graph(graph, stream=stream):\n        static_outputs = model(*static_inputs)\n    if not isinstance(static_outputs, (list, tuple)):\n        static_outputs = (static_outputs,)\n\n    if config.size_asserts:\n\n        def run(*new_inputs):\n            assert len(static_inputs) == len(new_inputs)\n            for idx, (dst, src) in enumerate(zip(static_inputs, new_inputs)):\n                if idx in static_input_idxs:\n                    assert dst.data_ptr() == src.data_ptr()\n                else:\n                    dst.copy_(src)\n            graph.replay()\n            return static_outputs\n\n    else:\n        copy_indices = [\n            idx for idx in range(len(static_inputs)) if idx not in static_input_idxs\n        ]\n\n        def run(*new_inputs):\n            for idx in copy_indices:\n                static_inputs[idx].copy_(new_inputs[idx])\n            graph.replay()\n            return static_outputs\n\n    return run\n\n\ndef count_tangents(fx_g: torch.fx.GraphModule):\n    \"\"\"\n    Infers which inputs are static for a backwards graph\n    \"\"\"\n\n    def is_not_gradout(x):\n        return \"tangents\" not in x.name\n\n    arg_count = 0\n    static_arg_idxs = []\n    for n in fx_g.graph.nodes:\n        if n.op == \"placeholder\":\n            if is_not_gradout(n):\n                static_arg_idxs.append(arg_count)\n            arg_count += 1\n\n    assert static_arg_idxs == list(range(len(static_arg_idxs)))\n    return len(static_arg_idxs)\n\n\ndef compile_fx_aot(model_: torch.fx.GraphModule, example_inputs_: List[torch.Tensor]):\n    \"\"\"Main entrypoint to a compile given FX graph\"\"\"\n    with overrides.patch_functions():\n        model_ = normalize_ir(model_, example_inputs_)\n        model_ = overrides.replace_fx(model_)\n    num_example_inputs = len(example_inputs_)\n    cudagraphs = BoxedBool(config.triton.cudagraphs)\n\n    def fw_compiler(model: torch.fx.GraphModule, example_inputs):\n        if config.debug:\n            print(\"FORWARD GRAPH:\")\n            model.graph.print_tabular()\n        fixed = len(example_inputs) - num_example_inputs\n        return compile_fx_inner(\n            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs\n        )\n\n    def bw_compiler(model: torch.fx.GraphModule, example_inputs):\n        if config.debug:\n            print(\"BACKWARD GRAPH:\")\n            model.graph.print_tabular()\n        fixed = count_tangents(model)\n        return compile_fx_inner(\n            model, example_inputs, num_fixed=fixed, cudagraphs=cudagraphs\n        )\n\n    with overrides.patch_functions():\n        return aot_autograd(\n            model_,\n            example_inputs_,\n            fw_compiler=fw_compiler,\n            bw_compiler=bw_compiler,\n            decompositions=decompositions,\n            partition_fn=functools.partial(\n                min_cut_rematerialization_partition, compiler=\"inductor\"\n            ),\n        )\n\n\ndef compile_fx(model_: torch.fx.GraphModule, example_inputs_: List[torch.Tensor]):\n    \"\"\"Main entrypoint to a compile given FX graph\"\"\"\n    logging.getLogger(\"torchinductor\").setLevel(\n        logging.DEBUG if config.debug else logging.WARNING\n    )\n    if config.aot_autograd:\n        return compile_fx_aot(model_, example_inputs_)\n    else:\n        return compile_fx_python_key(model_, example_inputs_)\n", "path": "torchinductor/compile_fx.py"}]}