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)) diff --git a/backends/cadence/fusion_g3/operators/op_cat.cpp b/backends/cadence/fusion_g3/operators/op_cat.cpp index b5fbf720187..b4c3730c4eb 100644 --- a/backends/cadence/fusion_g3/operators/op_cat.cpp +++ b/backends/cadence/fusion_g3/operators/op_cat.cpp @@ -33,6 +33,12 @@ namespace impl { namespace G3 { namespace native { +namespace { +// Upper bound on inputs for the xa_nn_cat path, whose per-input arrays live on +// the DSP thread stack. Larger cats take the generic loop. +constexpr size_t kMaxNnlibCatInputs = 16; +} // namespace + Tensor& cat_out( KernelRuntimeContext& ctx, ArrayRef tensors, @@ -42,7 +48,7 @@ Tensor& cat_out( dim += out.dim(); } - int kTensorDimensionLimit = executorch::runtime::kTensorDimensionLimit; + constexpr int kTensorDimensionLimit = executorch::runtime::kTensorDimensionLimit; #ifdef OP_ARG_CHECK @@ -76,34 +82,6 @@ Tensor& cat_out( return out; } - const signed char* inp_tensors[tensors.size()]; - const int* inp_tensors_shapes[tensors.size()]; - - int inp_shapes_size[tensors.size()]; - - int temp_sizes[tensors.size()][kTensorDimensionLimit]; - ArrayRef temp_size; - - for (int i = 0; i < tensors.size(); i++) { - inp_tensors[i] = tensors[i].const_data_ptr(); - temp_size = tensors[i].sizes(); - - for (int j = 0; j < temp_size.size(); j++) { - temp_sizes[i][j] = temp_size[j]; - } - inp_tensors_shapes[i] = temp_sizes[i]; // input shapes - inp_shapes_size[i] = temp_size.size(); // number of input dimensions - } - - signed char* out_data = out.mutable_data_ptr(); - - const ArrayRef out_size = out.sizes(); - int out_shapes[kTensorDimensionLimit]; - for (int i = 0; i < out_size.size(); i++) // output shapes - { - out_shapes[i] = out_size[i]; - } - bool all_same_dtype = true; for (int i = 0; i < tensors.size(); i++) { if (out.scalar_type() != tensors[i].scalar_type()) { @@ -120,7 +98,33 @@ Tensor& cat_out( out.scalar_type() == ScalarType::Byte || out.scalar_type() == ScalarType::Float; - if (all_same_dtype && supported_dtype) { + if (all_same_dtype && supported_dtype && + tensors.size() <= kMaxNnlibCatInputs) { + const signed char* inp_tensors[kMaxNnlibCatInputs]; + const int* inp_tensors_shapes[kMaxNnlibCatInputs]; + int inp_shapes_size[kMaxNnlibCatInputs]; + int temp_sizes[kMaxNnlibCatInputs][kTensorDimensionLimit]; + + for (int i = 0; i < tensors.size(); i++) { + inp_tensors[i] = tensors[i].const_data_ptr(); + const ArrayRef temp_size = tensors[i].sizes(); + + for (int j = 0; j < temp_size.size(); j++) { + temp_sizes[i][j] = temp_size[j]; + } + inp_tensors_shapes[i] = temp_sizes[i]; // input shapes + inp_shapes_size[i] = temp_size.size(); // number of input dimensions + } + + signed char* out_data = out.mutable_data_ptr(); + + const ArrayRef out_size = out.sizes(); + int out_shapes[kTensorDimensionLimit]; + for (int i = 0; i < out_size.size(); i++) // output shapes + { + out_shapes[i] = out_size[i]; + } + XT_KERNEL_CHECK( ctx, out,