diff --git a/exir/passes/memory_format_ops_pass.py b/exir/passes/memory_format_ops_pass.py index 13468dfd8d8..dc8781110b5 100644 --- a/exir/passes/memory_format_ops_pass.py +++ b/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 = DimOrderOpsMap.keys() + 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 = MemoryFormatOpsMap.keys() + 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, diff --git a/exir/passes/normalize_transpose_pass.py b/exir/passes/normalize_transpose_pass.py index b401d2ad9b7..8beed8e0004 100644 --- a/exir/passes/normalize_transpose_pass.py +++ b/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 ) diff --git a/exir/passes/remove_mixed_type_operators.py b/exir/passes/remove_mixed_type_operators.py index 86a71354337..0bf22da39e2 100644 --- a/exir/passes/remove_mixed_type_operators.py +++ b/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 diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 16ed5af4180..38cd3fe8782 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -22,7 +22,13 @@ ) from executorch.exir.pass_manager import ExportedProgramPassManager, PassManager from executorch.exir.passes import ScalarToTensorPass +from executorch.exir.passes.memory_format_ops_pass import ( + DimOrderOpsRevertPass, + MemoryFormatOpsPass, +) +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 +235,143 @@ def test_rejects_implicit_symbolic_scalar_coercions(self) -> None: float(ProxyValue(sym_float, torch.fx.Graph().placeholder("x"))) +class TestExportPassTargetedOps(unittest.TestCase): + def test_memory_format_passes_rewrite_only_targeted_ops(self) -> None: + class MemoryFormatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x.to(memory_format=torch.channels_last)) + + sample_input = torch.randn(1, 2, 3, 4) + exported = export(MemoryFormatModule(), (sample_input,), strict=True) + edge_program = to_edge( + exported, + compile_config=exir.EdgeCompileConfig(_skip_dim_order=True), + ).exported_program() + graph_module = edge_program.graph_module + + memory_format_result = MemoryFormatOpsPass()(graph_module).graph_module + self.assertEqual( + len( + memory_format_result.graph.find_nodes( + op="call_function", + target=exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + ) + ), + 1, + ) + self.assertEqual( + len( + memory_format_result.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.relu.default + ) + ), + 1, + ) + + reverted_result = DimOrderOpsRevertPass()(memory_format_result).graph_module + self.assertEqual( + len( + reverted_result.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten._to_copy.default + ) + ), + 1, + ) + self.assertEqual( + len( + reverted_result.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.relu.default + ) + ), + 1, + ) + torch.testing.assert_close( + reverted_result(sample_input)[0], + MemoryFormatModule()(sample_input), + ) + + 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: """