Properly set modified for RemoveNopSliceOrViewOpPass (pytorch#15470)

Summary:

Ensure modified is being correctly set for RemoveNopSliceOrViewOpPass. This diff also introduces a common base class that we can use for all remove ops passes.

Reviewed By: abeakkas

Differential Revision: D85815510

Author: Andrew Grebenisan

backends/cadence/aot/remove_ops.py

@@ -6,8 +6,8 @@
 
 # pyre-strict
 
-
 import logging
+from abc import abstractmethod
 from dataclasses import dataclass, field
 from typing import cast, List, Optional, Sequence, Set, Type
 
@@ -35,6 +35,45 @@
 from torch.fx.node import Argument, Node
 
 
+class RemovePassCommon(ExportPass):
+    @property
+    @abstractmethod
+    def targets(self) -> list[EdgeOpOverload]:
+        """
+        The list of targets to potentially remove.
+        """
+        raise NotImplementedError("`targets` must be defined")
+
+    @abstractmethod
+    def maybe_remove(self, node: Node) -> bool:
+        """
+        If the node should be removed, removes from the graph. Returns
+        True if the graph was modified, else False.
+        """
+        raise NotImplementedError("`maybe_remove` must be defined")
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        """
+        For each node in targets, if the node should be removed, removes from
+        the graph and returns the modified graph and modified set to True.
+        If no node should be removed, returns a pass result with the original
+        graph module and False for modified.
+        """
+        changed = False
+        for target in self.targets:
+            for module in graph_module.modules():
+                assert isinstance(module, torch.fx.GraphModule)
+                for node in module.graph.find_nodes(op="call_function", target=target):
+                    changed |= self.maybe_remove(node)
+
+        if changed:
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            return super().call(graph_module)
+
+        return PassResult(graph_module, False)
+
+
 @register_cadence_pass(CadencePassAttribute(opt_level=0))
 class RemoveCloneOpsTransformImported(ExportPass):
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
@@ -186,33 +225,27 @@ def call_operator(
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=1))
-class RemoveNopSliceOrViewOpPass(ExportPass):
+class RemoveNopSliceOrViewOpPass(RemovePassCommon):
     """
     Remove slice ops that are more like views, and view ops that do not change the shape
     """
 
-    def call_operator(
-        self,
-        op,  # pyre-ignore
-        args: tuple[Argument, ...],
-        kwargs: dict[str, Argument],
-        meta: NodeMetadata,
-    ) -> ProxyValue:
-        if op not in {
+    @property
+    def targets(self) -> list[EdgeOpOverload]:
+        return [
             exir_ops.edge.aten.slice_copy.Tensor,
             exir_ops.edge.aten.view_copy.default,
-        }:
-            return super().call_operator(op, args, kwargs, meta)
+        ]
 
-        arg0 = cast(ProxyValue, args[0])
-        out_shape = meta["val"].shape
+    def maybe_remove(self, node: Node) -> bool:
+        changed = False
+        input_node = node.args[0]
+        assert isinstance(input_node, Node)
+        if input_node.meta["val"].shape == node.meta["val"].shape:
+            node.replace_all_uses_with(input_node)
+            changed = True
 
-        # If both arg_shape and out_shape are the same, this slice is a nop
-        return (
-            arg0
-            if arg0.to_tensor().shape == out_shape
-            else super().call_operator(op, args, kwargs, meta)
-        )
+        return changed
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=1))

backends/cadence/aot/tests/test_remove_ops_passes.py

@@ -304,6 +304,22 @@ def test_remove_nop_slice(self) -> None:
             count_node(graph_after_passes, exir_ops.edge.aten.slice_copy.Tensor), 0
         )
 
+    def test_remove_nop_slice_or_view_not_modified(self) -> None:
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(3, 5, dtype=torch.float32))
+        abs_x = builder.call_operator(
+            op=exir_ops.edge.aten.abs.default,
+            args=(x,),
+        )
+        builder.output([abs_x])
+        original = builder.get_graph_module()
+        pass_result = cast(PassResult, RemoveNopSliceOrViewOpPass()(original))
+        self.assertFalse(pass_result.modified)
+        graph_after_passes = pass_result.graph_module
+        self.assertEqual(
+            count_node(graph_after_passes, exir_ops.edge.aten.abs.default), 1
+        )
+
     def test_remove_nop_select_before_view(self) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn(1, 5, 6, dtype=torch.float32))