From 8690ea3d21e6f98ef18ca6e2b0708693b46c0460 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:51:02 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- backends/cortex_m/ops/cortex_m_ops_common.h | 17 +- backends/cortex_m/ops/op_quantized_add.cpp | 129 +++++++++++++-- backends/cortex_m/ops/op_quantized_mul.cpp | 100 +++++++++++- backends/cortex_m/ops/operators.py | 151 ++++++++++++++++++ backends/cortex_m/ops/operators.yaml | 12 ++ backends/cortex_m/test/build_test_runner.sh | 2 + .../test/ops/test_explicit_nhwc_runtime.py | 60 +++++++ .../test/test_quantized_conv2d_layout.py | 58 +++++++ 8 files changed, 508 insertions(+), 21 deletions(-) diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index bcaed0a1bc7..d6086f0e9c1 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -149,7 +149,10 @@ inline bool is_channels_last_tensor(const Tensor& tensor) { return tensor.dim_order() == channels_last_order; } -inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { +inline bool is_channel_broadcast( + const Tensor& tensor1, + const Tensor& tensor2, + int64_t channel_dim) { if (tensor1.dim() != tensor2.dim()) { return false; } @@ -158,16 +161,22 @@ inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { return false; } - if (tensor1.size(1) != tensor2.size(1)) { + if (tensor1.size(channel_dim) != tensor2.size(channel_dim)) { return false; } - const bool tensor1_channels_only = tensor1.numel() == tensor1.size(1); - const bool tensor2_channels_only = tensor2.numel() == tensor2.size(1); + const bool tensor1_channels_only = + tensor1.numel() == tensor1.size(channel_dim); + const bool tensor2_channels_only = + tensor2.numel() == tensor2.size(channel_dim); return tensor1_channels_only || tensor2_channels_only; } +inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { + return is_channel_broadcast(tensor1, tensor2, 1); +} + inline bool check_int32_within_range( KernelRuntimeContext& context, const char* op_name, diff --git a/backends/cortex_m/ops/op_quantized_add.cpp b/backends/cortex_m/ops/op_quantized_add.cpp index f93bb6c1be9..43ddd2f86de 100644 --- a/backends/cortex_m/ops/op_quantized_add.cpp +++ b/backends/cortex_m/ops/op_quantized_add.cpp @@ -13,8 +13,7 @@ namespace cortex_m { namespace native { using KernelRuntimeContext = torch::executor::KernelRuntimeContext; -// cppcheck-suppress unusedFunction -Tensor& quantized_add_out( +static Tensor& quantized_add_out_impl( KernelRuntimeContext& context, const Tensor& input1_int8, const int64_t input1_zero_point, @@ -29,16 +28,51 @@ Tensor& quantized_add_out( const int64_t output_shift, const int64_t activation_min, const int64_t activation_max, + ActivationLayout layout, + const char* op_name, Tensor& out) { - // Validate tensor types and dim order - bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8); + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + bool channel_broadcast = + is_channel_broadcast(input1_int8, input2_int8, channel_dim); validate_cmsis_nn_tensor_requirements( input1_int8, input2_int8, out, ScalarType::Char, - /*require_channels_last=*/channel_broadcast, + /*require_channels_last=*/ + channel_broadcast && layout == ActivationLayout::NCHWLogical, /*require_same_sizes=*/!channel_broadcast); + if (layout == ActivationLayout::NHWCLogical) { + ET_CHECK_MSG( + input1_int8.dim() == 4 && input2_int8.dim() == 4 && out.dim() == 4, + "%s: tensors must be 4-D", + op_name); + ET_CHECK_MSG( + executorch::runtime::is_contiguous_dim_order( + input1_int8.dim_order().data(), input1_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + input2_int8.dim_order().data(), + input2_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + out.dim_order().data(), out.dim_order().size()), + "%s: tensors must use contiguous dimension order", + op_name); + } else if (channel_broadcast) { + ET_CHECK_MSG( + is_channels_last_tensor(input1_int8) && + is_channels_last_tensor(input2_int8) && + is_channels_last_tensor(out), + "%s: channel-broadcast tensors must use channels-last dimension order", + op_name); + } + if (channel_broadcast) { + const Tensor& full_input = + input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8; + ET_CHECK_MSG( + out.sizes() == full_input.sizes(), + "%s: output must have the broadcast result shape", + op_name); + } // Validate quantization parameters validate_quantization_params( @@ -54,7 +88,8 @@ Tensor& quantized_add_out( ET_LOG( Debug, - "quantized_add_out: input1_int8.sizes() = %zu", + "%s: input1_int8.sizes() = %zu", + op_name, input1_int8.sizes().size()); int32_t zp1 = static_cast(input1_zero_point); @@ -101,7 +136,7 @@ Tensor& quantized_add_out( std::swap(input1_shift_val, input2_shift_val); std::swap(input1_ptr, input2_ptr); } - adds_per_loop = input1_int8.size(1); + adds_per_loop = input1_int8.size(channel_dim); } else { adds_per_loop = out.numel(); } @@ -130,7 +165,8 @@ Tensor& quantized_add_out( if (status != ARM_CMSIS_NN_SUCCESS) { ET_LOG( Error, - "quantized_add_out: arm_elementwise_add_s8 failed with status [%d]", + "%s: arm_elementwise_add_s8 failed with status [%d]", + op_name, status); context.fail(Error::Internal); // Fail the execution context @@ -139,10 +175,85 @@ Tensor& quantized_add_out( } ET_LOG( Debug, - "quantized_add_out: Successfully completed with AoT-computed parameters!"); + "%s: Successfully completed with AoT-computed parameters!", + op_name); return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_add_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const int64_t input1_multiplier, + const int64_t input1_shift, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t input2_multiplier, + const int64_t input2_shift, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_add_out_impl( + context, + input1_int8, + input1_zero_point, + input1_multiplier, + input1_shift, + input2_int8, + input2_zero_point, + input2_multiplier, + input2_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ActivationLayout::NCHWLogical, + "quantized_add_out", + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_add_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const int64_t input1_multiplier, + const int64_t input1_shift, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t input2_multiplier, + const int64_t input2_shift, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_add_out_impl( + context, + input1_int8, + input1_zero_point, + input1_multiplier, + input1_shift, + input2_int8, + input2_zero_point, + input2_multiplier, + input2_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ActivationLayout::NHWCLogical, + "quantized_add_nhwc_out", + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_mul.cpp b/backends/cortex_m/ops/op_quantized_mul.cpp index 93ce2303d64..5b5f660b7d1 100644 --- a/backends/cortex_m/ops/op_quantized_mul.cpp +++ b/backends/cortex_m/ops/op_quantized_mul.cpp @@ -18,8 +18,7 @@ constexpr int32_t kInt8ActivationMax = std::numeric_limits::max(); using KernelRuntimeContext = torch::executor::KernelRuntimeContext; -// cppcheck-suppress unusedFunction -Tensor& quantized_mul_out( +static Tensor& quantized_mul_out_impl( KernelRuntimeContext& context, const Tensor& input1_int8, const int64_t input1_zero_point, @@ -28,17 +27,51 @@ Tensor& quantized_mul_out( const int64_t output_zero_point, const int64_t output_multiplier, const int64_t output_shift, + ActivationLayout layout, + const char* op_name, Tensor& out) { - // Validate tensor types and quantization parameters - - bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8); + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + bool channel_broadcast = + is_channel_broadcast(input1_int8, input2_int8, channel_dim); validate_cmsis_nn_tensor_requirements( input1_int8, input2_int8, out, ScalarType::Char, - /*require_channels_last=*/channel_broadcast, + /*require_channels_last=*/ + channel_broadcast && layout == ActivationLayout::NCHWLogical, /*require_same_sizes=*/!channel_broadcast); + if (layout == ActivationLayout::NHWCLogical) { + ET_CHECK_MSG( + input1_int8.dim() == 4 && input2_int8.dim() == 4 && out.dim() == 4, + "%s: tensors must be 4-D", + op_name); + ET_CHECK_MSG( + executorch::runtime::is_contiguous_dim_order( + input1_int8.dim_order().data(), input1_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + input2_int8.dim_order().data(), + input2_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + out.dim_order().data(), out.dim_order().size()), + "%s: tensors must use contiguous dimension order", + op_name); + } else if (channel_broadcast) { + ET_CHECK_MSG( + is_channels_last_tensor(input1_int8) && + is_channels_last_tensor(input2_int8) && + is_channels_last_tensor(out), + "%s: channel-broadcast tensors must use channels-last dimension order", + op_name); + } + if (channel_broadcast) { + const Tensor& full_input = + input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8; + ET_CHECK_MSG( + out.sizes() == full_input.sizes(), + "%s: output must have the broadcast result shape", + op_name); + } const int32_t kIdentityMultiplier(/*value=*/1); const int32_t kZeroShift(/*value=*/0); @@ -70,7 +103,7 @@ Tensor& quantized_mul_out( std::swap(input1_ptr, input2_ptr); } - muls_per_loop = input1_int8.size(1); + muls_per_loop = input1_int8.size(channel_dim); } else { muls_per_loop = out.numel(); } @@ -108,7 +141,8 @@ Tensor& quantized_mul_out( if (status != ARM_CMSIS_NN_SUCCESS) { ET_LOG( Error, - "quantized_mul_out: arm_elementwise_mul_s8 failed with status [%d]", + "%s: arm_elementwise_mul_s8 failed with status [%d]", + op_name, status); context.fail(Error::Internal); return out; @@ -117,5 +151,55 @@ Tensor& quantized_mul_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_mul_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + Tensor& out) { + return quantized_mul_out_impl( + context, + input1_int8, + input1_zero_point, + input2_int8, + input2_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ActivationLayout::NCHWLogical, + "quantized_mul_out", + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_mul_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + Tensor& out) { + return quantized_mul_out_impl( + context, + input1_int8, + input1_zero_point, + input2_int8, + input2_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ActivationLayout::NHWCLogical, + "quantized_mul_nhwc_out", + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 79c7dddcc30..bcc78fafddd 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -137,6 +137,23 @@ def dequantize_per_tensor_impl( "*, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_add_nhwc(" + "Tensor self, int self_zero_point, int self_multiplier, int self_shift, " + "Tensor other, int other_zero_point, int other_multiplier, int other_shift, " + "int output_zero_point, int output_multiplier, int output_shift, " + "int activation_min, int activation_max) -> Tensor" +) + +lib.define( + "quantized_add_nhwc.out(" + "Tensor self, int self_zero_point, int self_multiplier, int self_shift, " + "Tensor other, int other_zero_point, int other_multiplier, int other_shift, " + "int output_zero_point, int output_multiplier, int output_shift, " + "int activation_min, int activation_max, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + @register_fake("cortex_m::quantized_add") # type: ignore[misc] def quantized_add_meta( @@ -199,6 +216,78 @@ def quantized_add_impl( return result +@register_fake("cortex_m::quantized_add_nhwc") # type: ignore[misc] +def quantized_add_nhwc_meta( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_add_nhwc expects 4D inputs") + result = quantized_add_meta( + self.permute(0, 3, 1, 2), + self_zero_point, + self_multiplier, + self_shift, + other.permute(0, 3, 1, 2), + other_zero_point, + other_multiplier, + other_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ) + return result.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_add_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_add_nhwc_impl( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_add_nhwc expects 4D inputs") + result = quantized_add_impl( + self.permute(0, 3, 1, 2), + self_zero_point, + self_multiplier, + self_shift, + other.permute(0, 3, 1, 2), + other_zero_point, + other_multiplier, + other_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ) + return result.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED MUL OPERATION DEFINITION # =================================================================== @@ -216,6 +305,20 @@ def quantized_add_impl( "*, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_mul_nhwc(" + "Tensor self, int self_zero_point, " + "Tensor other, int other_zero_point, " + "int output_zero_point, int output_multiplier, int output_shift) -> Tensor" +) +lib.define( + "quantized_mul_nhwc.out(" + "Tensor self, int self_zero_point, " + "Tensor other, int other_zero_point, " + "int output_zero_point, int output_multiplier, int output_shift, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + @register_fake("cortex_m::quantized_mul") # type: ignore[misc] def quantized_mul_meta( @@ -264,6 +367,54 @@ def quantized_mul_impl( return result +@register_fake("cortex_m::quantized_mul_nhwc") # type: ignore[misc] +def quantized_mul_nhwc_meta( + self: torch.Tensor, + self_zero_point: int, + other: torch.Tensor, + other_zero_point: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_mul_nhwc expects 4D inputs") + result = quantized_mul_meta( + self.permute(0, 3, 1, 2), + self_zero_point, + other.permute(0, 3, 1, 2), + other_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ) + return result.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_mul_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_mul_nhwc_impl( + self: torch.Tensor, + self_zero_point: int, + other: torch.Tensor, + other_zero_point: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_mul_nhwc expects 4D inputs") + result = quantized_mul_impl( + self.permute(0, 3, 1, 2), + self_zero_point, + other.permute(0, 3, 1, 2), + other_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ) + return result.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED DIV OPERATION DEFINITION # =================================================================== diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index e91aaca3569..43b75c784c4 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -23,12 +23,24 @@ - arg_meta: null kernel_name: cortex_m::quantized_add_out +- func: cortex_m::quantized_add_nhwc.out(Tensor self, int self_zero_point, int self_multiplier, int self_shift, Tensor other, int other_zero_point, int other_multiplier, int other_shift, int output_zero_point, int output_multiplier, int output_shift, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_add_nhwc_out + - func: cortex_m::quantized_mul.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, int output_multiplier, int output_shift, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_mul_out +- func: cortex_m::quantized_mul_nhwc.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, int output_multiplier, int output_shift, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_mul_nhwc_out + - func: cortex_m::quantized_div.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, float output_scale, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index ad91eef264a..f1e301f1bcd 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -56,8 +56,10 @@ ops_list=( cortex_m::quantize_per_tensor.out cortex_m::dequantize_per_tensor.out cortex_m::quantized_add.out + cortex_m::quantized_add_nhwc.out cortex_m::quantized_div.out cortex_m::quantized_mul.out + cortex_m::quantized_mul_nhwc.out cortex_m::quantized_activation.out cortex_m::minimum.out cortex_m::maximum.out diff --git a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py index 36b92bbd9b5..aa2ad366829 100644 --- a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py +++ b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py @@ -197,6 +197,46 @@ def forward(self, x): ) +class AddNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", _int8_values((1, 1, 1, 3))) + + def forward(self, x): + return torch.ops.cortex_m.quantized_add_nhwc.default( + x, + 0, + 1 << 30, + -1, + self.bias, + 0, + 1 << 30, + -1, + 0, + 1 << 30, + -1, + -128, + 127, + ) + + +class MulNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", _int8_values((1, 1, 1, 3))) + + def forward(self, x): + return torch.ops.cortex_m.quantized_mul_nhwc.default( + x, + 0, + self.bias, + 0, + 0, + 1 << 30, + -1, + ) + + def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): _run_on_fvp( Conv2dNhwc(), @@ -254,3 +294,23 @@ def test_pad_nhwc_runs_on_fvp_with_singleton_height(cortex_m_target): exir_ops.edge.cortex_m.pad_nhwc.default, cortex_m_target, ) + + +def test_channel_broadcast_add_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + AddNhwc(), + _int8_values((1, 5, 7, 3)), + exir_ops.edge.cortex_m.quantized_add_nhwc.default, + cortex_m_target, + atol=1, + ) + + +def test_channel_broadcast_mul_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + MulNhwc(), + _int8_values((1, 5, 7, 3)), + exir_ops.edge.cortex_m.quantized_mul_nhwc.default, + cortex_m_target, + atol=1, + ) diff --git a/backends/cortex_m/test/test_quantized_conv2d_layout.py b/backends/cortex_m/test/test_quantized_conv2d_layout.py index 891ea750686..7a99a372e92 100644 --- a/backends/cortex_m/test/test_quantized_conv2d_layout.py +++ b/backends/cortex_m/test/test_quantized_conv2d_layout.py @@ -104,6 +104,28 @@ def _run_max_pool2d(op, x): ) +def _run_add(op, x, bias): + return op( + x, + 0, + 1 << 30, + -1, + bias, + 0, + 1 << 30, + -1, + 0, + 1 << 30, + -1, + -128, + 127, + ) + + +def _run_mul(op, x, bias): + return op(x, 0, bias, 0, 0, 1 << 30, -1) + + def test_nhwc_conv2d_matches_legacy_layout(): torch.manual_seed(0) x = torch.randint(-8, 8, (1, 3, 8, 8), dtype=torch.int8) @@ -200,6 +222,42 @@ def test_nhwc_max_pool2d_matches_legacy_layout(): torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) +def test_nhwc_channel_broadcast_add_matches_legacy_layout(): + x = torch.randint(-8, 8, (1, 4, 5, 7), dtype=torch.int8) + bias = torch.randint(-4, 4, (1, 4, 1, 1), dtype=torch.int8) + + legacy = _run_add( + torch.ops.cortex_m.quantized_add, + x.to(memory_format=torch.channels_last), + bias.to(memory_format=torch.channels_last), + ) + explicit = _run_add( + torch.ops.cortex_m.quantized_add_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + bias.permute(0, 2, 3, 1).contiguous(), + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + +def test_nhwc_channel_broadcast_mul_matches_legacy_layout(): + x = torch.randint(-8, 8, (1, 4, 5, 7), dtype=torch.int8) + bias = torch.randint(-4, 4, (1, 4, 1, 1), dtype=torch.int8) + + legacy = _run_mul( + torch.ops.cortex_m.quantized_mul, + x.to(memory_format=torch.channels_last), + bias.to(memory_format=torch.channels_last), + ) + explicit = _run_mul( + torch.ops.cortex_m.quantized_mul_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + bias.permute(0, 2, 3, 1).contiguous(), + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + def test_nhwc_conv2d_fake_shape_is_logical_nhwc(): with FakeTensorMode(): output = _run_conv2d( From 98aa5eeaac004e1d30aeaf79cd4f5e315f4be72b Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 11:55:12 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- backends/cortex_m/ops/op_pad.cpp | 80 ++++++++++++++++--- backends/cortex_m/ops/operators.py | 64 +++++++++++++++ backends/cortex_m/ops/operators.yaml | 6 ++ backends/cortex_m/test/build_test_runner.sh | 1 + .../test/ops/test_explicit_nhwc_runtime.py | 18 +++++ backends/cortex_m/test/ops/test_pad.py | 13 +++ 6 files changed, 173 insertions(+), 9 deletions(-) diff --git a/backends/cortex_m/ops/op_pad.cpp b/backends/cortex_m/ops/op_pad.cpp index 57b5257873e..2b2268e8480 100644 --- a/backends/cortex_m/ops/op_pad.cpp +++ b/backends/cortex_m/ops/op_pad.cpp @@ -17,21 +17,19 @@ namespace { constexpr size_t kMaxSupportedDims = 4; -} // namespace - -// cppcheck-suppress unusedFunction -Tensor& pad_out( +Tensor& pad_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef pre_pad, const Int64ArrayRef post_pad, int64_t pad_value, + bool require_contiguous, Tensor& out) { if (input.scalar_type() != ScalarType::Char || out.scalar_type() != ScalarType::Char) { ET_LOG( Error, - "pad_out: only int8 tensors are supported (input=%d, out=%d)", + "cortex_m::pad: only int8 tensors are supported (input=%d, out=%d)", static_cast(input.scalar_type()), static_cast(out.scalar_type())); context.fail(Error::InvalidArgument); @@ -42,22 +40,48 @@ Tensor& pad_out( if (rank == 0 || rank > kMaxSupportedDims) { ET_LOG( Error, - "pad_out: expected tensor rank in [1, %zu], got %zu", + "cortex_m::pad: expected tensor rank in [1, %zu], got %zu", kMaxSupportedDims, rank); context.fail(Error::InvalidArgument); return out; } + if (pre_pad.size() != kMaxSupportedDims || + post_pad.size() != kMaxSupportedDims) { + ET_LOG(Error, "cortex_m::pad: pre_pad and post_pad must have length 4"); + context.fail(Error::InvalidArgument); + return out; + } + + if (require_contiguous) { + // This entry point infers nothing: it requires the dim order to say the + // tensor is contiguous, and then indexes the padding by logical axis. + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + out.dim_order().data(), out.dim_order().size())) { + ET_LOG( + Error, + "cortex_m::pad_contiguous: input and output must use contiguous dim order"); + context.fail(Error::InvalidArgument); + return out; + } + } // Permute logical sizes to physical memory order. // Padding is already in physical order from the AOT pass. constexpr size_t kNhwcDimOrder[] = {0, 2, 3, 1}; const size_t offset = kMaxSupportedDims - rank; - const bool nhwc = is_channels_last_tensor(input); + // Only the legacy entry point infers the layout. Its predicate is tolerant on + // purpose: the tolerance short-circuits before the dim order is consulted, + // which is what keeps it agreeing with the AOT pass for shapes whose + // serialized dim order cannot name the channel axis. + const bool legacy_channels_last = + !require_contiguous && is_channels_last_tensor(input); int32_t dims[kMaxSupportedDims] = {1, 1, 1, 1}; for (size_t i = 0; i < rank; ++i) { - const size_t src = nhwc ? kNhwcDimOrder[offset + i] : i; + const size_t src = legacy_channels_last ? kNhwcDimOrder[offset + i] : i; dims[offset + i] = static_cast(input.size(src)); } @@ -87,7 +111,7 @@ Tensor& pad_out( if (status != ARM_CMSIS_NN_SUCCESS) { ET_LOG( Error, - "pad_out: arm_pad_s8 failed with status [%d]", + "cortex_m::pad: arm_pad_s8 failed with status [%d]", static_cast(status)); context.fail(Error::Internal); return out; @@ -96,5 +120,43 @@ Tensor& pad_out( return out; } +} // namespace + +// cppcheck-suppress unusedFunction +Tensor& pad_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef pre_pad, + const Int64ArrayRef post_pad, + int64_t pad_value, + Tensor& out) { + return pad_out_impl( + context, + input, + pre_pad, + post_pad, + pad_value, + /*require_contiguous=*/false, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& pad_contiguous_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef pre_pad, + const Int64ArrayRef post_pad, + int64_t pad_value, + Tensor& out) { + return pad_out_impl( + context, + input, + pre_pad, + post_pad, + pad_value, + /*require_contiguous=*/true, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 6c6ab804b9a..96a7bdc0165 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -663,6 +663,18 @@ def transpose_impl(input: torch.Tensor, perm: Sequence[int]) -> torch.Tensor: "pad.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, " "*, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "pad_contiguous(Tensor input, int[] pre_pad, int[] post_pad, int pad_value) -> Tensor" +) +lib.define( + "pad_contiguous.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +_NHWC_INV_ORDER = [0, 3, 1, 2] + + def _pad_to_logical_order(physical_pad: list[int], input: torch.Tensor) -> list[int]: """Inverse of _to_physical_order: map physical-order padding back to logical.""" if not is_channels_last(input): @@ -717,6 +729,58 @@ def pad_impl( return F.pad(input, padding, mode="constant", value=pad_value) +@register_fake("cortex_m::pad_contiguous") # type: ignore[misc] +def pad_contiguous_meta( + input: torch.Tensor, + pre_pad: list[int], + post_pad: list[int], + pad_value: int, +) -> torch.Tensor: + del pad_value + rank = input.dim() + if rank == 0 or rank > 4: + raise RuntimeError( + f"cortex_m.pad_contiguous expects a rank in [1, 4], got {rank}" + ) + if len(pre_pad) != 4 or len(post_pad) != 4: + raise RuntimeError( + "cortex_m.pad_contiguous expects four padding values per side" + ) + offset = 4 - rank + output_shape = [ + input.shape[dim] + pre_pad[offset + dim] + post_pad[offset + dim] + for dim in range(rank) + ] + return torch.empty(output_shape, dtype=input.dtype, device=input.device) + + +@impl(lib, "pad_contiguous", "CompositeExplicitAutograd") # type: ignore[misc] +def pad_contiguous_impl( + input: torch.Tensor, + pre_pad: list[int], + post_pad: list[int], + pad_value: int, +) -> torch.Tensor: + rank = input.dim() + if rank == 0 or rank > 4: + raise RuntimeError( + f"cortex_m.pad_contiguous expects a rank in [1, 4], got {rank}" + ) + if len(pre_pad) != 4 or len(post_pad) != 4: + raise RuntimeError( + "cortex_m.pad_contiguous expects four padding values per side" + ) + offset = 4 - rank + padding = [] + for dim in reversed(range(rank)): + padding.extend([pre_pad[offset + dim], post_pad[offset + dim]]) + return F.pad(input, padding, mode="constant", value=pad_value) + + +# =================================================================== +# QUANTIZED CONV2D OPERATION DEFINITION +# =================================================================== + lib.define( "quantized_conv2d(" "Tensor input, " diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index 15d7f97b929..93fdd83835b 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -77,6 +77,12 @@ - arg_meta: null kernel_name: cortex_m::pad_out +- func: cortex_m::pad_contiguous.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::pad_contiguous_out + - func: cortex_m::quantized_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index c597b222ca5..dddef3c9ed4 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -65,6 +65,7 @@ ops_list=( cortex_m::softmax.out cortex_m::transpose.out cortex_m::pad.out + cortex_m::pad_contiguous.out cortex_m::quantized_conv2d.out cortex_m::quantized_conv2d_nhwc.out cortex_m::quantized_depthwise_conv2d.out diff --git a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py index 6edc6000cef..a64e583be49 100644 --- a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py +++ b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py @@ -217,6 +217,16 @@ def forward(self, x): ) +class PadNhwc(torch.nn.Module): + def forward(self, x): + return torch.ops.cortex_m.pad_contiguous.default( + x, + [0, 1, 2, 0], + [0, 2, 1, 0], + -7, + ) + + def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): _run_on_fvp( Conv2dNhwc(), @@ -276,3 +286,11 @@ def test_max_pool2d_nhwc_runs_on_fvp(cortex_m_target): cortex_m_target, ) + +def test_pad_contiguous_runs_on_fvp_with_singleton_height(cortex_m_target): + _run_on_fvp( + PadNhwc(), + _int8_values((1, 1, 7, 3)), + exir_ops.edge.cortex_m.pad_contiguous.default, + cortex_m_target, + ) diff --git a/backends/cortex_m/test/ops/test_pad.py b/backends/cortex_m/test/ops/test_pad.py index f1bf5f4a568..0fd182f1bed 100644 --- a/backends/cortex_m/test/ops/test_pad.py +++ b/backends/cortex_m/test/ops/test_pad.py @@ -77,6 +77,19 @@ def forward(self, x): CortexMPad((1, 2, 3, 4)), (ramp_tensor(-1.0, 1.0, (1, 3, 4, 5)).to(memory_format=torch.channels_last),), ), + # A channels-last tensor with one channel serializes its dim order as + # (0, 2, 1, 3), which names no channel axis. Deriving the physical sizes + # from it instead of from the shape sizes the pad write wrongly. + "pad_rank4_single_channel_channels_last": McuTestCase( + CortexMPad((1, 1, 2, 2)), + (ramp_tensor(-0.5, 0.5, (1, 1, 3, 4)).to(memory_format=torch.channels_last),), + ), + # With one channel and unit width the dim order collapses all the way to + # (0, 1, 2, 3), making the tensor indistinguishable from a contiguous one. + "pad_rank4_single_channel_unit_width_channels_last": McuTestCase( + CortexMPad((0, 0, 2, 2)), + (ramp_tensor(-0.5, 0.5, (1, 1, 8, 1)).to(memory_format=torch.channels_last),), + ), }