Update on "Add reinplacing pass for scatters + incremental fake tensor updating"

Chillee · Chillee · commit aecd43348db4 · 2023-08-25T14:11:12.000-07:00
mutation for params)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 aakhundov anijain2305

[ghstack-poisoned]
diff --git a/test/inductor/test_perf.py b/test/inductor/test_perf.py
@@ -545,6 +545,21 @@ def f4(a):
         self.assertExpectedInline(count_numel(f3, inp), """20""")
         self.assertExpectedInline(count_numel(f4, inp), """20""")
 
+    def test_noop_cat(self):
+        def f1(a):
+            b = torch.cat([a])
+            return unfusible(b)
+
+        inp = T(10)
+        self.assertExpectedInline(count_numel(f1, inp), """20""")
+
+        def f2(a):
+            b = torch.cat([a])
+            c = torch.cat([b])
+            return c
+
+        self.assertExpectedInline(count_numel(f2, inp), """20""")
+
 
 class InplacingTests(TestCase):
     def test_inplace_scatter(self):
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
@@ -11,7 +11,6 @@
 import torch._inductor as inductor
 from torch._decomp import register_decomposition
 from torch._prims_common import is_integer_dtype
-from torch.utils._pytree import tree_map
 
 from .. import config, ir, pattern_matcher
 from ..fx_utils import FakeTensorUpdater, get_fake_args_kwargs, get_node_storage
@@ -511,10 +510,10 @@ def same_layout(node1: torch.fx.Node, node2: torch.fx.Node):
 noop_registry: Dict[Any, Any] = {}
 
 
-def register_noop_decomp(targets, arg_num=0):
+def register_noop_decomp(targets, nop_arg=0):
     def register_fun(cond):
         register_decomposition(targets, registry=noop_registry, unsafe=True)(
-            (cond, arg_num)
+            (cond, nop_arg)
         )
 
     return register_fun
@@ -566,6 +565,11 @@ def pow_noop(a, b):
     return isinstance(b, int) and b == 1
 
 
+@register_noop_decomp([aten.cat], lambda args: args[0][0])
+def cat_noop(inputs, dim=0):
+    return len(inputs) == 1
+
+
 @register_noop_decomp([aten.clone, aten.alias])
 def true_noop(*args, **kwargs):
     return True
@@ -593,7 +597,10 @@ def remove_noop_ops(graph: torch.fx.Graph):
     for node in graph.nodes:
         if node.target in noop_registry:
             cond, src_index = noop_registry[node.target]
-            src = node.args[src_index]
+            if isinstance(src_index, int):
+                src = node.args[src_index]
+            else:
+                src = src_index(node.args)
             # See fx_passes/README.md for a discussion of why this is
             # necessary.
             if (
