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

Chillee · Chillee · commit b426d0f02986 · 2023-08-25T13:12:34.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
@@ -500,6 +500,51 @@ def f(a):
         inp = T(10)
         self.assertExpectedInline(count_numel(f, inp), """20""")
 
+    def test_noop_dtype_conversion(self):
+        def f(a):
+            b = torch.ops.prims.convert_element_type(a, torch.float32)
+            c = unfusible(b)
+            return c
+
+        inp = T(10)
+        self.assertExpectedInline(count_numel(f, inp), """20""")
+
+    def test_noop_device_conversion(self):
+        def f(a):
+            b = torch.ops.prims.device_put(a, "cuda")
+            c = unfusible(b)
+            return c
+
+        inp = T(10)
+        self.assertExpectedInline(count_numel(f, inp), """20""")
+
+    def test_noop_int_ops(self):
+        def f1(a):
+            b = torch.ceil(a)
+            c = unfusible(b)
+            return c
+
+        def f2(a):
+            d = torch.floor(a)
+            e = unfusible(d)
+            return e
+
+        def f3(a):
+            f = torch.round(a)
+            g = unfusible(f)
+            return g
+
+        def f4(a):
+            f = torch.pow(a, 1)
+            g = unfusible(f)
+            return g
+
+        inp = TI(10)
+        self.assertExpectedInline(count_numel(f1, inp), """20""")
+        self.assertExpectedInline(count_numel(f2, inp), """20""")
+        self.assertExpectedInline(count_numel(f3, inp), """20""")
+        self.assertExpectedInline(count_numel(f4, 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
@@ -10,9 +10,11 @@
 import torch
 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_node_storage
+from ..fx_utils import FakeTensorUpdater, get_fake_args_kwargs, get_node_storage
 
 from ..lowering import lowerings as L
 from ..pattern_matcher import (
@@ -31,6 +33,7 @@
     register_graph_pattern,
     stable_topological_sort,
 )
+from ..utils import decode_device
 from ..virtualized import V
 from .group_batch_fusion import group_batch_fusion_post_grad_passes
 
@@ -543,6 +546,26 @@ def constant_pad_nd(x, padding, fill_value=0):
         return True
 
 
+@register_noop_decomp(torch.ops.prims.convert_element_type)
+def convert_element_type_noop(x, dtype: torch.dtype):
+    return x.dtype == dtype
+
+
+@register_noop_decomp(torch.ops.prims.device_put)
+def device_put_noop(x, device):
+    return x.device == decode_device(device)
+
+
+@register_noop_decomp([aten.ceil, aten.floor, aten.round, aten.trunc])
+def int_noop(x):
+    return is_integer_dtype(x.dtype)
+
+
+@register_noop_decomp([aten.pow])
+def pow_noop(a, b):
+    return isinstance(b, int) and b == 1
+
+
 @register_noop_decomp([aten.clone, aten.alias])
 def true_noop(*args, **kwargs):
     return True
@@ -578,7 +601,10 @@ def remove_noop_ops(graph: torch.fx.Graph):
                 and get_node_storage(src) in input_storages
             ):
                 continue
-            if same_layout(node, src) and cond(*node.args, **node.kwargs):
+            is_valid, args, kwargs = get_fake_args_kwargs(node)
+            if not is_valid:
+                continue
+            if same_layout(node, src) and cond(*args, **kwargs):
                 node.replace_all_uses_with(src)
                 graph.erase_node(node)
 
diff --git a/torch/_inductor/fx_utils.py b/torch/_inductor/fx_utils.py
@@ -1,8 +1,9 @@
 from collections import defaultdict
+from typing import Optional, Tuple
 
 import torch
 import torch.fx
-from torch.utils._pytree import tree_map
+from torch.utils._pytree import tree_flatten, tree_map
 from .virtualized import V
 
 
@@ -63,7 +64,8 @@ def __init__(self, graph: torch.fx.Graph):
             self.processed_hashes.add(self.hash_node(node))
 
     def hash_node(self, node: torch.fx.Node):
-        return (node, node.target, node.args, node.kwargs)
+        # todo(chilli): Not a great hash function
+        return (node, node.target, id(node.args), id(node.kwargs))
 
     def incremental_update(self):
         processed = set()
@@ -100,11 +102,6 @@ def is_fake_tensor_same(new, old):
             if self.hash_node(node) in self.processed_hashes:
                 continue
 
-            def get_fake_tensor(x):
-                if isinstance(x, torch.fx.Node):
-                    return x.meta["val"]
-                return x
-
             def is_aten_node(node):
                 return node.op == "call_function" and isinstance(
                     node.target, torch._ops.OpOverload
@@ -121,9 +118,9 @@ def is_aten_node(node):
                 if is_aten_node(updating_node):
                     continue
 
-                args, kwargs = tree_map(
-                    get_fake_tensor, (updating_node.args, updating_node.kwargs)
-                )
+                is_valid, args, kwargs = get_fake_args_kwargs(updating_node)
+                if not is_valid:
+                    continue
                 with V.fake_mode:
                     new_fake_tensor = updating_node.target(*args, **kwargs)
                 if "val" in updating_node.meta and is_fake_tensor_same(
@@ -139,15 +136,33 @@ def is_aten_node(node):
                 self.processed_hashes.add(self.hash_node(updating_node))
 
 
-def get_storage(t):
+def get_storage(t: torch.Tensor) -> int:
     return t.untyped_storage()._cdata
 
 
-def get_node_storage(node):
+def get_node_storage(node: torch.Tensor) -> Optional[int]:
     if "val" not in node.meta:
         return None
     if not isinstance(node.meta["val"], torch.Tensor):
         return None
     if not torch._C._has_storage(node.meta["val"]):
         return None
     return get_storage(node.meta["val"])
+
+
+def get_fake(x):
+    if isinstance(x, torch.fx.Node):
+        if "val" not in x.meta:
+            return x
+        return x.meta["val"]
+    return x
+
+
+def get_fake_args_kwargs(x: torch.fx.Node) -> Tuple[bool, tuple, dict]:
+    """
+    First value returns a boolean if any of the input nodes don't have a faketensor.
+    """
+    args, kwargs = tree_map(get_fake, (x.args, x.kwargs))
+    if any(isinstance(a, torch.fx.Node) for a in tree_flatten((args, kwargs))[0]):
+        return False, args, kwargs
+    return True, args, kwargs
