Add targeted_ops to core ExportPass passes for fast-copy optimization (#21702)

exir/passes/memory_format_ops_pass.py

@@ -29,8 +29,11 @@ class MemoryFormatOpsPass(ExportPass):
     the aten op and the new edge dialect dim_order op.
     """
 
+    enable_fast_copy = True
+    targeted_ops = frozenset(DimOrderOpsMap)
+
     def call_operator(self, op, args, kwargs, meta):
-        if not (isinstance(op, EdgeOpOverload) and op in DimOrderOpsMap):
+        if not (isinstance(op, EdgeOpOverload) and op in self.targeted_ops):
             return super().call_operator(
                 op,
                 args,
@@ -96,8 +99,11 @@ class DimOrderOpsRevertPass(ExportPass):
     This pass is to revert the dim_order ops back to the memory format ops.
     """
 
+    enable_fast_copy = True
+    targeted_ops = frozenset(MemoryFormatOpsMap)
+
     def call_operator(self, op, args, kwargs, meta):
-        if not (isinstance(op, EdgeOpOverload) and op in MemoryFormatOpsMap):
+        if not (isinstance(op, EdgeOpOverload) and op in self.targeted_ops):
             return super().call_operator(
                 op,
                 args,

exir/passes/normalize_transpose_pass.py

@@ -13,11 +13,15 @@ class NormalizeTransposePass(ExportPass):
     Even with functionalization on, we still get graph with
     torch.ops.aten.t.default op. Ideally we should fix functionalization.
     TODO: once we have that, we should remove this pass.
-    Check test_normalize_transpose_op in test_passes.py for more details
+    Check test_normalize_transpose_rewrites_transpose_to_copy in test_pass_infra.py
+    for more details.
     """
 
+    enable_fast_copy = True
+    targeted_ops = frozenset({torch.ops.aten.t.default})
+
     def call_operator(self, op, args, kwargs, meta):
-        if op == torch.ops.aten.t.default:
+        if op in self.targeted_ops:
             return super().call_operator(
                 torch.ops.aten.t_copy.default, (args[0],), kwargs, meta
             )

exir/passes/remove_mixed_type_operators.py

@@ -6,6 +6,8 @@
 
 # pyre-strict
 
+from types import MappingProxyType
+
 import torch
 from executorch.exir.pass_base import ExportPass, map_args, NodeMetadata, ProxyValue
 from torch import SymBool, SymFloat, SymInt
@@ -14,13 +16,8 @@
 
 
 class RemoveMixedTypeOperators(ExportPass):
-    # pyre-ignore
-    def call_operator(self, op, args, kwargs, meta: NodeMetadata):  # noqa: C901
-        if len(args) <= 1:
-            # Unary Operators are not mixed type
-            return super().call_operator(op, args, kwargs, meta)
-
-        promotion_type_allow_list = {
+    promotion_type_allow_list = MappingProxyType(
+        {
             torch.ops.aten.add.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
             torch.ops.aten.mul.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
             torch.ops.aten.sub.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
@@ -30,9 +27,18 @@ def call_operator(self, op, args, kwargs, meta: NodeMetadata):  # noqa: C901
             torch.ops.aten.div.Tensor_mode: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
             torch.ops.aten.minimum.default: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
         }
+    )
+    enable_fast_copy = True
+    targeted_ops = frozenset(promotion_type_allow_list)
+
+    # pyre-ignore
+    def call_operator(self, op, args, kwargs, meta: NodeMetadata):  # noqa: C901
+        if len(args) <= 1:
+            # Unary Operators are not mixed type
+            return super().call_operator(op, args, kwargs, meta)
 
-        if op in promotion_type_allow_list:
-            promotion_kind = promotion_type_allow_list[op]
+        if op in self.promotion_type_allow_list:
+            promotion_kind = self.promotion_type_allow_list[op]
             if (
                 op == torch.ops.aten.div.Tensor_mode
                 and kwargs.get("rounding_mode") is None

exir/tests/test_pass_infra.py

@@ -22,7 +22,9 @@
 )
 from executorch.exir.pass_manager import ExportedProgramPassManager, PassManager
 from executorch.exir.passes import ScalarToTensorPass
+from executorch.exir.passes.normalize_transpose_pass import NormalizeTransposePass
 from executorch.exir.passes.pass_registry import PassRegistry
+from executorch.exir.passes.remove_mixed_type_operators import RemoveMixedTypeOperators
 from executorch.exir.program import to_edge
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.export import Dim, export, ExportedProgram
@@ -229,6 +231,89 @@ def test_rejects_implicit_symbolic_scalar_coercions(self) -> None:
             float(ProxyValue(sym_float, torch.fx.Graph().placeholder("x")))
 
 
+class TestExportPassTargetedOps(unittest.TestCase):
+    def test_normalize_transpose_rewrites_transpose_to_copy(self) -> None:
+        class TransposeModule(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return torch.relu(torch.ops.aten.t.default(x))
+
+        graph_module = export(
+            TransposeModule(), (torch.randn(3, 4),), strict=True
+        ).module()
+        self.assertEqual(
+            len(
+                graph_module.graph.find_nodes(
+                    op="call_function", target=torch.ops.aten.t.default
+                )
+            ),
+            1,
+        )
+
+        new_graph_module = NormalizeTransposePass()(graph_module).graph_module
+
+        self.assertEqual(
+            len(
+                new_graph_module.graph.find_nodes(
+                    op="call_function", target=torch.ops.aten.t.default
+                )
+            ),
+            0,
+        )
+        self.assertEqual(
+            len(
+                new_graph_module.graph.find_nodes(
+                    op="call_function", target=torch.ops.aten.t_copy.default
+                )
+            ),
+            1,
+        )
+        self.assertEqual(
+            len(
+                new_graph_module.graph.find_nodes(
+                    op="call_function", target=torch.ops.aten.relu.default
+                )
+            ),
+            1,
+        )
+
+    def test_remove_mixed_type_operators_promotes_operands(self) -> None:
+        class AddModule(torch.nn.Module):
+            def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+                return torch.relu(x + y)
+
+        int_tensor = torch.tensor([[1, 2, 3]], dtype=torch.int64)
+        float_tensor = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float)
+        graph_module = export(
+            AddModule(), (int_tensor, float_tensor), strict=True
+        ).module()
+        add_node = graph_module.graph.find_nodes(
+            op="call_function", target=torch.ops.aten.add.Tensor
+        )[0]
+        self.assertEqual(add_node.args[0].meta["val"].dtype, torch.int64)
+        self.assertEqual(add_node.args[1].meta["val"].dtype, torch.float)
+
+        new_graph_module = RemoveMixedTypeOperators()(graph_module).graph_module
+
+        add_nodes = new_graph_module.graph.find_nodes(
+            op="call_function", target=torch.ops.aten.add.Tensor
+        )
+        to_copy_nodes = new_graph_module.graph.find_nodes(
+            op="call_function", target=torch.ops.aten._to_copy.default
+        )
+        relu_nodes = new_graph_module.graph.find_nodes(
+            op="call_function", target=torch.ops.aten.relu.default
+        )
+        self.assertEqual(len(add_nodes), 1)
+        self.assertEqual(len(to_copy_nodes), 1)
+        self.assertEqual(len(relu_nodes), 1)
+        for arg in add_nodes[0].args:
+            self.assertEqual(arg.meta["val"].dtype, torch.float)
+        torch.testing.assert_close(
+            new_graph_module(int_tensor, float_tensor),
+            AddModule()(int_tensor, float_tensor),
+        )
+
+
 class TestExportedProgramPassManager(unittest.TestCase):
     def test_runs_graph_module_passes_on_exported_program(self) -> None:
         """