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

Chillee · Chillee · commit b118ec6bffbe · 2023-08-23T15:47: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
@@ -8,7 +8,7 @@
 import torch._inductor.config as config
 from torch._dynamo.backends.registry import register_backend
 from torch._inductor import metrics
-from torch._inductor.compile_fx import compile_fx, count_bytes_inner
+import torch._inductor.compile_fx
 from torch.testing._internal.common_utils import IS_WINDOWS, TestCase as TorchTestCase
 from torch.testing._internal.inductor_utils import HAS_CUDA
 
diff --git a/test/inductor/test_snode_runtime.py b/test/inductor/test_snode_runtime.py
@@ -2,7 +2,7 @@
 
 
 import torch
-from torch._dynamo.backends.registry import register_backend
+import torch._inductor.compile_fx
 from torch._inductor import metrics
 from torch._inductor.compile_fx import compile_fx, count_bytes_inner
 from torch.testing._internal.common_utils import TestCase as TorchTestCase
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -276,6 +276,24 @@ def fmax(self, other):
     return torch.where(torch.isnan(other) | (other < self), self, other)
 
 
+@register_decomposition(aten.slice)
+def slice(self, dim=0, start=None, end=None, step=1):
+    if start is None or end is None:
+        return NotImplemented
+    if start == 0 and end >= 2**63 - 1 and step == 1:
+        return aten.alias(self)
+    return NotImplemented
+
+
+@register_decomposition(aten.slice_scatter)
+def slice_scatter(self, src, dim=0, start=None, end=None, step=1):
+    if start is None or end is None:
+        return NotImplemented
+    if start == 0 and end >= 2**63 - 1 and step == 1:
+        return aten.clone(src)
+    return NotImplemented
+
+
 @register_decomposition([aten.narrow_copy])
 def narrow_copy(self, dim, start, length):
     return torch.narrow(self, dim, start, length).clone()
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
@@ -9,9 +9,9 @@
 
 import torch
 import torch._inductor as inductor
-from torch.utils._pytree import tree_map
 
 from .. import config, ir, pattern_matcher
+from ..fx_utils import FakeTensorUpdater, get_node_storage
 
 from ..lowering import lowerings as L
 from ..pattern_matcher import (
@@ -48,96 +48,6 @@
 inference_patterns = PatternMatcherPass()
 
 
-def get_storage(t):
-    return t.untyped_storage()._cdata
-
-
-def get_node_storage(node):
-    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"])
-
-
-class FakeTensorUpdater:
-    def __init__(self, graph):
-        self.processed_hashes = set()
-        self.graph = graph
-
-        for node in self.graph.nodes:
-            self.processed_hashes.add(self.hash_node(node))
-
-    def hash_node(self, node):
-        return (node, node.target, id(node.args), id(node.kwargs))
-
-    def incremental_update(self):
-        """
-        Not fully featured yet. Mainly added to ensure that reinplacing pass is
-        sound.
-
-        Todo things:
-        1. We probably want to persist this across more passes.
-        """
-        processed = set()
-        existing_tensors = defaultdict(int)
-        for node in self.graph.nodes:
-            existing_tensors[get_node_storage(node)] += 1
-
-        def is_fake_tensor_same(new, old):
-            if type(new) != type(old):
-                return False
-            if isinstance(new, (list, tuple)):
-                return all(is_fake_tensor_same(new_i, old_i) for new_i, old_i in zip(new, old))
-            if new.shape != old.shape or new.stride() != old.stride():
-                return False
-            if get_storage(new) == get_storage(old):
-                return True
-            if (
-                existing_tensors[get_storage(old)] == 1
-                and get_storage(new) not in existing_tensors
-            ):
-                return True
-            return False
-
-        for node in self.graph.nodes:
-            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
-
-            if node.op == "call_function" and isinstance(
-                node.target, torch._ops.OpOverload
-            ):
-                processing = [node]
-                while len(processing) > 0:
-                    updating_node = processing.pop()
-                    if updating_node in processed:
-                        continue
-                    if updating_node.op != "call_function" or not isinstance(
-                        updating_node.target, torch._ops.OpOverload
-                    ):
-                        continue
-
-                    args, kwargs = tree_map(
-                        get_fake_tensor, (updating_node.args, updating_node.kwargs)
-                    )
-                    new_fake_tensor = updating_node.target(*args, **kwargs)
-                    if "val" in updating_node.meta and is_fake_tensor_same(new_fake_tensor, updating_node.meta["val"]):
-                        continue
-                    updating_node.meta["val"] = new_fake_tensor
-                    processed.add(updating_node)
-                    for user in updating_node.users:
-                        processing.append(user)
-
-                self.processed_hashes.add(self.hash_node(updating_node))
-
-
 def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
     """
     Passes that run on after grad.  This is called once on the forwards
@@ -594,23 +504,11 @@ def same_layout(node1: torch.fx.Node, node2: torch.fx.Node):
 
 
 def remove_noop_ops(graph: torch.fx.Graph):
-    def true(*args, **kwargs):
-        return True
-
-    def slice_scatter_cond(self, src, dim=0, start=None, end=None, step=1):
-        if start is None or end is None:
-            return False
-        if start == 0 and end >= 2**63 - 1 and step == 1:
-            return True
-        return False
-
-    def slice_cond(self, dim=0, start=None, end=None, step=1):
-        if start is None or end is None:
-            return False
-        if start == 0 and end >= 2**63 - 1 and step == 1:
-            return True
-        return False
+    """
+    Removes aten.clone and aten.alias ops from the graph when it's safe.
 
+    Other no-ops should be done as decompositions that selectively turn into aten.clone or aten.alias
+    """
     input_storages = set()
     output_storages = set()
     for node in graph.nodes:
@@ -623,27 +521,22 @@ def slice_cond(self, dim=0, start=None, end=None, step=1):
         if isinstance(out, torch.fx.Node):
             output_storages.add(get_node_storage(out))
 
-    noop_conditions = {
-        aten.slice.Tensor: (0, slice_cond),
-        aten.clone.default: (0, true),
-        aten.slice_scatter.default: (1, slice_scatter_cond),
-    }
-
     for node in graph.nodes:
-        if node.target in noop_conditions:
-            cloned_arg_index, cond = noop_conditions[node.target]
-            src = node.args[cloned_arg_index]
+        if node.target in (aten.alias.default, aten.clone.default):
+            src = node.args[0]
             if (
                 get_node_storage(node) in output_storages
                 and get_node_storage(src) in input_storages
             ):
                 continue
-            if same_layout(node, src) and cond(*node.args, **node.kwargs):
+            if same_layout(node, src):
                 node.replace_all_uses_with(src)
                 graph.erase_node(node)
 
 
 InplaceableOp = namedtuple("InplaceableOp", ["inplace_op", "mutated_arg"])
+
+
 def reinplace_scatters(graph):
     """
     Reinplaces scatter operations in easy cases where the node being mutated
@@ -660,10 +553,9 @@ def reinplace_scatters(graph):
         storage_to_nodes[get_node_storage(node)].append(node)
         if node.target == aten.copy_.default:
             copy_nodes[(node.args[0], node.args[1])] = node
-            assert node.args[0].op == 'placeholder'
+            assert node.args[0].op == "placeholder"
             mutated_inputs.add(node.args[0])
 
-
     inplaceable_ops = {
         aten.index_put.default: InplaceableOp(aten.index_put_.default, 0),
     }
diff --git a/torch/_inductor/fx_utils.py b/torch/_inductor/fx_utils.py
@@ -1,4 +1,9 @@
+from collections import defaultdict
+
 import torch
+import torch.fx
+from torch.utils._pytree import tree_map
+from .virtualized import V
 
 
 # Check the pattern: (nn.module, F.function/torch.Tensor.method) matched.
@@ -28,3 +33,110 @@ def matches_module_function_pattern(pattern, node, modules):
     if len(node.args[0].users) > 1:
         return False
     return True
+
+
+class FakeTensorUpdater:
+    """
+    The main idea here is that it's difficult to maintain accurate fake
+    tensors (our primary form of metadata) for each node in our graph as we
+    transform it.
+
+    The most reliable way to obtain this information is by rerunning
+    faketensor propagation. However, in general, faketensor propagation is
+    fairly expensive. So, instead we'd like to only rerun faketensor
+    propagation on nodes that have changed.
+
+    In order to detect which nodes have changed, we first hash its node,
+    target, and argument lists (which are immutable in FX).
+
+    Then, whenever we call incremental_update, we check which FX nodes have a
+    new hash, and recompute the faketensor metadata for that node. Then, we
+    continue to recursively compute the faketensors for all users until the
+    fake tensors stop changing.
+    """
+
+    def __init__(self, graph: torch.fx.Graph):
+        self.processed_hashes = set()
+        self.graph = graph
+
+        for node in self.graph.nodes:
+            self.processed_hashes.add(self.hash_node(node))
+
+    def hash_node(self, node: torch.fx.Node):
+        return (node, node.target, id(node.args), id(node.kwargs))
+
+    def incremental_update(self):
+        processed = set()
+        existing_tensors = defaultdict(int)
+        for node in self.graph.nodes:
+            existing_tensors[get_node_storage(node)] += 1
+
+        def is_fake_tensor_same(new, old):
+            if type(new) != type(old):
+                return False
+            if isinstance(new, (list, tuple)):
+                return all(
+                    is_fake_tensor_same(new_i, old_i) for new_i, old_i in zip(new, old)
+                )
+            if new.shape != old.shape or new.stride() != old.stride():
+                return False
+            if get_storage(new) == get_storage(old):
+                return True
+            if (
+                existing_tensors[get_storage(old)] == 1
+                and get_storage(new) not in existing_tensors
+            ):
+                return True
+            return False
+
+        for node in self.graph.nodes:
+            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
+
+            if node.op == "call_function" and isinstance(
+                node.target, torch._ops.OpOverload
+            ):
+                processing = [node]
+                while len(processing) > 0:
+                    updating_node = processing.pop()
+                    if updating_node in processed:
+                        continue
+                    if updating_node.op != "call_function" or not isinstance(
+                        updating_node.target, torch._ops.OpOverload
+                    ):
+                        continue
+
+                    args, kwargs = tree_map(
+                        get_fake_tensor, (updating_node.args, updating_node.kwargs)
+                    )
+                    with V.fake_mode:
+                        new_fake_tensor = updating_node.target(*args, **kwargs)
+                    if "val" in updating_node.meta and is_fake_tensor_same(
+                        new_fake_tensor, updating_node.meta["val"]
+                    ):
+                        continue
+                    updating_node.meta["val"] = new_fake_tensor
+                    processed.add(updating_node)
+                    for user in updating_node.users:
+                        processing.append(user)
+
+                self.processed_hashes.add(self.hash_node(updating_node))
+
+
+def get_storage(t):
+    return t.untyped_storage()._cdata
+
+
+def get_node_storage(node):
+    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"])
