diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index 2e3f49dd861..685559bb420 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -143,6 +143,45 @@ inline bool is_channels_last_tensor(const Tensor& tensor) { return tensor.dim_order() == channels_last_order; } +// Strict channels_last check: true only when the tensor explicitly carries the +// channels_last dim_order [0, 2, 3, 1]. Unlike is_channels_last_tensor(), it +// does NOT treat unit-dim tensors as ambiguously channels_last, so a +// dim-order-free contiguous NHWC tensor (dim_order [0, 1, 2, 3]) is never +// misclassified. +inline bool has_channels_last_dim_order(const Tensor& tensor) { + if (tensor.dim() != 4) { + return false; + } + constexpr executorch::aten::DimOrderType kChannelsLastDimOrder[] = { + 0, 2, 3, 1}; + executorch::aten::ArrayRef + channels_last_order(kChannelsLastDimOrder, 4); + return tensor.dim_order() == channels_last_order; +} + +// Logical (batch, height, width, channels) of a 4-D activation tensor under +// either supported contract: +// - legacy channels_last: logical [N, C, H, W] (dim_order [0, 2, 3, 1]); +// - dim-order-free: logical [N, H, W, C] (plain contiguous). +// The CMSIS-NN kernels always operate on physically-NHWC data, so reading dims +// this way lets them serve both previously-generated channels_last .pte files +// and new contiguous ones through one code path. +struct NhwcDims { + int64_t n; + int64_t h; + int64_t w; + int64_t c; +}; + +inline NhwcDims read_nhwc_dims(const Tensor& tensor) { + if (has_channels_last_dim_order(tensor)) { + return NhwcDims{ + tensor.size(0), tensor.size(2), tensor.size(3), tensor.size(1)}; + } + return NhwcDims{ + tensor.size(0), tensor.size(1), tensor.size(2), tensor.size(3)}; +} + inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { if (tensor1.dim() != tensor2.dim()) { return false; @@ -203,7 +242,6 @@ inline bool prepare_cmsis_pool2d_config( int64_t activation_min, int64_t activation_max, CmsisPool2DConfig& config, - bool require_channels_last = true, bool allow_ceil_mode = false) { if (input.dim() != 4 || output.dim() != 4) { ET_LOG(Error, "%s: tensors must be 4-D", op_name); @@ -218,7 +256,11 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (input.size(0) != output.size(0) || input.size(1) != output.size(1)) { + // Read dims under either the legacy channels_last or the dim-order-free + // contiguous contract (both physically NHWC for CMSIS-NN). + const NhwcDims in_dims = read_nhwc_dims(input); + const NhwcDims out_dims = read_nhwc_dims(output); + if (in_dims.n != out_dims.n || in_dims.c != out_dims.c) { ET_LOG( Error, "%s: batch and channel dimensions must match between input and output", @@ -227,15 +269,6 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (require_channels_last) { - if (!is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { - ET_LOG( - Error, "%s: tensors must use channels_last dimension order", op_name); - context.fail(Error::InvalidArgument); - return false; - } - } - auto check_tuple_len = [&](const Int64ArrayRef& arr, const char* name) -> bool { if (arr.size() != 2) { @@ -314,17 +347,17 @@ inline bool prepare_cmsis_pool2d_config( int32_t batch, channels, input_h, input_w, output_h, output_w; if (!check_int32_within_range( - context, op_name, input.size(0), "input batch", batch) || + context, op_name, in_dims.n, "input batch", batch) || !check_int32_within_range( - context, op_name, input.size(1), "input channels", channels) || + context, op_name, in_dims.h, "input height", input_h) || !check_int32_within_range( - context, op_name, input.size(2), "input height", input_h) || + context, op_name, in_dims.w, "input width", input_w) || !check_int32_within_range( - context, op_name, input.size(3), "input width", input_w) || + context, op_name, in_dims.c, "input channels", channels) || !check_int32_within_range( - context, op_name, output.size(2), "output height", output_h) || + context, op_name, out_dims.h, "output height", output_h) || !check_int32_within_range( - context, op_name, output.size(3), "output width", output_w)) { + context, op_name, out_dims.w, "output width", output_w)) { return false; } diff --git a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp index 39b6432c45a..1d15d4f843a 100644 --- a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp @@ -97,8 +97,7 @@ Tensor& quantized_avg_pool2d_out( activation_min, activation_max, pool_config, - true, - true)) { + /*allow_ceil_mode=*/true)) { return out; } diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 0c731911e44..da2a8b05fa2 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -1317,19 +1317,23 @@ def quantized_avg_pool2d_meta( kernel = _ensure_tuple2(kernel_size) stride_vals = _ensure_tuple2(stride) padding_vals = _ensure_tuple2(padding) + # Input is NHWC-shaped [N, H, W, C] and plain contiguous. Derive the output + # spatial dims via the NCHW pooling helper, then emit an NHWC-shaped output. + n, h, w, c = input.shape + nchw = torch.empty((n, c, h, w), dtype=torch.float, device=input.device) output = F.avg_pool2d( - input.to(torch.float), + nchw, kernel, stride=stride_vals, padding=padding_vals, ceil_mode=ceil_mode, count_include_pad=False, ) + out_n, out_c, out_h, out_w = output.shape return torch.empty( - output.shape, + (out_n, out_h, out_w, out_c), dtype=torch.int8, device=input.device, - memory_format=torch.channels_last, ) @@ -1346,6 +1350,8 @@ def quantized_avg_pool2d_impl( scratch: torch.Tensor, ) -> torch.Tensor: dequant_input = dequantize_per_tensor_cmsis(input, zero_point, multiplier, shift) + # Input is NHWC-shaped [N, H, W, C]; convert to NCHW for the reference pool. + dequant_input = dequant_input.permute(0, 3, 1, 2).contiguous() kernel = _ensure_tuple2(kernel_size) stride_vals = _ensure_tuple2(stride) @@ -1362,7 +1368,8 @@ def quantized_avg_pool2d_impl( ) result = quantize_per_tensor_cmsis(result, zero_point, multiplier, shift) output = torch.clamp(result, -128, 127) - return output.to(torch.int8) + # result is NCHW; return NHWC-shaped output to match the kernel contract. + return output.to(torch.int8).permute(0, 2, 3, 1).contiguous() # =================================================================== @@ -1416,10 +1423,11 @@ def _compute_max_pool2d_output_shape( padding: Sequence[int], dilation: Sequence[int], ) -> torch.Size: + # Input is NHWC-shaped [N, H, W, C]. batch = input_shape[0] - channels = input_shape[1] - in_height = input_shape[2] - in_width = input_shape[3] + in_height = input_shape[1] + in_width = input_shape[2] + channels = input_shape[3] kernel_height, kernel_width = kernel_size stride_h, stride_w = stride @@ -1432,7 +1440,7 @@ def _compute_max_pool2d_output_shape( out_width = ( in_width + 2 * pad_w - dilation_w * (kernel_width - 1) - 1 ) // stride_w + 1 - return torch.Size([batch, channels, out_height, out_width]) + return torch.Size([batch, out_height, out_width, channels]) @register_fake("cortex_m::quantized_max_pool2d") # type: ignore[misc] @@ -1453,6 +1461,7 @@ def quantized_max_pool2d_meta( padding_vals = _ensure_tuple2(padding) dilation_vals = _ensure_tuple2(dilation) + # Input/output are NHWC-shaped [N, H, W, C] and plain contiguous. output_shape = _compute_max_pool2d_output_shape( input.shape, kernel, stride_vals, padding_vals, dilation_vals ) @@ -1460,7 +1469,6 @@ def quantized_max_pool2d_meta( output_shape, dtype=torch.int8, device=input.device, - memory_format=torch.channels_last, ) @@ -1495,8 +1503,10 @@ def quantized_max_pool2d_impl( if ceil_mode: raise RuntimeError("quantized_max_pool2d does not support ceil_mode=True") + # Input is NHWC-shaped [N, H, W, C]; convert to NCHW for the reference pool. + input_nchw = input.permute(0, 3, 1, 2).contiguous() result = F.max_pool2d( - input, + input_nchw, kernel, stride=stride_vals, padding=padding_vals, @@ -1504,4 +1514,5 @@ def quantized_max_pool2d_impl( ceil_mode=ceil_mode, ) result = torch.clamp(result, activation_min, activation_max) - return result.to(torch.int8).contiguous(memory_format=torch.channels_last) + # result is NCHW; return NHWC-shaped output to match the kernel contract. + return result.to(torch.int8).permute(0, 2, 3, 1).contiguous() diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index a3af7b5bc1b..fd561ea8a68 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -828,23 +828,32 @@ def _get_avg_pool2d_replacement( input_scale = node.meta["input_qparams"][0].scale output_mult, output_shift = quantize_multiplier_aot(input_scale) + # The cortex_m pool op consumes/produces NHWC-shaped contiguous tensors, so + # permute the NCHW graph input to NHWC before the op. + with node.graph.inserting_before(node): + pool_input = node.graph.create_node( + "call_function", + target=exir_ops.edge.aten.permute_copy.default, + args=(input_node, [0, 2, 3, 1]), + ) + avg_padding = padding if count_include_pad: pad_h, pad_w = padding - input_tensor = get_first_fake_tensor(input_node) - pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) + # Input is now NHWC [N, H, W, C]; pad H and W in NHWC physical order. + pre_pad = post_pad = [0, int(pad_h), int(pad_w), 0] with node.graph.inserting_before(node): - input_node = node.graph.create_node( + pool_input = node.graph.create_node( "call_function", target=exir_ops.edge.cortex_m.pad.default, - args=(input_node, pre_pad, post_pad, int(input_zp)), + args=(pool_input, pre_pad, post_pad, int(input_zp)), ) avg_padding = [0, 0] scratch = _create_uninitialized_alloc_node(node, exported_program) new_args = ( - input_node, + pool_input, kernel_size, stride, avg_padding, @@ -854,8 +863,19 @@ def _get_avg_pool2d_replacement( int(output_shift), scratch, ) + with node.graph.inserting_before(node): + pool_node = node.graph.create_node( + "call_function", + target=exir_ops.edge.cortex_m.quantized_avg_pool2d.default, + args=new_args, + ) + for key in ("input_qparams", "output_qparams"): + if key in node.meta: + pool_node.meta[key] = node.meta[key] + + # Permute the NHWC-shaped op output back to NCHW so downstream stays NCHW. return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_avg_pool2d.default, new_args + exir_ops.edge.aten.permute_copy.default, (pool_node, [0, 3, 1, 2]) ) @@ -1090,19 +1110,39 @@ def _get_max_pool2d_replacement( if quantized_op is None: return None - args = ( - node.args[0], - kernel_size, - stride, - padding, - dilation, - ceil_mode, - input_zero_point, - output_zero_point, - activation_min, - activation_max, + input_node = cast(Node, node.args[0]) + # The cortex_m pool op consumes/produces NHWC-shaped contiguous tensors, so + # permute the NCHW graph input to NHWC before the op. + with node.graph.inserting_before(node): + pool_input = node.graph.create_node( + "call_function", + target=exir_ops.edge.aten.permute_copy.default, + args=(input_node, [0, 2, 3, 1]), + ) + pool_node = node.graph.create_node( + "call_function", + target=quantized_op.default, + args=( + pool_input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ), + ) + for key in ("input_qparams", "output_qparams"): + if key in node.meta: + pool_node.meta[key] = node.meta[key] + + # Permute the NHWC-shaped op output back to NCHW so downstream stays NCHW. + return DialectNodeSpec( + exir_ops.edge.aten.permute_copy.default, (pool_node, [0, 3, 1, 2]) ) - return DialectNodeSpec(quantized_op.default, args) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.minimum.default) diff --git a/backends/cortex_m/passes/scratch_buffer_sizes.py b/backends/cortex_m/passes/scratch_buffer_sizes.py index b247e2be944..ef70b274b9e 100644 --- a/backends/cortex_m/passes/scratch_buffer_sizes.py +++ b/backends/cortex_m/passes/scratch_buffer_sizes.py @@ -251,10 +251,10 @@ def cmsis_nn_avgpool_buffer_size( ) -> list[int]: x = cast(torch.fx.Node, pool_node.args[0]) - # Input is NCHW (PyTorch); CMSIS-NN's avgpool buffer sizer only needs the - # input channel count and output width. - _, c_in, _, _ = _shape_from_node(x) - _, _, _, out_w = _shape_from_node(pool_node) + # The cortex_m pool op consumes/produces NHWC-shaped [N, H, W, C] tensors. + # The avgpool buffer sizer only needs the input channel count and output width. + _, _, _, c_in = _shape_from_node(x) + _, _, out_w, _ = _shape_from_node(pool_node) return [ int( diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index 4d8502ec59e..ed82519c4ee 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -43,10 +43,14 @@ join_by_comma() { ops_list=( aten::add.out aten::clamp.out + aten::clone.out aten::mul.out aten::convolution.out aten::max_pool2d_with_indices.out + # Keep legacy PTEs generated before _skip_dim_order=True loadable. dim_order_ops::_clone_dim_order.out + dim_order_ops::_empty_dim_order.out + dim_order_ops::_to_dim_order_copy.out aten::cat.out aten::full.out aten::ge.Tensor_out diff --git a/backends/cortex_m/test/misc/test_portable_int8.py b/backends/cortex_m/test/misc/test_portable_int8.py index 6efeec9e5b0..838a660ef29 100644 --- a/backends/cortex_m/test/misc/test_portable_int8.py +++ b/backends/cortex_m/test/misc/test_portable_int8.py @@ -17,9 +17,10 @@ from executorch.backends.arm.quantizer.arm_quantizer_utils import SharedQspecQuantizer from executorch.backends.arm.test.common import parametrize, xfail_type from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer -from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.cortex_m.test.tester import CortexMTester, CortexMToEdge from executorch.backends.test.harness.stages import StageType from executorch.exir import EdgeCompileConfig +from executorch.exir._serialize import _deserialize_pte_binary from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -784,6 +785,28 @@ def test_shared_qspec_portable_int8_ops_fvp(op_case: OpCase) -> None: tester.test_implementation() +def test_legacy_dim_order_pte_runs_on_cortex_m() -> None: + op_case = OP_CASES["clone"] + tester = CortexMTester(op_case.module, op_case.example_inputs) + tester.quantize() + tester.export() + tester.to_edge(CortexMToEdge(skip_dim_order=False)) + tester.check_count( + {"executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1} + ) + tester.run_passes() + tester.to_executorch() + tester.serialize() + + pte = _deserialize_pte_binary(tester.get_artifact(StageType.SERIALIZE)) + serialized_ops = { + (op.name, op.overload) for op in pte.program.execution_plan[0].operators + } + assert ("dim_order_ops::_clone_dim_order", "out") in serialized_ops + + tester.run_method_and_compare_outputs(inputs=tester.example_inputs) + + def test_shared_qspec_ops_default_covered() -> None: expected = set(SharedQspecQuantizer.SHARED_QSPEC_OPS_DEFAULT) covered = {case.target for case in OP_CASES.values()} diff --git a/backends/cortex_m/test/misc/test_quantization.py b/backends/cortex_m/test/misc/test_quantization.py index 6f27d1a6c28..c73ac8430a6 100644 --- a/backends/cortex_m/test/misc/test_quantization.py +++ b/backends/cortex_m/test/misc/test_quantization.py @@ -21,7 +21,7 @@ class SharedQspecMulipleClusters(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_add_Tensor": 2, "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 4, + "executorch_exir_dialects_edge__ops_aten_clone_default": 4, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 8, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 8, } @@ -30,7 +30,7 @@ class SharedQspecMulipleClusters(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 2, "executorch_exir_dialects_edge__ops_cortex_m_transpose_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 4, + "executorch_exir_dialects_edge__ops_aten_clone_default": 4, } def forward(self, x): @@ -71,7 +71,7 @@ class SharedQspecInputForkShared(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_minimum_default": 1, "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 5, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 5, } @@ -80,7 +80,7 @@ class SharedQspecInputForkShared(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_minimum_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, "executorch_exir_dialects_edge__ops_cortex_m_transpose_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } def forward(self, x, y): @@ -96,7 +96,7 @@ class SharedQspecInputForkXShared(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_maximum_default": 1, "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 4, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 4, } @@ -105,7 +105,7 @@ class SharedQspecInputForkXShared(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_maximum_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, "executorch_exir_dialects_edge__ops_cortex_m_transpose_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } def forward(self, x, y): @@ -120,7 +120,7 @@ class SharedQspecInputForkYShared(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_minimum_default": 1, "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 5, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 5, } @@ -129,7 +129,7 @@ class SharedQspecInputForkYShared(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_minimum_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } def forward(self, x, y): @@ -187,7 +187,7 @@ class SharedQspecOutputForkShared(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 6, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 4, } @@ -196,7 +196,7 @@ class SharedQspecOutputForkShared(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_transpose_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } def forward(self, x): @@ -213,7 +213,7 @@ class SharedQspecManyForks(torch.nn.Module): "executorch_exir_dialects_edge__ops_aten_maximum_default": 2, "executorch_exir_dialects_edge__ops_aten_minimum_default": 1, "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 9, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 6, } @@ -223,7 +223,7 @@ class SharedQspecManyForks(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_minimum_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_transpose_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } def forward(self, x): @@ -239,7 +239,7 @@ class SharedQspecSurroundedQuantizedOp(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_add_Tensor": 1, "executorch_exir_dialects_edge__ops_aten_maximum_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 5, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 4, } @@ -248,7 +248,7 @@ class SharedQspecSurroundedQuantizedOp(torch.nn.Module): "executorch_exir_dialects_edge__ops_cortex_m_maximum_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } def forward(self, x): diff --git a/backends/cortex_m/test/models/test_ds_cnn.py b/backends/cortex_m/test/models/test_ds_cnn.py index 206af19a61e..264c655a106 100644 --- a/backends/cortex_m/test/models/test_ds_cnn.py +++ b/backends/cortex_m/test/models/test_ds_cnn.py @@ -15,7 +15,7 @@ "executorch_exir_dialects_edge__ops_aten_linear_default": 1, "executorch_exir_dialects_edge__ops_aten_relu_default": 9, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 2, + "executorch_exir_dialects_edge__ops_aten_clone_default": 2, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 18, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 17, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 15, @@ -30,7 +30,7 @@ "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 4, "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 5, "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 2, + "executorch_exir_dialects_edge__ops_aten_clone_default": 2, } test_cases = { diff --git a/backends/cortex_m/test/models/test_mobilenet_v2.py b/backends/cortex_m/test/models/test_mobilenet_v2.py index 67f0937a006..c2f121d5f16 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v2.py +++ b/backends/cortex_m/test/models/test_mobilenet_v2.py @@ -20,7 +20,7 @@ "executorch_exir_dialects_edge__ops_aten_hardtanh_default": 35, "executorch_exir_dialects_edge__ops_aten_linear_default": 1, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 104, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 79, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 67, @@ -35,7 +35,7 @@ "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 35, "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 17, "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, + "executorch_exir_dialects_edge__ops_aten_clone_default": 1, } # Use larger sample set for calibration to get better quantization diff --git a/backends/cortex_m/test/ops/test_avg_pool2d.py b/backends/cortex_m/test/ops/test_avg_pool2d.py index 2fac40cdd9c..8fb8f89ce9e 100644 --- a/backends/cortex_m/test/ops/test_avg_pool2d.py +++ b/backends/cortex_m/test/ops/test_avg_pool2d.py @@ -72,6 +72,18 @@ def forward(self, x): # noqa: D102 CortexMAvgPool2d(kernel_size=3, stride=2, padding=1, count_include_pad=True), (ramp_tensor(0, 15, (1, 1, 4, 4)),), ), + # Multi-channel coverage in both memory formats. The cortex_m pool op contract + # is NHWC-shaped contiguous; the explicit permute inserted by the AoT pass + # normalizes both contiguous and channels_last inputs, so both should lower and + # run identically. + "avgpool_2x2_mc": McuTestCase( + CortexMAvgPool2d(kernel_size=2, stride=2), + (ramp_tensor(0, 47, (1, 3, 4, 4)),), + ), + "avgpool_2x2_mc_channels_last": McuTestCase( + CortexMAvgPool2d(kernel_size=2, stride=2), + (ramp_tensor(0, 47, (1, 3, 4, 4)).to(memory_format=torch.channels_last),), + ), } @@ -136,13 +148,15 @@ def test_dialect_avg_pool2d(test_case, cortex_m_target): scratch_size = scratch_arg.args[0][0][0] input_node = pool_node.args[0] + # The cortex_m pool op consumes/produces NHWC-shaped [N, H, W, C] tensors, so + # channels are at index 3 and output width at index 2. input_shape = input_node.meta["val"].shape output_shape = pool_node.meta["val"].shape expected_size = cmsis_nn.avgpool_buffer_size( cortex_m_target.backend, cmsis_nn.DataType.A8W8, - dim_dst_width=int(output_shape[3]), - ch_src=int(input_shape[1]), + dim_dst_width=int(output_shape[2]), + ch_src=int(input_shape[3]), ) assert ( scratch_size == expected_size diff --git a/backends/cortex_m/test/ops/test_max_pool2d.py b/backends/cortex_m/test/ops/test_max_pool2d.py index a67dd6b6e01..54399efbf5a 100644 --- a/backends/cortex_m/test/ops/test_max_pool2d.py +++ b/backends/cortex_m/test/ops/test_max_pool2d.py @@ -64,6 +64,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: CortexMMaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)), (ramp_tensor(-16, 16, (1, 1, 6, 6)),), ), + # Multi-channel coverage in both memory formats (NHWC-shaped contract; the AoT + # permute normalizes contiguous and channels_last inputs identically). + "maxpool_2x2_mc": McuTestCase( + CortexMMaxPool2d(kernel_size=2, stride=2), + (ramp_tensor(-50, 50, (1, 3, 6, 6)),), + ), + "maxpool_2x2_mc_channels_last": McuTestCase( + CortexMMaxPool2d(kernel_size=2, stride=2), + (ramp_tensor(-50, 50, (1, 3, 6, 6)).to(memory_format=torch.channels_last),), + ), "maxpool_2x2_indices": McuTestCase( CortexMMaxPool2dIndices(kernel_size=2, stride=2), (ramp_tensor(-50, 50, (1, 1, 6, 6)),), diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index 1f7d7f3059d..7020850e607 100644 --- a/backends/cortex_m/test/tester.py +++ b/backends/cortex_m/test/tester.py @@ -33,28 +33,36 @@ def __init__(self, calibration_samples=None): super().__init__(quantizer, calibration_samples=calibration_samples) +def get_cortex_m_edge_compile_config( + *, skip_dim_order: bool = True +) -> EdgeCompileConfig: + return EdgeCompileConfig( + preserve_ops=[ + torch.ops.aten.linear.default, + torch.ops.aten.hardsigmoid.default, + torch.ops.aten.hardsigmoid_.default, + torch.ops.aten.hardswish.default, + torch.ops.aten.hardswish_.default, + # silu naturally decomposes to sigmoid*x at the to_edge step. + # Preserve it so the LUT lowering can collapse it into a single + # cortex_m.quantized_activation call rather than emitting an + # extra elementwise mul. Set globally because no per-test + # opt-out exists today; any new cortex_m test that uses SiLU + # must therefore expect a single aten.silu op in the edge graph + # (not sigmoid+mul). + torch.ops.aten.silu.default, + ], + _check_ir_validity=False, + _skip_dim_order=skip_dim_order, + _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], + ) + + class CortexMToEdge(ToEdge): - def __init__(self): - config = EdgeCompileConfig( - preserve_ops=[ - torch.ops.aten.linear.default, - torch.ops.aten.hardsigmoid.default, - torch.ops.aten.hardsigmoid_.default, - torch.ops.aten.hardswish.default, - torch.ops.aten.hardswish_.default, - # silu naturally decomposes to sigmoid*x at the to_edge step. - # Preserve it so the LUT lowering can collapse it into a single - # cortex_m.quantized_activation call rather than emitting an - # extra elementwise mul. Set globally because no per-test - # opt-out exists today; any new cortex_m test that uses SiLU - # must therefore expect a single aten.silu op in the edge graph - # (not sigmoid+mul). - torch.ops.aten.silu.default, - ], - _check_ir_validity=False, - _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], + def __init__(self, *, skip_dim_order: bool = True): + super().__init__( + get_cortex_m_edge_compile_config(skip_dim_order=skip_dim_order) ) - super().__init__(config) class CortexMRunPasses(RunPasses):