Remove pass-local ExportPass replays (pytorch#22091)

Differential Revision: D115472301

Pull Request resolved: pytorch#22091

Author: mcremon-meta

backends/cadence/aot/passes.py

@@ -68,6 +68,7 @@
 from executorch.exir.passes.scalar_to_tensor_pass import ScalarToTensorPass
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 from torch.export.exported_program import ExportedProgram
+from torch.fx.passes.infra.pass_base import PassBase
 
 
 class InitializePipeline(ExportPass):
@@ -104,7 +105,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
 # The passes that must run for the graph to be legal on the target, regardless
 # of compile mode. Everything else in the pipeline is an optimization.
-REQUIRED_PASSES: frozenset[Type[ExportPass]] = frozenset(
+REQUIRED_PASSES: frozenset[Type[PassBase]] = frozenset(
     {
         InitializePipeline,
         FinalizePipeline,
@@ -139,7 +140,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 )
 
 
-def _get_pipeline() -> list[Type[ExportPass]]:
+def _get_pipeline() -> list[Type[PassBase]]:
     """The full ordered pass pipeline.
 
     Order is load-bearing and levels are interleaved, so this list is the single
@@ -168,7 +169,7 @@ def _get_pipeline() -> list[Type[ExportPass]]:
 def get_passes(
     mode: CompileMode | str,
     edge_passes_config: Optional[EdgePassesConfig] = None,
-) -> list[Type[ExportPass]]:
+) -> list[Type[PassBase]]:
     # Coerce at the choke point: modes arrive from CLIs and JSON as plain
     # strings, and a str silently matches none of the checks below.
     mode = CompileMode(mode)

backends/cadence/aot/remove_ops.py

@@ -40,6 +40,7 @@
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind, OutputKind
 from torch.fx.node import Node
+from torch.fx.passes.infra.pass_base import PassBase
 from torch.utils import _pytree as pytree
 
 
@@ -727,7 +728,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         return PassResult(graph_module, False)
 
 
-class RemoveBranchedQuantDequant(ExportPass):
+class RemoveBranchedQuantDequant(PassBase):
     """
     This pass looks for adjacent quant and dequant nodes with identical
     parameters, where the quant node has other users in addition to the
@@ -755,10 +756,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
         if modified:
             graph_module.graph.eliminate_dead_code()
-            result = super().call(graph_module)
-            return result
+            graph_module.recompile()
 
-        return PassResult(graph_module, False)
+        return PassResult(graph_module, modified)
 
     def remove_branched(
         self,
@@ -848,7 +848,7 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
 
 
 class CommonRemovePasses:
-    passes: List[Type[ExportPass]] = [
+    passes: List[Type[PassBase]] = [
         # Canonicalise squeeze/unsqueeze to view_copy first: the nop-view and
         # permute passes below both reason about view_copy only.
         ReplaceSqueezeAndUnsqueezeWithViewPassImported,
@@ -964,7 +964,7 @@ def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
 
 
 class CadenceRemoveNops:
-    passes: List[Type[ExportPass]] = CommonRemovePasses.passes + [
+    passes: List[Type[PassBase]] = CommonRemovePasses.passes + [
         SimplifySliceOpPass,
         RemoveNopRequantizeOpPass,
         RemoveZeroSizedConstantPadNd,

backends/cadence/aot/replace_ops.py

@@ -37,7 +37,8 @@
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
-from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.pass_base import PassResult
+from torch.fx.passes.infra.pass_base import PassBase
 
 # A map to represent ops that:
 # (a) are functionally equivalent; and
@@ -1981,47 +1982,37 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
         return True
 
 
-class ReplaceEmptyTensorsWithFullPass(ExportPass):
+class ReplaceEmptyTensorsWithFullPass(PassBase):
     """Replaces nodes that produce empty tensors with full nodes."""
 
-    def call_operator(self, op, args, kwargs, meta):
-        val = meta.data.get("val", None)
-        if isinstance(val, torch.Tensor) and val.numel() == 0:
-            return super().call_operator(
-                exir_ops.edge.aten.full.default,
-                args=(val.shape, 0),
-                kwargs={"dtype": val.dtype},
-                meta=meta,
-            )
-        return super().call_operator(op, args, kwargs, meta)
-
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         changed = False
         for module in filter(
             lambda m: isinstance(m, torch.fx.GraphModule), graph_module.modules()
         ):
-            module = cast(torch.fx.GraphModule, module)
-            for node in module.graph.nodes:
+            graph_module_to_update = cast(torch.fx.GraphModule, module)
+            module_changed = False
+            for node in graph_module_to_update.graph.nodes:
                 if node.op != "call_function":
                     continue
                 val = node.meta.get("val", None)
                 if isinstance(val, torch.Tensor) and val.numel() == 0:
-                    with module.graph.inserting_before(node):
-                        new_node = module.graph.call_function(
+                    with graph_module_to_update.graph.inserting_before(node):
+                        new_node = graph_module_to_update.graph.call_function(
                             exir_ops.edge.aten.full.default,
                             args=(val.shape, 0),
                             kwargs={"dtype": val.dtype},
                         )
-                        new_node.meta = node.meta
+                        new_node.meta = node.meta.copy()
                     node.replace_all_uses_with(new_node)
-                    changed = True
+                    module_changed = True
 
-        if changed:
-            graph_module.graph.eliminate_dead_code()
-            graph_module.recompile()
-            return super().call(graph_module)
+            if module_changed:
+                graph_module_to_update.graph.eliminate_dead_code()
+                graph_module_to_update.recompile()
+                changed = True
 
-        return PassResult(graph_module, False)
+        return PassResult(graph_module, changed)
 
 
 class ReplaceWhereWithFullArgsWithWhereScalar(RemoveOrReplacePassInterface):

backends/cadence/aot/simplify_ops.py

@@ -17,8 +17,9 @@
 from executorch.backends.cadence.aot.utils import rebind
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
-from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.pass_base import PassResult
 from torch.fx import Node
+from torch.fx.passes.infra.pass_base import PassBase
 
 
 class SimplifySliceOpPass(RemoveOrReplacePassInterface):
@@ -169,7 +170,7 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
         return True
 
 
-class BindOptionalArgsPass(ExportPass):
+class BindOptionalArgsPass(PassBase):
     """Bind all optional args and kwargs."""
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
@@ -181,7 +182,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         for module in filter(
             lambda m: isinstance(m, torch.fx.GraphModule), graph_module.modules()
         ):
-            for node in cast(torch.fx.GraphModule, module).graph.nodes:
+            graph_module_to_update = cast(torch.fx.GraphModule, module)
+            module_modified = False
+            for node in graph_module_to_update.graph.nodes:
                 if node.op != "call_function":
                     continue
                 if not isinstance(node.target, EdgeOpOverload):
@@ -198,11 +201,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 if new_args != node.args or new_kwargs != node.kwargs:
                     node.args = new_args
                     node.kwargs = new_kwargs
-                    modified = True
+                    module_modified = True
 
-        if modified:
-            graph_module.recompile()
-            return super().call(graph_module)
+            if module_modified:
+                graph_module_to_update.recompile()
+                modified = True
 
         return PassResult(graph_module, modified)
 

backends/cadence/aot/tests/test_remove_ops_passes.py

@@ -913,9 +913,28 @@ def test_remove_dequant_on_branch(self) -> None:
         )
         builder.output([x1_output, y1_output])
         original = builder.get_graph_module()
+        original_nodes = list(original.graph.nodes)
+        original_abs_nodes = original.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.abs.default
+        )
+        x1_output_node = original_abs_nodes[1]
+        x1_output_node.meta["remove_branched_test"] = "preserved"
+        x1_output_meta = x1_output_node.meta
         pass_result = cast(PassResult, RemoveBranchedQuantDequant()(original))
         self.assertTrue(pass_result.modified)
         graph_after_passes = pass_result.graph_module
+        self.assertIs(graph_after_passes, original)
+        self.assertIs(x1_output_node.meta, x1_output_meta)
+        self.assertEqual(x1_output_node.meta["remove_branched_test"], "preserved")
+        self.assertIs(
+            next(
+                node
+                for node in graph_after_passes.graph.nodes
+                if node.name == x1_output_node.name
+            ),
+            x1_output_node,
+        )
+        self.assertLess(len(list(graph_after_passes.graph.nodes)), len(original_nodes))
         self.assertEqual(
             count_node(
                 graph_after_passes,

backends/cadence/aot/tests/test_replace_ops_passes.py

@@ -3179,10 +3179,32 @@ def test_empty_slice(self) -> None:
 
         # Deepcopy before the pass
         gm_before = copy.deepcopy(gm)
+        original_graph_module = gm
+        original_nodes = list(gm.graph.nodes)
+        empty_slice = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+        )[0]
+        empty_slice.meta["replace_empty_test"] = "preserved"
+        empty_slice_meta = empty_slice.meta
+        unchanged_slice = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+        )[1]
+        unchanged_slice_meta = unchanged_slice.meta
 
         result = ReplaceEmptyTensorsWithFullPass().call(gm)
         self.assertTrue(result.modified)
         updated_gm = result.graph_module
+        self.assertIs(updated_gm, original_graph_module)
+        self.assertIs(unchanged_slice.meta, unchanged_slice_meta)
+        self.assertIn(unchanged_slice, updated_gm.graph.nodes)
+        self.assertNotIn(empty_slice, updated_gm.graph.nodes)
+        full = updated_gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.full.default
+        )[0]
+        self.assertIsNot(full.meta, empty_slice_meta)
+        self.assertEqual(full.meta, empty_slice_meta)
+        self.assertEqual(full.meta["replace_empty_test"], "preserved")
+        self.assertEqual(len(list(updated_gm.graph.nodes)), len(original_nodes))
 
         # Validate numerical accuracy
         inputs = [x_input]

backends/cadence/aot/tests/test_simplify_ops_passes.py

@@ -147,14 +147,49 @@ def test_simplify_slice_op_args(self) -> None:
             args=(x, 1),
             kwargs={"end": 3},
         )
+        original_graph_module = gm
         original_slice_copy = list(gm.graph.nodes)[1]
+        original_metadata = original_slice_copy.meta
+        original_slice_copy.meta["bind_optional_args_test"] = "preserved"
         self.assertEqual(original_slice_copy.args[1:], (1,))
         self.assertEqual(original_slice_copy.kwargs, {"end": 3})
+
         result = transform_and_check_numerics(
             gm, (x,), BindOptionalArgsPass(), "BindOptionalArgsPass"
         )
+
         self.assertTrue(result.modified)
-        gm = result.graph_module
-        modified_slice_copy = list(gm.graph.nodes)[1]
+        self.assertIs(result.graph_module, original_graph_module)
+        modified_slice_copy = list(result.graph_module.graph.nodes)[1]
+        self.assertIs(modified_slice_copy, original_slice_copy)
+        self.assertIs(modified_slice_copy.meta, original_metadata)
+        self.assertEqual(
+            modified_slice_copy.meta["bind_optional_args_test"], "preserved"
+        )
         self.assertEqual(modified_slice_copy.args[1:], (1, None, 3, 1))
         self.assertEqual(modified_slice_copy.kwargs, {})
+
+    def test_bind_optional_args_no_change(self) -> None:
+        x = torch.rand(4, 5)
+        gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 1, None, 3, 1),
+        )
+        original_nodes = list(gm.graph.nodes)
+        slice_copy = original_nodes[1]
+        original_metadata = slice_copy.meta
+        slice_copy.meta["bind_optional_args_test"] = "preserved"
+
+        result = BindOptionalArgsPass().call(gm)
+
+        self.assertFalse(result.modified)
+        self.assertIs(result.graph_module, gm)
+        self.assertEqual(
+            len(list(result.graph_module.graph.nodes)), len(original_nodes)
+        )
+        self.assertIs(list(result.graph_module.graph.nodes)[1], slice_copy)
+        self.assertIs(slice_copy.meta, original_metadata)
+        self.assertEqual(slice_copy.meta["bind_optional_args_test"], "preserved")
+        self.assertEqual(slice_copy.args[1:], (1, None, 3, 1))
+        self.assertEqual(slice_copy.kwargs, {})

backends/transforms/permute_pass_utils.py

@@ -157,8 +157,11 @@ def _apply_hierarchical_inplace(self, graph_module: torch.fx.GraphModule) -> boo
             modified |= self._apply_flat_inplace(module)
         return modified
 
+    def apply_inplace(self, graph_module: torch.fx.GraphModule) -> bool:
+        return self._apply_hierarchical_inplace(graph_module)
+
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
-        modified = self._apply_hierarchical_inplace(graph_module)
+        modified = self.apply_inplace(graph_module)
         if modified:
             graph_module.graph.eliminate_dead_code()
             graph_module.recompile()

backends/transforms/test/test_permute_optimization_passes.py

@@ -125,6 +125,42 @@ def test_permute_transpose_fusion(self) -> None:
             "FuseCascadedTransposeOrPermuteOps",
         )
 
+    def test_apply_inplace_preserves_identity_and_output_metadata(self) -> None:
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(2, 3, 4))
+        permute1 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1])
+        )
+        permute2 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [2, 1, 0])
+        )
+        builder.output([permute2])
+        graph_module = builder.get_graph_module()
+        placeholder = next(
+            node for node in graph_module.graph.nodes if node.op == "placeholder"
+        )
+        placeholder.meta["apply_inplace_test"] = "preserved"
+        placeholder_meta = placeholder.meta
+        output_permute = graph_module.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.permute_copy.default
+        )[1]
+        output_permute.meta["apply_inplace_test"] = "output"
+        output_meta = output_permute.meta
+
+        modified = FuseCascadedTransposeOrPermuteOps().apply_inplace(graph_module)
+        graph_module.graph.eliminate_dead_code()
+        graph_module.recompile()
+
+        self.assertTrue(modified)
+        self.assertIn(placeholder, graph_module.graph.nodes)
+        self.assertIs(placeholder.meta, placeholder_meta)
+        self.assertEqual(placeholder.meta["apply_inplace_test"], "preserved")
+        fused_permute = graph_module.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.permute_copy.default
+        )[0]
+        self.assertIs(fused_permute.meta, output_meta)
+        self.assertEqual(fused_permute.meta["apply_inplace_test"], "output")
+
     def test_cascaded_permutes_multiple_users(self) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn(2, 3, 4, 5))