Properly set modified for RemoveNopSliceOrViewOpPass (#15470)

Summary:

Ensure modified is being correctly set for RemoveNopSliceOrViewOpPass

Differential Revision: D85815510

Author: Andrew Grebenisan

backends/cadence/aot/remove_ops.py

@@ -191,6 +191,11 @@ class RemoveNopSliceOrViewOpPass(ExportPass):
     Remove slice ops that are more like views, and view ops that do not change the shape
     """
 
+    @staticmethod
+    def _is_nop(input_shape: tuple[int, ...], output_shape: tuple[int, ...]) -> bool:
+        # If both arg_shape and out_shape are the same, this slice is a nop
+        return input_shape == output_shape
+
     def call_operator(
         self,
         op,  # pyre-ignore
@@ -207,13 +212,27 @@ def call_operator(
         arg0 = cast(ProxyValue, args[0])
         out_shape = meta["val"].shape
 
-        # If both arg_shape and out_shape are the same, this slice is a nop
         return (
             arg0
-            if arg0.to_tensor().shape == out_shape
+            if self._is_nop(arg0.to_tensor().shape, out_shape)
             else super().call_operator(op, args, kwargs, meta)
         )
 
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        for target in [
+            exir_ops.edge.aten.slice_copy.Tensor,
+            exir_ops.edge.aten.view_copy.default,
+        ]:
+            for node in graph_module.graph.find_nodes(
+                op="call_function", target=target
+            ):
+                input_node = node.args[0]
+                assert isinstance(input_node, torch.fx.Node)
+                if self._is_nop(input_node.meta["val"].shape, node.meta["val"].shape):
+                    return super().call(graph_module)
+
+        return PassResult(graph_module, False)
+
 
 @register_cadence_pass(CadencePassAttribute(opt_level=1))
 class RemoveNopLinalgVectorNormOpPass(ExportPass):

backends/cadence/aot/tests/test_remove_ops_passes.py

@@ -279,6 +279,7 @@ def test_remove_nop_view(self, shape: Tuple[int], new_shape: List[int]) -> None:
         graph_after_passes = cast(
             PassResult, RemoveNopSliceOrViewOpPass()(original)
         ).graph_module
+        assert original is not graph_after_passes
         self.assertEqual(
             count_node(graph_after_passes, exir_ops.edge.aten.view_copy.default), 0
         )
@@ -300,10 +301,29 @@ def test_remove_nop_slice(self) -> None:
         graph_after_passes = cast(
             PassResult, RemoveNopSliceOrViewOpPass()(original)
         ).graph_module
+        assert original is not graph_after_passes
         self.assertEqual(
             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()
+        graph_after_passes = cast(
+            PassResult, RemoveNopSliceOrViewOpPass()(original)
+        ).graph_module
+        # If the graph is not modified, the same object should be returned
+        assert original is graph_after_passes
+        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))