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 2753fab7e1a..be2f4f87607 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -663,6 +663,13 @@ 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] @@ -722,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 0db25ae508f..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(), @@ -275,3 +285,12 @@ def test_max_pool2d_nhwc_runs_on_fvp(cortex_m_target): exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default, 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),), + ), } diff --git a/backends/cortex_m/test/test_quantized_conv2d_layout.py b/backends/cortex_m/test/test_quantized_conv2d_layout.py index c6aa57d7248..6061b827fa0 100644 --- a/backends/cortex_m/test/test_quantized_conv2d_layout.py +++ b/backends/cortex_m/test/test_quantized_conv2d_layout.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import pytest import torch from executorch.backends.cortex_m.passes.scratch_buffer_sizes import ( @@ -212,6 +213,59 @@ def test_nhwc_conv2d_fake_shape_is_logical_nhwc(): assert output.dim_order() == (0, 1, 2, 3) +def test_pad_contiguous_preserves_singleton_height_layout(): + x = torch.arange(1 * 1 * 5 * 3, dtype=torch.int8).reshape(1, 1, 5, 3) + pre_pad = [0, 0, 1, 0] + post_pad = [0, 0, 2, 0] + + actual = torch.ops.cortex_m.pad_contiguous(x, pre_pad, post_pad, -7) + expected = torch.nn.functional.pad(x, (0, 0, 1, 2, 0, 0, 0, 0), value=-7) + + assert actual.shape == torch.Size([1, 1, 8, 3]) + torch.testing.assert_close(actual, expected) + + +def test_pad_contiguous_handles_every_supported_rank(): + """Ranks below four are contiguous by construction, so this entry point + covers them too. That is what lets it eventually replace ``cortex_m::pad`` + outright rather than sitting beside it forever.""" + for shape, pad in ( + ((2, 3, 4, 5), [0, 0, 1, 2]), + ((2, 3, 4), [0, 0, 1, 2]), + ((3, 5), [0, 0, 1, 2]), + ((6,), [0, 0, 0, 2]), + ): + x = torch.randint(-8, 8, shape, dtype=torch.int8) + rank = len(shape) + offset = 4 - rank + actual = torch.ops.cortex_m.pad_contiguous(x, pad, pad, -7) + flat = [] + for dim in reversed(range(rank)): + flat.extend([pad[offset + dim], pad[offset + dim]]) + expected = torch.nn.functional.pad(x, flat, value=-7) + assert actual.shape == expected.shape, shape + torch.testing.assert_close(actual, expected) + + +def test_pad_contiguous_rejects_invalid_padding_length(): + with pytest.raises(RuntimeError, match="expects four padding values per side"): + torch.ops.cortex_m.pad_contiguous( + torch.zeros((1, 1, 5, 3), dtype=torch.int8), + [0, 0, 0], + [0, 0, 0, 0], + 0, + ) + + with FakeTensorMode(): + with pytest.raises(RuntimeError, match="expects four padding values per side"): + torch.ops.cortex_m.pad_contiguous( + torch.zeros((1, 1, 5, 3), dtype=torch.int8), + [0, 0, 0], + [0, 0, 0, 0], + 0, + ) + + def test_nhwc_and_legacy_scratch_sizes_match(): backends = tuple( CortexMTargetConfig(cpu=cpu).backend for cpu in (CortexM.M33, CortexM.M55)