diff --git a/backends/cadence/aot/BUCK b/backends/cadence/aot/BUCK index 765f7b7f1e9..151127f77a3 100644 --- a/backends/cadence/aot/BUCK +++ b/backends/cadence/aot/BUCK @@ -502,6 +502,7 @@ fbcode_target(_kind = python_unittest, "//executorch/backends/cadence/aot:pass_utils", "//executorch/backends/cadence/aot:remove_ops", "//executorch/backends/cadence/aot/quantizer:quantizer", + "//executorch/exir:lib", "//executorch/exir/dialects:lib", ], ) diff --git a/backends/cadence/aot/passes.py b/backends/cadence/aot/passes.py index 73bc3fbb686..63701e26c6d 100644 --- a/backends/cadence/aot/passes.py +++ b/backends/cadence/aot/passes.py @@ -21,7 +21,7 @@ from executorch.backends.cadence.aot.remove_ops import ( CadenceRemoveNops, RemoveAliasCopyOpPass, - RemoveCloneOpsTransformImported, + RemoveCloneOpsPass, RemoveDetachCopyPass, RemoveNopExpandOpPass, RemoveNopSliceOrViewOpPass, @@ -114,7 +114,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: FinalizePipeline, FuseSliceSameDimPass, RemoveAliasCopyOpPass, - RemoveCloneOpsTransformImported, + RemoveCloneOpsPass, RemoveDetachCopyPass, RemoveNopExpandOpPass, RemoveToOpsPass, diff --git a/backends/cadence/aot/remove_ops.py b/backends/cadence/aot/remove_ops.py index 78144cf583c..80b591cc1b2 100644 --- a/backends/cadence/aot/remove_ops.py +++ b/backends/cadence/aot/remove_ops.py @@ -35,7 +35,6 @@ ExportPass, PassResult, ) -from executorch.exir.pass_manager import PassManager, PassType from executorch.exir.passes import dead_code_elimination_pass from torch.export import ExportedProgram from torch.export.graph_signature import InputKind, OutputKind @@ -44,12 +43,32 @@ from torch.utils import _pytree as pytree -class RemoveCloneOpsTransformImported(ExportPass): +class RemoveCloneOpsPass(RemoveCloneOpsTransform): + """ + Remove clones, including those that only request a memory format. + + Cadence lowers with ``_skip_dim_order=True``: every edge tensor is + contiguous and a memory format has no runtime representation, so such a + clone is a plain copy. Left in the graph, the ExecuTorch emitter rejects + any format other than contiguous or preserve. + + Before: + clone(x, memory_format=torch.channels_last) + + After: + x + """ + + def __init__(self) -> None: + super().__init__(eliminate_quant_dequant_pairs=False) + + def _is_non_identity_clone(self, node: torch.fx.Node) -> bool: + if node.target == exir_ops.edge.aten.clone.default: + return False + return super()._is_non_identity_clone(node) + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - finalize_passes: List[PassType] = [ - RemoveCloneOpsTransform(eliminate_quant_dequant_pairs=False), - ] - result = PassManager(passes=finalize_passes)(graph_module) + result = super().call(graph_module) dead_code_elimination_pass(result.graph_module) return result @@ -911,7 +930,7 @@ class CommonRemovePasses: FuseTransposeOrPermuteOpPairsPass, RemoveSqueezeViewBeforeElementwiseOps, RemoveCatFromSliceCopyPass, - RemoveCloneOpsTransformImported, + RemoveCloneOpsPass, ] diff --git a/backends/cadence/aot/tests/test_remove_ops_passes.py b/backends/cadence/aot/tests/test_remove_ops_passes.py index 4143656ab36..1bbb3671731 100644 --- a/backends/cadence/aot/tests/test_remove_ops_passes.py +++ b/backends/cadence/aot/tests/test_remove_ops_passes.py @@ -20,7 +20,7 @@ RemoveAliasCopyOpPass, RemoveBranchedQuantDequant, RemoveCatFromSliceCopyPass, - RemoveCloneOpsTransformImported, + RemoveCloneOpsPass, RemoveContiguousOpPass, RemoveDetachCopyPass, RemoveNopAddOpPass, @@ -38,6 +38,7 @@ ) from executorch.backends.cadence.aot.typing_stubs import expand from executorch.backends.test.graph_builder import GraphBuilder +from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops from pyre_extensions import none_throws @@ -282,12 +283,66 @@ def test_remove_clone(self) -> None: clone = builder.call_operator(op=exir_ops.edge.aten.clone.default, args=(x,)) builder.output([clone]) original = builder.get_graph_module() - p = RemoveCloneOpsTransformImported() + p = RemoveCloneOpsPass() graph_after_passes = cast(PassResult, p(original)).graph_module self.assertEqual( count_node(graph_after_passes, torch.ops.aten.clone.default), 0 ) + @expand( + [ + ["channels_last", torch.channels_last], + ["contiguous", torch.contiguous_format], + ["preserve", torch.preserve_format], + ] + ) + def test_remove_clone_with_memory_format( + self, _name: str, memory_format: torch.memory_format + ) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn([1, 8, 4, 4], dtype=torch.float32)) + clone = builder.call_operator( + op=exir_ops.edge.aten.clone.default, + args=(x,), + kwargs={"memory_format": memory_format}, + ) + builder.output([clone]) + original = builder.get_graph_module() + p = RemoveCloneOpsPass() + result = cast(PassResult, p(copy.deepcopy(original))) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.clone.default), 0 + ) + validate( + original, + result.graph_module, + (torch.randn(1, 8, 4, 4),), + "RemoveCloneOpsPass", + ) + + def test_remove_clone_lets_channels_last_program_emit(self) -> None: + # rlx_arbitrary_image_stylization: a CPU-side + # clone(memory_format=torch.channels_last) failed in the emitter. + class ChannelsLastCopy(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.contiguous(memory_format=torch.channels_last) + 1.0 + + ep = torch.export.export(ChannelsLastCopy(), (torch.randn(1, 8, 4, 4),)) + config = EdgeCompileConfig(_skip_dim_order=True) + with self.assertRaisesRegex(Exception, "memory_format"): + to_edge(copy.deepcopy(ep), compile_config=config).to_executorch() + + edge = to_edge(ep, compile_config=config).transform([RemoveCloneOpsPass()]) + self.assertEqual( + count_node( + edge.exported_program().graph_module, + exir_ops.edge.aten.clone.default, + ), + 0, + ) + self.assertGreater(len(edge.to_executorch().buffer), 0) + def test_remove_contiguous(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn([3, 5], dtype=torch.float32))