From 6923cfaf648443c83cedb52b2f9ad870326a1e30 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 12:13:01 -0700 Subject: [PATCH] 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),), + ), }