From 335c69e51e1d3ecc836bc0bf90cd489d15ab4625 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:50:57 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cortex_m/ops/op_pad.cpp | 82 +++- .../cortex_m/ops/op_quantized_avg_pool2d.cpp | 28 ++ backends/cortex_m/ops/op_quantized_conv2d.cpp | 36 ++ .../ops/op_quantized_depthwise_conv2d.cpp | 38 ++ .../cortex_m/ops/op_quantized_max_pool2d.cpp | 30 ++ .../ops/op_quantized_transpose_conv2d.cpp | 40 ++ backends/cortex_m/ops/operators.py | 445 ++++++++++++++++++ backends/cortex_m/ops/operators.yaml | 36 ++ backends/cortex_m/passes/BUCK | 15 +- .../cortex_m/passes/scratch_buffer_sizes.py | 48 +- backends/cortex_m/test/build_test_runner.sh | 6 + .../test/ops/test_explicit_nhwc_runtime.py | 256 ++++++++++ backends/cortex_m/test/targets.bzl | 18 +- .../test/test_quantized_conv2d_layout.py | 381 +++++++++++++++ 14 files changed, 1437 insertions(+), 22 deletions(-) create mode 100644 backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py create mode 100644 backends/cortex_m/test/test_quantized_conv2d_layout.py diff --git a/backends/cortex_m/ops/op_pad.cpp b/backends/cortex_m/ops/op_pad.cpp index 57b5257873e..76782b430a8 100644 --- a/backends/cortex_m/ops/op_pad.cpp +++ b/backends/cortex_m/ops/op_pad.cpp @@ -17,21 +17,21 @@ 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, + ActivationLayout layout, + const char* op_name, 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)", + "%s: only int8 tensors are supported (input=%d, out=%d)", + op_name, static_cast(input.scalar_type()), static_cast(out.scalar_type())); context.fail(Error::InvalidArgument); @@ -42,22 +42,45 @@ Tensor& pad_out( if (rank == 0 || rank > kMaxSupportedDims) { ET_LOG( Error, - "pad_out: expected tensor rank in [1, %zu], got %zu", + "%s: expected tensor rank in [1, %zu], got %zu", + op_name, kMaxSupportedDims, rank); context.fail(Error::InvalidArgument); return out; } + if (pre_pad.size() != kMaxSupportedDims || + post_pad.size() != kMaxSupportedDims) { + ET_LOG(Error, "%s: pre_pad and post_pad must have length 4", op_name); + context.fail(Error::InvalidArgument); + return out; + } + + if (layout == ActivationLayout::NHWCLogical) { + if (rank != kMaxSupportedDims || + !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, + "%s: input and output must be contiguous 4-D tensors", + op_name); + 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); + const bool legacy_channels_last = + layout == ActivationLayout::NCHWLogical && 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 +110,8 @@ Tensor& pad_out( if (status != ARM_CMSIS_NN_SUCCESS) { ET_LOG( Error, - "pad_out: arm_pad_s8 failed with status [%d]", + "%s: arm_pad_s8 failed with status [%d]", + op_name, static_cast(status)); context.fail(Error::Internal); return out; @@ -96,5 +120,45 @@ 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, + ActivationLayout::NCHWLogical, + "pad_out", + out); +} + +// cppcheck-suppress unusedFunction +Tensor& pad_nhwc_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, + ActivationLayout::NHWCLogical, + "pad_nhwc_out", + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp index 66940f18997..0a3080cb607 100644 --- a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp @@ -184,5 +184,33 @@ Tensor& quantized_avg_pool2d_out( out); } +// cppcheck-suppress unusedFunction +Tensor& quantized_avg_pool2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const bool ceil_mode, + const int64_t zero_point, + const int64_t multiplier, + const int64_t shift, + const Tensor& scratch, + Tensor& out) { + return quantized_avg_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_conv2d.cpp b/backends/cortex_m/ops/op_quantized_conv2d.cpp index 7865b50e486..979e4b7c94f 100644 --- a/backends/cortex_m/ops/op_quantized_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_conv2d.cpp @@ -280,5 +280,41 @@ Tensor& quantized_conv2d_out( out); } +// cppcheck-suppress unusedFunction +Tensor& quantized_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp index 4aa58bb33dd..8f4da4f452d 100644 --- a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp @@ -331,5 +331,43 @@ Tensor& quantized_depthwise_conv2d_out( out); } +// cppcheck-suppress unusedFunction +Tensor& quantized_depthwise_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t depth_multiplier, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_depthwise_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp index 68caa764ad5..1118a0640b4 100644 --- a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp @@ -129,5 +129,35 @@ Tensor& quantized_max_pool2d_out( out); } +// cppcheck-suppress unusedFunction +Tensor& quantized_max_pool2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const bool ceil_mode, + const int64_t input_zero_point, + const int64_t output_zero_point, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_max_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp index fcfe78ce48d..4ac9b2338e6 100644 --- a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp @@ -302,5 +302,45 @@ Tensor& quantized_transpose_conv2d_out( out); } +// cppcheck-suppress unusedFunction +Tensor& quantized_transpose_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef output_padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + const Tensor& output_scratch, + Tensor& out) { + return quantized_transpose_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 44e47087c11..79c7dddcc30 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -669,6 +669,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_nhwc(Tensor input, int[] pre_pad, int[] post_pad, int pad_value) -> Tensor" +) +lib.define( + "pad_nhwc.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, " + "*, Tensor(a!) out) -> Tensor(a!)" +) _NHWC_INV_ORDER = [0, 3, 1, 2] @@ -720,6 +727,39 @@ def pad_impl( return F.pad(input, padding, mode="constant", value=pad_value) +@register_fake("cortex_m::pad_nhwc") # type: ignore[misc] +def pad_nhwc_meta( + input: torch.Tensor, + pre_pad: list[int], + post_pad: list[int], + pad_value: int, +) -> torch.Tensor: + del pad_value + if input.dim() != 4: + raise RuntimeError("cortex_m.pad_nhwc expects a 4D input tensor") + if len(pre_pad) != 4 or len(post_pad) != 4: + raise RuntimeError("cortex_m.pad_nhwc expects four padding values per side") + output_shape = [input.shape[dim] + pre_pad[dim] + post_pad[dim] for dim in range(4)] + return torch.empty(output_shape, dtype=input.dtype, device=input.device) + + +@impl(lib, "pad_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def pad_nhwc_impl( + input: torch.Tensor, + pre_pad: list[int], + post_pad: list[int], + pad_value: int, +) -> torch.Tensor: + if input.dim() != 4: + raise RuntimeError("cortex_m.pad_nhwc expects a 4D input tensor") + if len(pre_pad) != 4 or len(post_pad) != 4: + raise RuntimeError("cortex_m.pad_nhwc expects four padding values per side") + padding = [] + for dim in reversed(range(4)): + padding.extend([pre_pad[dim], post_pad[dim]]) + return F.pad(input, padding, mode="constant", value=pad_value) + + # =================================================================== # QUANTIZED CONV2D OPERATION DEFINITION # =================================================================== @@ -915,6 +955,91 @@ def quantized_conv2d_impl( return result.to(torch.int8, memory_format=torch.channels_last) +lib.define( + "quantized_conv2d_nhwc(" + "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" +) +lib.define( + "quantized_conv2d_nhwc.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!)" +) + + +@register_fake("cortex_m::quantized_conv2d_nhwc") # type: ignore[misc] +def quantized_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED DEPTHWISE CONV2D OPERATION DEFINITION # =================================================================== @@ -1062,6 +1187,95 @@ def quantized_depthwise_conv2d_impl( return result.to(torch.int8, memory_format=torch.channels_last) +lib.define( + "quantized_depthwise_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int depth_multiplier, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch) -> Tensor" +) +lib.define( + "quantized_depthwise_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int depth_multiplier, 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!)" +) + + +@register_fake("cortex_m::quantized_depthwise_conv2d_nhwc") # type: ignore[misc] +def quantized_depthwise_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + depth_multiplier: int, + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_depthwise_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_depthwise_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_depthwise_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + depth_multiplier: int, + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_depthwise_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED TRANSPOSE_CONV2D OPERATION DEFINITION # =================================================================== @@ -1270,6 +1484,100 @@ def quantized_transpose_conv2d_impl( return result.to(torch.int8).to(memory_format=torch.channels_last) +lib.define( + "quantized_transpose_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] output_padding, int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "Tensor output_scratch) -> Tensor" +) +lib.define( + "quantized_transpose_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] output_padding, int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_transpose_conv2d_nhwc") # type: ignore[misc] +def quantized_transpose_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + output_padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, + output_scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_transpose_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_transpose_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_transpose_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + output_padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, + output_scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_transpose_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED AVG_POOL2D OPERATION DEFINITION # =================================================================== @@ -1365,6 +1673,72 @@ def quantized_avg_pool2d_impl( return output.to(torch.int8) +lib.define( + "quantized_avg_pool2d_nhwc(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "bool ceil_mode, int zero_point, int multiplier, int shift, " + "Tensor scratch) -> Tensor" +) +lib.define( + "quantized_avg_pool2d_nhwc.out(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "bool ceil_mode, int zero_point, int multiplier, int shift, " + "Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_avg_pool2d_nhwc") # type: ignore[misc] +def quantized_avg_pool2d_nhwc_meta( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + ceil_mode: bool, + zero_point: int, + multiplier: int, + shift: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_avg_pool2d_meta( + input.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_avg_pool2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_avg_pool2d_nhwc_impl( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + ceil_mode: bool, + zero_point: int, + multiplier: int, + shift: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_avg_pool2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED MAX POOL2D OPERATION DEFINITION # =================================================================== @@ -1520,3 +1894,74 @@ def quantized_max_pool2d_impl( ) result = torch.clamp(result, activation_min, activation_max) return result.to(torch.int8).contiguous(memory_format=torch.channels_last) + + +lib.define( + "quantized_max_pool2d_nhwc(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "int[] dilation, bool ceil_mode, int input_zero_point, " + "int output_zero_point, int activation_min, int activation_max) -> Tensor" +) +lib.define( + "quantized_max_pool2d_nhwc.out(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "int[] dilation, bool ceil_mode, int input_zero_point, " + "int output_zero_point, int activation_min, int activation_max, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_max_pool2d_nhwc") # type: ignore[misc] +def quantized_max_pool2d_nhwc_meta( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + input_zero_point: int, + output_zero_point: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + nchw = quantized_max_pool2d_meta( + input.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_max_pool2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_max_pool2d_nhwc_impl( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + input_zero_point: int, + output_zero_point: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + nchw = quantized_max_pool2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ) + return nchw.permute(0, 2, 3, 1).contiguous() diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index 2c85325f854..e91aaca3569 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -77,12 +77,23 @@ - arg_meta: null kernel_name: cortex_m::pad_out +- func: cortex_m::pad_nhwc.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_nhwc_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: - arg_meta: null kernel_name: cortex_m::quantized_conv2d_out +- func: cortex_m::quantized_conv2d_nhwc.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: + - arg_meta: null + kernel_name: cortex_m::quantized_conv2d_nhwc_out - func: cortex_m::quantized_depthwise_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, 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 @@ -90,23 +101,48 @@ - arg_meta: null kernel_name: cortex_m::quantized_depthwise_conv2d_out +- func: cortex_m::quantized_depthwise_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, 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: + - arg_meta: null + kernel_name: cortex_m::quantized_depthwise_conv2d_nhwc_out + - func: cortex_m::quantized_transpose_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_transpose_conv2d_out +- func: cortex_m::quantized_transpose_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_transpose_conv2d_nhwc_out + - func: cortex_m::quantized_avg_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_avg_pool2d_out + +- func: cortex_m::quantized_avg_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_avg_pool2d_nhwc_out + - func: cortex_m::quantized_max_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, int input_zero_point, int output_zero_point, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_max_pool2d_out +- func: cortex_m::quantized_max_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, int input_zero_point, int output_zero_point, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_max_pool2d_nhwc_out + - func: cortex_m::quantized_batch_matmul.out(Tensor lhs, int lhs_zero_point, Tensor rhs_transposed, int rhs_zero_point, int output_zero_point, int output_multiplier, int output_shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index d301e14823c..09dfb5942f4 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -37,7 +37,6 @@ fbcode_target(_kind = runtime.python_library, "decompose_mean_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", - "scratch_buffer_sizes.py", ], deps=[ "//caffe2:torch", @@ -47,6 +46,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/cortex_m/ops:ops", "//executorch/backends/cortex_m/passes:passes_utils", "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", + "//executorch/backends/cortex_m/passes:scratch_buffer_sizes", "//executorch/backends/transforms:aten_to_dialect_pass", "//executorch/backends/transforms:remove_getitem_op", "//executorch/backends/transforms:replace_scalar_with_tensor", @@ -58,6 +58,19 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name="scratch_buffer_sizes", + srcs=[ + "scratch_buffer_sizes.py", + ], + deps=[ + "//caffe2:torch", + "//executorch/backends/cortex_m:cmsis_nn", + "//executorch/backends/cortex_m/ops:ops", + "//executorch/exir/dialects:lib", + ], +) + fbcode_target(_kind = runtime.python_library, name="passes_utils", srcs=[ diff --git a/backends/cortex_m/passes/scratch_buffer_sizes.py b/backends/cortex_m/passes/scratch_buffer_sizes.py index b247e2be944..65a3a178757 100644 --- a/backends/cortex_m/passes/scratch_buffer_sizes.py +++ b/backends/cortex_m/passes/scratch_buffer_sizes.py @@ -4,6 +4,7 @@ # LICENSE file in the root directory of this source tree. from collections.abc import Callable +from functools import partial from typing import Any, cast import executorch.backends.cortex_m.ops.operators # noqa @@ -37,6 +38,7 @@ def _shape_from_node(node: torch.fx.Node) -> torch.Size: def _get_common_conv_buffer_size_inputs( conv_node: torch.fx.Node, *, + nhwc_logical: bool = False, stride_arg_idx: int = 3, padding_arg_idx: int = 4, dilation_arg_idx: int = 5, @@ -54,13 +56,14 @@ def _get_common_conv_buffer_size_inputs( padding = cast(list[int], conv_node.args[padding_arg_idx]) dilation = cast(list[int], conv_node.args[dilation_arg_idx]) - # Input is NCHW (PyTorch); CMSIS-NN wants NHWC dims. - n, c_in, height, width = _shape_from_node(x) - weight_shape = _shape_from_node(weight) - # Output is NCHW; convert to NHWC dims. - out_n, out_c, out_h, out_w = _shape_from_node(conv_node) + if nhwc_logical: + n, height, width, c_in = _shape_from_node(x) + out_n, out_h, out_w, out_c = _shape_from_node(conv_node) + else: + n, c_in, height, width = _shape_from_node(x) + out_n, out_c, out_h, out_w = _shape_from_node(conv_node) input_nhwc = [n, height, width, c_in] output_nhwc = [out_n, out_h, out_w, out_c] @@ -81,6 +84,7 @@ def _get_common_conv_buffer_size_inputs( def cmsis_nn_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -89,7 +93,9 @@ def cmsis_nn_conv_buffer_size( stride_hw, padding_hw, dilation_hw, - ) = _get_common_conv_buffer_size_inputs(conv_node=conv_node) + ) = _get_common_conv_buffer_size_inputs( + conv_node=conv_node, nhwc_logical=nhwc_logical + ) input_offset = cast(int, conv_node.args[6]) output_offset = cast(int, conv_node.args[7]) output_qmin = cast(int, conv_node.args[10]) @@ -122,6 +128,7 @@ def cmsis_nn_conv_buffer_size( def cmsis_nn_depthwise_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -130,7 +137,9 @@ def cmsis_nn_depthwise_conv_buffer_size( stride_hw, padding_hw, dilation_hw, - ) = _get_common_conv_buffer_size_inputs(conv_node=conv_node) + ) = _get_common_conv_buffer_size_inputs( + conv_node=conv_node, nhwc_logical=nhwc_logical + ) depth_multiplier = cast(int, conv_node.args[6]) input_offset = cast(int, conv_node.args[7]) output_offset = cast(int, conv_node.args[8]) @@ -185,6 +194,7 @@ def cmsis_nn_batch_matmul_buffer_size( def cmsis_nn_transpose_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -195,6 +205,7 @@ def cmsis_nn_transpose_conv_buffer_size( dilation_hw, ) = _get_common_conv_buffer_size_inputs( conv_node=conv_node, + nhwc_logical=nhwc_logical, stride_arg_idx=3, padding_arg_idx=4, dilation_arg_idx=6, @@ -248,13 +259,16 @@ def cmsis_nn_transpose_conv_buffer_size( def cmsis_nn_avgpool_buffer_size( backend: cmsis_nn.Backend, pool_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> 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) + if nhwc_logical: + _, _, _, c_in = _shape_from_node(x) + _, _, out_w, _ = _shape_from_node(pool_node) + else: + _, c_in, _, _ = _shape_from_node(x) + _, _, _, out_w = _shape_from_node(pool_node) return [ int( @@ -270,10 +284,22 @@ def cmsis_nn_avgpool_buffer_size( _target_to_buffer_sizes_registry: dict[Any, BufferSizeFunction] = { exir_ops.edge.cortex_m.quantized_conv2d.default: cmsis_nn_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: partial( + cmsis_nn_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default: cmsis_nn_depthwise_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default: partial( + cmsis_nn_depthwise_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_batch_matmul.default: cmsis_nn_batch_matmul_buffer_size, exir_ops.edge.cortex_m.quantized_transpose_conv2d.default: cmsis_nn_transpose_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default: partial( + cmsis_nn_transpose_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_avg_pool2d.default: cmsis_nn_avgpool_buffer_size, + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default: partial( + cmsis_nn_avgpool_buffer_size, nhwc_logical=True + ), } diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index 4d8502ec59e..ad91eef264a 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -65,11 +65,17 @@ ops_list=( cortex_m::softmax.out cortex_m::transpose.out cortex_m::pad.out + cortex_m::pad_nhwc.out cortex_m::quantized_conv2d.out + cortex_m::quantized_conv2d_nhwc.out cortex_m::quantized_depthwise_conv2d.out + cortex_m::quantized_depthwise_conv2d_nhwc.out cortex_m::quantized_transpose_conv2d.out + cortex_m::quantized_transpose_conv2d_nhwc.out cortex_m::quantized_avg_pool2d.out + cortex_m::quantized_avg_pool2d_nhwc.out cortex_m::quantized_max_pool2d.out + cortex_m::quantized_max_pool2d_nhwc.out cortex_m::quantized_batch_matmul.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 new file mode 100644 index 00000000000..36b92bbd9b5 --- /dev/null +++ b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py @@ -0,0 +1,256 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import torch + +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.passes.scratch_buffer_sizes import ( + required_cmsis_nn_buffer_sizes, +) +from executorch.backends.cortex_m.target_config import CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.test.harness.stages import RunPasses, StageType +from executorch.exir.dialects._ops import ops as exir_ops + + +def _int8_values(shape): + values = torch.arange(math.prod(shape), dtype=torch.int32) + return (values.remainder(7) - 3).to(torch.int8).reshape(shape) + + +def _run_on_fvp( + module, + x, + target, + target_config: CortexMTargetConfig, + scratch_count=0, + atol=1e-3, +): + sizing_inputs = (x,) + tuple( + torch.empty(0, dtype=torch.uint8) for _ in range(scratch_count) + ) + sizing_tester = CortexMTester(module, sizing_inputs, target_config=target_config) + sizing_tester.export().to_edge() + sizing_program = sizing_tester.get_artifact(StageType.TO_EDGE).exported_program() + [node] = [ + node + for node in sizing_program.graph.nodes + if node.op == "call_function" and node.target == target + ] + + if scratch_count: + scratch_sizes = required_cmsis_nn_buffer_sizes(node, target_config.backend) + assert scratch_sizes is not None + assert len(scratch_sizes) == scratch_count + else: + scratch_sizes = [] + + inputs = (x,) + tuple( + torch.empty(size, dtype=torch.uint8) for size in scratch_sizes + ) + tester = CortexMTester(module, inputs, target_config=target_config) + tester.export().to_edge() + program = tester.get_artifact(StageType.TO_EDGE).exported_program() + assert ( + sum( + node.op == "call_function" and node.target == target + for node in program.graph.nodes + ) + == 1 + ) + # The graph already contains the runtime operator; only advance the harness stage. + tester.run_passes(RunPasses(CortexMPassManager, pass_list=[])) + tester.to_executorch().serialize() + tester.run_method_and_compare_outputs(inputs=inputs, atol=atol) + + +class Conv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", _int8_values((4, 3, 2, 3))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [2, 1], + [1, 0], + [1, 1], + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + ) + + +class DepthwiseConv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", _int8_values((1, 3, 2, 4))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_depthwise_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [2, 1], + [1, 0], + [1, 1], + 1, + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + ) + + +class TransposeConv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", _int8_values((4, 2, 4, 2))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch, output_scratch): + return torch.ops.cortex_m.quantized_transpose_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [1, 1], + [0, 0], + [0, 0], + [1, 1], + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + output_scratch, + ) + + +class AvgPool2dNhwc(torch.nn.Module): + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_avg_pool2d_nhwc.default( + x, + [2, 3], + [2, 1], + [0, 1], + False, + 0, + 1 << 30, + 1, + scratch, + ) + + +class MaxPool2dNhwc(torch.nn.Module): + def forward(self, x): + return torch.ops.cortex_m.quantized_max_pool2d_nhwc.default( + x, + [2, 3], + [2, 1], + [0, 1], + [1, 1], + False, + 0, + 0, + -128, + 127, + ) + + +class PadNhwc(torch.nn.Module): + def forward(self, x): + return torch.ops.cortex_m.pad_nhwc.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(), + _int8_values((1, 7, 10, 3)), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + cortex_m_target, + scratch_count=1, + ) + + +def test_depthwise_conv2d_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + DepthwiseConv2dNhwc(), + _int8_values((1, 7, 10, 4)), + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, + cortex_m_target, + scratch_count=1, + ) + + +def test_transpose_conv2d_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + TransposeConv2dNhwc(), + _int8_values((1, 5, 6, 2)), + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default, + cortex_m_target, + scratch_count=2, + ) + + +def test_avg_pool2d_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + AvgPool2dNhwc(), + _int8_values((1, 7, 9, 3)), + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default, + cortex_m_target, + scratch_count=1, + atol=1, + ) + + +def test_max_pool2d_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + MaxPool2dNhwc(), + _int8_values((1, 7, 9, 3)), + exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default, + cortex_m_target, + ) + + +def test_pad_nhwc_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_nhwc.default, + cortex_m_target, + ) diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index b639aaebed1..b80f8d3f31c 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//tools/build_defs:platform_defs.bzl", "CXX") @@ -50,4 +51,19 @@ def define_common_targets(is_fbcode = False): ], ) - + python_pytest( + name = "test_quantized_conv2d_layout", + srcs = [ + "test_quantized_conv2d_layout.py", + ], + compile = "with-source", + typing = False, + deps = [ + "//caffe2:torch", + "//executorch/backends/cortex_m:target_config", + "//executorch/backends/cortex_m/ops:ops", + "//executorch/backends/cortex_m/passes:scratch_buffer_sizes", + "//executorch/exir:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) diff --git a/backends/cortex_m/test/test_quantized_conv2d_layout.py b/backends/cortex_m/test/test_quantized_conv2d_layout.py new file mode 100644 index 00000000000..891ea750686 --- /dev/null +++ b/backends/cortex_m/test/test_quantized_conv2d_layout.py @@ -0,0 +1,381 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# 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 ( + required_cmsis_nn_buffer_sizes, +) +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode + + +def _run_conv2d(op, x, weight, bias): + out_channels = weight.shape[0] + return op( + x, + weight, + bias, + [1, 1], + [1, 1], + [1, 1], + 5, + -3, + torch.full((out_channels,), 1 << 30, dtype=torch.int32), + torch.full((out_channels,), -2, dtype=torch.int32), + -128, + 127, + torch.zeros(0, dtype=torch.uint8), + ) + + +def _run_depthwise_conv2d(op, x, weight, bias): + out_channels = weight.shape[3] + return op( + x, + weight, + bias, + [1, 1], + [1, 1], + [1, 1], + 1, + 5, + -3, + torch.full((out_channels,), 1 << 30, dtype=torch.int32), + torch.full((out_channels,), -2, dtype=torch.int32), + -128, + 127, + torch.zeros(0, dtype=torch.uint8), + ) + + +def _run_transpose_conv2d(op, x, weight, bias): + out_channels = weight.shape[0] + return op( + x, + weight, + bias, + [2, 2], + [1, 1], + [0, 0], + [1, 1], + 5, + -3, + torch.full((out_channels,), 1 << 30, dtype=torch.int32), + torch.full((out_channels,), -2, dtype=torch.int32), + -128, + 127, + torch.zeros(0, dtype=torch.uint8), + torch.zeros(0, dtype=torch.uint8), + ) + + +def _run_avg_pool2d(op, x): + return op( + x, + [2, 2], + [2, 2], + [0, 0], + False, + 0, + 1 << 30, + 1, + torch.zeros(0, dtype=torch.uint8), + ) + + +def _run_max_pool2d(op, x): + return op( + x, + [2, 2], + [2, 2], + [0, 0], + [1, 1], + False, + 0, + 0, + -128, + 127, + ) + + +def test_nhwc_conv2d_matches_legacy_layout(): + torch.manual_seed(0) + x = torch.randint(-8, 8, (1, 3, 8, 8), dtype=torch.int8) + weight = torch.randint(-4, 4, (4, 3, 3, 3), dtype=torch.int8) + bias = torch.randint(-50, 50, (4,), dtype=torch.int32) + + legacy = _run_conv2d( + torch.ops.cortex_m.quantized_conv2d, + x.to(memory_format=torch.channels_last), + weight, + bias, + ) + explicit = _run_conv2d( + torch.ops.cortex_m.quantized_conv2d_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + weight, + bias, + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + +def test_nhwc_depthwise_conv2d_matches_legacy_layout(): + torch.manual_seed(0) + x = torch.randint(-8, 8, (1, 4, 8, 8), dtype=torch.int8) + weight = torch.randint(-4, 4, (1, 3, 3, 4), dtype=torch.int8) + bias = torch.randint(-50, 50, (4,), dtype=torch.int32) + + legacy = _run_depthwise_conv2d( + torch.ops.cortex_m.quantized_depthwise_conv2d, + x.to(memory_format=torch.channels_last), + weight, + bias, + ) + explicit = _run_depthwise_conv2d( + torch.ops.cortex_m.quantized_depthwise_conv2d_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + weight, + bias, + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + +def test_nhwc_transpose_conv2d_matches_legacy_layout(): + torch.manual_seed(0) + x = torch.randint(-8, 8, (1, 3, 6, 6), dtype=torch.int8) + weight = torch.randint(-4, 4, (4, 3, 3, 3), dtype=torch.int8) + bias = torch.randint(-50, 50, (4,), dtype=torch.int32) + + legacy = _run_transpose_conv2d( + torch.ops.cortex_m.quantized_transpose_conv2d, + x.to(memory_format=torch.channels_last), + weight, + bias, + ) + explicit = _run_transpose_conv2d( + torch.ops.cortex_m.quantized_transpose_conv2d_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + weight, + bias, + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + +def test_nhwc_avg_pool2d_matches_legacy_layout(): + x = torch.randint(-8, 8, (1, 4, 8, 8), dtype=torch.int8) + + legacy = _run_avg_pool2d( + torch.ops.cortex_m.quantized_avg_pool2d, + x.to(memory_format=torch.channels_last), + ) + explicit = _run_avg_pool2d( + torch.ops.cortex_m.quantized_avg_pool2d_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + +def test_nhwc_max_pool2d_matches_legacy_layout(): + x = torch.randint(-8, 8, (1, 4, 8, 8), dtype=torch.int8) + + legacy = _run_max_pool2d( + torch.ops.cortex_m.quantized_max_pool2d, + x.to(memory_format=torch.channels_last), + ) + explicit = _run_max_pool2d( + torch.ops.cortex_m.quantized_max_pool2d_nhwc, + x.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( + torch.ops.cortex_m.quantized_conv2d_nhwc, + torch.empty(2, 10, 6, 3, dtype=torch.int8), + torch.empty(5, 3, 3, 3, dtype=torch.int8), + torch.empty(5, dtype=torch.int32), + ) + + assert output.shape == torch.Size([2, 10, 6, 5]) + assert output.dim_order() == (0, 1, 2, 3) + + +def test_nhwc_pad_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_nhwc(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_nhwc_pad_rejects_non_4d_input(): + with pytest.raises(RuntimeError, match="expects a 4D input tensor"): + torch.ops.cortex_m.pad_nhwc( + torch.zeros((1, 5, 3), dtype=torch.int8), + [0, 0, 0, 0], + [0, 0, 0, 0], + 0, + ) + + with FakeTensorMode(): + with pytest.raises(RuntimeError, match="expects a 4D input tensor"): + torch.ops.cortex_m.pad_nhwc( + torch.zeros((1, 5, 3), dtype=torch.int8), + [0, 0, 0, 0], + [0, 0, 0, 0], + 0, + ) + + +def test_nhwc_pad_rejects_invalid_padding_length(): + with pytest.raises(RuntimeError, match="expects four padding values per side"): + torch.ops.cortex_m.pad_nhwc( + 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_nhwc( + 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) + ) + + def make_node( + target, + input_shape, + output_shape, + weight_shape, + trailing_args, + ): + graph = torch.fx.Graph() + with FakeTensorMode() as mode: + input_node = graph.placeholder("input") + input_node.meta["val"] = mode.from_tensor( + torch.empty(input_shape, dtype=torch.int8) + ) + args = [input_node] + if weight_shape is not None: + weight_node = graph.placeholder("weight") + weight_node.meta["val"] = mode.from_tensor( + torch.empty(weight_shape, dtype=torch.int8) + ) + args.extend((weight_node, None)) + args.extend(trailing_args) + node = graph.call_function(target, args=tuple(args)) + node.meta["val"] = mode.from_tensor( + torch.empty(output_shape, dtype=torch.int8) + ) + return node + + cases = ( + ( + exir_ops.edge.cortex_m.quantized_conv2d.default, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + (1, 3, 10, 6), + (1, 10, 6, 3), + (1, 4, 10, 6), + (1, 10, 6, 4), + (4, 3, 3, 3), + ([1, 1], [1, 1], [1, 1], 5, -3, None, None, -128, 127, None), + ), + ( + exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, + (1, 4, 10, 6), + (1, 10, 6, 4), + (1, 4, 5, 5), + (1, 5, 5, 4), + (1, 3, 2, 4), + ([2, 1], [1, 0], [1, 1], 1, 5, -3, None, None, -128, 127, None), + ), + ( + exir_ops.edge.cortex_m.quantized_transpose_conv2d.default, + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default, + (1, 3, 5, 6), + (1, 5, 6, 3), + (1, 4, 9, 8), + (1, 9, 8, 4), + (4, 2, 3, 3), + ( + [2, 1], + [1, 0], + [0, 0], + [1, 1], + 5, + -3, + None, + None, + -128, + 127, + None, + None, + ), + ), + ( + exir_ops.edge.cortex_m.quantized_avg_pool2d.default, + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default, + (1, 4, 10, 6), + (1, 10, 6, 4), + (1, 4, 5, 3), + (1, 5, 3, 4), + None, + ([2, 2], [2, 2], [0, 0], False, 0, 1 << 30, 1, None), + ), + ) + + for ( + legacy_target, + explicit_target, + legacy_input_shape, + explicit_input_shape, + legacy_output_shape, + explicit_output_shape, + weight_shape, + trailing_args, + ) in cases: + legacy = make_node( + legacy_target, + legacy_input_shape, + legacy_output_shape, + weight_shape, + trailing_args, + ) + explicit = make_node( + explicit_target, + explicit_input_shape, + explicit_output_shape, + weight_shape, + trailing_args, + ) + for backend in backends: + assert required_cmsis_nn_buffer_sizes( + legacy, backend + ) == required_cmsis_nn_buffer_sizes(explicit, backend)