From ec4b3270aae0a15b8830dfc7c379271f111556f7 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:50:52 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- backends/cortex_m/ops/cortex_m_ops_common.h | 48 ++++++--- .../cortex_m/ops/op_quantized_avg_pool2d.cpp | 35 ++++++- backends/cortex_m/ops/op_quantized_conv2d.cpp | 96 +++++++++++++----- .../ops/op_quantized_depthwise_conv2d.cpp | 99 +++++++++++++++---- .../cortex_m/ops/op_quantized_max_pool2d.cpp | 38 ++++++- .../ops/op_quantized_transpose_conv2d.cpp | 94 ++++++++++++++---- .../cortex_m/test/models/test_mobilenet_v3.py | 4 - 7 files changed, 334 insertions(+), 80 deletions(-) diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index 2e3f49dd861..bcaed0a1bc7 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -36,6 +37,11 @@ using KernelRuntimeContext = torch::executor::KernelRuntimeContext; // 16-byte alignment for MVE vector operations. constexpr size_t kCortexMMveAlignment = 16; +enum class ActivationLayout { + NCHWLogical, + NHWCLogical, +}; + // Basic tensor type / layout validation and dimension order checking inline void validate_cmsis_nn_tensor_requirements( const Tensor& input1, @@ -203,7 +209,7 @@ inline bool prepare_cmsis_pool2d_config( int64_t activation_min, int64_t activation_max, CmsisPool2DConfig& config, - bool require_channels_last = true, + ActivationLayout layout, bool allow_ceil_mode = false) { if (input.dim() != 4 || output.dim() != 4) { ET_LOG(Error, "%s: tensors must be 4-D", op_name); @@ -218,7 +224,9 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (input.size(0) != output.size(0) || input.size(1) != output.size(1)) { + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + if (input.size(0) != output.size(0) || + input.size(channel_dim) != output.size(channel_dim)) { ET_LOG( Error, "%s: batch and channel dimensions must match between input and output", @@ -227,13 +235,21 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (require_channels_last) { - if (!is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { - ET_LOG( - Error, "%s: tensors must use channels_last dimension order", op_name); + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG(Error, "%s: tensors must use contiguous dimension order", op_name); context.fail(Error::InvalidArgument); return false; } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { + ET_LOG( + Error, "%s: tensors must use channels_last dimension order", op_name); + context.fail(Error::InvalidArgument); + return false; } auto check_tuple_len = [&](const Int64ArrayRef& arr, @@ -312,19 +328,29 @@ inline bool prepare_cmsis_pool2d_config( return false; } + const int64_t height_dim = layout == ActivationLayout::NHWCLogical ? 1 : 2; + const int64_t width_dim = layout == ActivationLayout::NHWCLogical ? 2 : 3; int32_t batch, channels, input_h, input_w, output_h, output_w; if (!check_int32_within_range( context, op_name, input.size(0), "input batch", batch) || !check_int32_within_range( - context, op_name, input.size(1), "input channels", channels) || + context, + op_name, + input.size(channel_dim), + "input channels", + channels) || !check_int32_within_range( - context, op_name, input.size(2), "input height", input_h) || + context, op_name, input.size(height_dim), "input height", input_h) || !check_int32_within_range( - context, op_name, input.size(3), "input width", input_w) || + context, op_name, input.size(width_dim), "input width", input_w) || !check_int32_within_range( - context, op_name, output.size(2), "output height", output_h) || + context, + op_name, + output.size(height_dim), + "output height", + output_h) || !check_int32_within_range( - context, op_name, output.size(3), "output width", output_w)) { + context, op_name, output.size(width_dim), "output width", output_w)) { return false; } diff --git a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp index 39b6432c45a..66940f18997 100644 --- a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp @@ -1,4 +1,6 @@ /* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * Copyright 2025-2026 Arm Limited and/or its affiliates. * * This source code is licensed under the BSD-style license found in the @@ -66,7 +68,7 @@ bool validate_avg_pool2d_output_size( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_avg_pool2d_out( +static Tensor& quantized_avg_pool2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef kernel_size, @@ -77,6 +79,7 @@ Tensor& quantized_avg_pool2d_out( const int64_t multiplier, const int64_t shift, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { constexpr int32_t activation_min = std::numeric_limits::min(); constexpr int32_t activation_max = std::numeric_limits::max(); @@ -97,7 +100,7 @@ Tensor& quantized_avg_pool2d_out( activation_min, activation_max, pool_config, - true, + layout, true)) { return out; } @@ -153,5 +156,33 @@ Tensor& quantized_avg_pool2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_avg_pool2d_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::NCHWLogical, + 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 204a2b8369b..7865b50e486 100644 --- a/backends/cortex_m/ops/op_quantized_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_conv2d.cpp @@ -1,10 +1,14 @@ /* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * Copyright 2025-2026 Arm Limited and/or its affiliates. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ +#include + #include "cortex_m_ops_common.h" namespace cortex_m { @@ -25,7 +29,8 @@ bool validate_conv2d_arguments( const Int64ArrayRef& padding, const Int64ArrayRef& dilation, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvDim || weight.dim() != kConvDim || output.dim() != kConvDim) { ET_LOG(Error, "quantized_conv2d_out: tensors must be 4-D"); @@ -33,20 +38,22 @@ bool validate_conv2d_arguments( return false; } - // Check for channels_last dim_order (NHWC: 0, 2, 3, 1) - // Skip check if channels == 1, as dim_order is ambiguous in that case - if (input.size(1) > 1 && !is_channels_last_tensor(input)) { - ET_LOG( - Error, - "quantized_conv2d_out: input must have channels_last dim_order (NHWC)"); - context.fail(Error::InvalidArgument); - return false; - } - - if (output.size(1) > 1 && !is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( Error, - "quantized_conv2d_out: output must have channels_last dim_order (NHWC)"); + "quantized_conv2d_out: input and output must have channels_last dim_order"); context.fail(Error::InvalidArgument); return false; } @@ -78,7 +85,8 @@ bool validate_conv2d_arguments( return false; } - const int64_t out_channels = output.size(1); + const int64_t out_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (requantize_multipliers.size(0) != out_channels || requantize_shifts.size(0) != out_channels) { ET_LOG( @@ -94,7 +102,7 @@ bool validate_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_conv2d_out( +static Tensor& quantized_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -109,6 +117,7 @@ Tensor& quantized_conv2d_out( const int64_t activation_min, const int64_t activation_max, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { if (!validate_conv2d_arguments( context, @@ -120,23 +129,30 @@ Tensor& quantized_conv2d_out( padding, dilation, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t kernel_output_channels = static_cast(weight.size(0)); const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); const int32_t kernel_input_channels = static_cast(weight.size(3)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t input_offset_val = static_cast(input_offset); const int32_t output_offset_val = static_cast(output_offset); @@ -228,5 +244,41 @@ Tensor& quantized_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_conv2d_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::NCHWLogical, + 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 0793606de44..4aa58bb33dd 100644 --- a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp @@ -1,10 +1,14 @@ /* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * Copyright 2025-2026 Arm Limited and/or its affiliates. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ +#include + #include "cortex_m_ops_common.h" namespace cortex_m { @@ -26,7 +30,8 @@ bool validate_depthwise_conv2d_arguments( const Int64ArrayRef& dilation, const int64_t depth_multiplier, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvDim || weight.dim() != kConvDim || output.dim() != kConvDim) { ET_LOG(Error, "quantized_depthwise_conv2d_out: tensors must be 4-D"); @@ -55,7 +60,8 @@ bool validate_depthwise_conv2d_arguments( } const int64_t weight_output_channels = weight.size(3); - const int64_t output_channels = output.size(1); + const int64_t output_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (weight_output_channels != output_channels) { ET_LOG( Error, @@ -66,16 +72,22 @@ bool validate_depthwise_conv2d_arguments( return false; } - if (!is_channels_last_tensor(input)) { - ET_LOG( - Error, "quantized_depthwise_conv2d_out: input must be channels_last"); - context.fail(Error::InvalidArgument); - return false; - } - - if (!is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_depthwise_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( - Error, "quantized_depthwise_conv2d_out: output must be channels_last"); + Error, + "quantized_depthwise_conv2d_out: input and output must be channels_last"); context.fail(Error::InvalidArgument); return false; } @@ -108,7 +120,8 @@ bool validate_depthwise_conv2d_arguments( return false; } - const int64_t input_channels = input.size(1); + const int64_t input_channels = + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); // output_channels already extracted above for weight validation if (output_channels != input_channels * depth_multiplier) { ET_LOG( @@ -136,7 +149,7 @@ bool validate_depthwise_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_depthwise_conv2d_out( +static Tensor& quantized_depthwise_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -152,6 +165,7 @@ Tensor& quantized_depthwise_conv2d_out( const int64_t activation_min, const int64_t activation_max, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { if (!validate_depthwise_conv2d_arguments( context, @@ -164,23 +178,30 @@ Tensor& quantized_depthwise_conv2d_out( dilation, depth_multiplier, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); // Weight is in IHWO layout after permutation in the pass: [1, H, W, C_OUT] // For depthwise conv, this matches CMSIS-NN's expected format const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t depth_multiplier_val = static_cast(depth_multiplier); @@ -272,5 +293,43 @@ Tensor& quantized_depthwise_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_depthwise_conv2d_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::NCHWLogical, + 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 ca1b00ff340..68caa764ad5 100644 --- a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp @@ -1,4 +1,6 @@ /* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * Copyright 2026 Arm Limited and/or its affiliates. * * This source code is licensed under the BSD-style license found in the @@ -11,7 +13,7 @@ namespace cortex_m { namespace native { // cppcheck-suppress unusedFunction -Tensor& quantized_max_pool2d_out( +static Tensor& quantized_max_pool2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef kernel_size, @@ -23,6 +25,7 @@ Tensor& quantized_max_pool2d_out( const int64_t output_zero_point, const int64_t activation_min, const int64_t activation_max, + ActivationLayout layout, Tensor& out) { CmsisPool2DConfig pool_config; if (!prepare_cmsis_pool2d_config( @@ -37,7 +40,8 @@ Tensor& quantized_max_pool2d_out( ceil_mode, activation_min, activation_max, - pool_config)) { + pool_config, + layout)) { return out; } @@ -95,5 +99,35 @@ Tensor& quantized_max_pool2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_max_pool2d_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::NCHWLogical, + 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 04d57d4c693..fcfe78ce48d 100644 --- a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp @@ -24,7 +24,8 @@ bool validate_transpose_conv2d_arguments( const std::optional& bias, const Tensor& output, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvTransposeDim || weight.dim() != kConvTransposeDim || output.dim() != kConvTransposeDim) { ET_LOG(Error, "quantized_transpose_conv2d_out: tensors must be 4-D"); @@ -32,16 +33,22 @@ bool validate_transpose_conv2d_arguments( return false; } - if (!is_channels_last_tensor(input)) { - ET_LOG( - Error, "quantized_transpose_conv2d_out: input must be channels_last"); - context.fail(Error::InvalidArgument); - return false; - } - - if (!is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_transpose_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( - Error, "quantized_transpose_conv2d_out: output must be channels_last"); + Error, + "quantized_transpose_conv2d_out: input and output must be channels_last"); context.fail(Error::InvalidArgument); return false; } @@ -68,7 +75,8 @@ bool validate_transpose_conv2d_arguments( return false; } - const int64_t out_channels = output.size(1); + const int64_t out_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (requantize_multipliers.size(0) != out_channels || requantize_shifts.size(0) != out_channels) { ET_LOG( @@ -84,7 +92,7 @@ bool validate_transpose_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_transpose_conv2d_out( +static Tensor& quantized_transpose_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -101,6 +109,7 @@ Tensor& quantized_transpose_conv2d_out( const int64_t activation_max, const Tensor& scratch, const Tensor& output_scratch, + ActivationLayout layout, Tensor& out) { if (!validate_transpose_conv2d_arguments( context, @@ -109,23 +118,30 @@ Tensor& quantized_transpose_conv2d_out( bias, out, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t kernel_output_channels = static_cast(weight.size(0)); const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); const int32_t kernel_input_channels = static_cast(weight.size(3)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); if (kernel_output_channels != output_channels) { ET_LOG( @@ -246,5 +262,45 @@ Tensor& quantized_transpose_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_transpose_conv2d_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::NCHWLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/test/models/test_mobilenet_v3.py b/backends/cortex_m/test/models/test_mobilenet_v3.py index 08633d54dd6..2fccc89c131 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v3.py +++ b/backends/cortex_m/test/models/test_mobilenet_v3.py @@ -59,10 +59,6 @@ @parametrize( "test_case", test_cases, - xfails={ - "mobilenet_v3_small": "MLETORCH-1821 - Investigate mobilenet_v3_small flakyness" - }, - strict=False, ) def test_dialect_mv3(test_case): inputs = test_case.get_example_inputs() From a429bee4c75a2a2c944b43cee80a7582d2b765c8 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 11:54:58 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- .../absorb_boundary_layout_copies.py | 208 ++++++++++++++++ backends/transforms/targets.bzl | 34 +++ .../test_absorb_boundary_layout_copies.py | 228 ++++++++++++++++++ 3 files changed, 470 insertions(+) create mode 100644 backends/transforms/absorb_boundary_layout_copies.py create mode 100644 backends/transforms/test/test_absorb_boundary_layout_copies.py diff --git a/backends/transforms/absorb_boundary_layout_copies.py b/backends/transforms/absorb_boundary_layout_copies.py new file mode 100644 index 00000000000..d97ef62e900 --- /dev/null +++ b/backends/transforms/absorb_boundary_layout_copies.py @@ -0,0 +1,208 @@ +# 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. + +# pyre-unsafe + +from dataclasses import dataclass, field + +import torch + +from executorch.backends.transforms.channels_last_layout import is_layout_copy +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +# Elementwise per-tensor quantization is layout-agnostic, so a layout copy on +# the far side of one is still a boundary copy. Per-channel quantization is not: +# its axis is dimension-dependent. +_LAYOUT_AGNOSTIC_TARGETS = frozenset( + { + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + } +) + + +def _inverse(dims: tuple[int, ...]) -> list[int]: + inverse = [0] * len(dims) + for position, dim in enumerate(dims): + inverse[dim] = position + return inverse + + +@dataclass(frozen=True) +class BoundaryLayoutContract: + """Which method inputs and outputs changed layout, and to what. + + An entry ``{0: (0, 2, 3, 1)}`` in ``inputs`` means argument 0 must now be + passed as ``argument.permute(0, 2, 3, 1)``. An entry in ``outputs`` means + the returned tensor needs the same permutation applied to recover what the + method used to return. + """ + + inputs: dict[int, tuple[int, ...]] = field(default_factory=dict) + outputs: dict[int, tuple[int, ...]] = field(default_factory=dict) + + def __bool__(self) -> bool: + return bool(self.inputs or self.outputs) + + +class AbsorbBoundaryLayoutCopies(ExportPass): + """Move layout copies that sit on the method boundary into the signature. + + A layout region formed by ``ToContiguousChannelsLastPass`` is bracketed by + ``channels_last.permute_copy``. Those in the interior cancel against each + other; the ones on the boundary have nothing to cancel against and survive. + Deleting them and declaring the corresponding method input or output to be + channels-last moves the transpose to the caller, which is free whenever the + caller already has the data in that layout. + + Run this *after* region formation. Permuting the boundary first and hoping + the copies cancel is measurably worse: it inserts copies into graphs that + have no anchors at all, where nothing can cancel them. + + A copy need not touch the boundary directly. Quantized graphs put a + per-tensor ``quantize``/``dequantize`` in between, which reorders nothing, + so the search walks through those and relabels them on the way. + + Changing a method's layout is caller-visible, so the applied changes are + reported in ``contract`` rather than assumed. + """ + + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__() + self.exported_program = exported_program + self.contract = BoundaryLayoutContract() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + if graph_module is not self.exported_program.graph_module: + raise RuntimeError( + "AbsorbBoundaryLayoutCopies rewrites the ExportedProgram's graph " + "and signature together; run it as its own transform rather than " + "after a pass that replaces the graph module." + ) + + inputs = self._absorb_inputs(graph_module) + outputs = self._absorb_outputs(graph_module) + self.contract = BoundaryLayoutContract(inputs=inputs, outputs=outputs) + + modified = bool(self.contract) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + def _forward_to_copies(self, node: torch.fx.Node): + """Follow every path out of ``node`` until it reaches a layout copy. + + Returns the layout-agnostic nodes crossed, the copies terminating the + paths, and the permutation they agree on; ``None`` if any path ends + somewhere else or the copies disagree. + """ + interior: list[torch.fx.Node] = [] + copies: list[torch.fx.Node] = [] + dims: tuple[int, ...] | None = None + frontier = [node] + + while frontier: + users = list(frontier.pop().users) + if not users: + return None + for user in users: + if is_layout_copy(user): + user_dims = tuple(user.args[1]) + if dims is not None and user_dims != dims: + # A residual block feeds one value to several branches, + # each with its own copy. They collapse to a single + # contract entry only if they agree. + return None + dims = user_dims + copies.append(user) + elif user.target in _LAYOUT_AGNOSTIC_TARGETS: + if user in interior: + continue + interior.append(user) + frontier.append(user) + else: + return None + + return None if dims is None else (interior, copies, dims) + + def _backward_to_copy(self, result: torch.fx.Node): + """Walk back from a returned value to the layout copy that produced it.""" + interior: list[torch.fx.Node] = [] + current = result + + while True: + if is_layout_copy(current): + if len(current.users) != 1: + return None + source = current.args[0] + if not isinstance(source, torch.fx.Node): + return None + return interior, current, source, tuple(current.args[1]) + if ( + current.target not in _LAYOUT_AGNOSTIC_TARGETS + or len(current.users) != 1 + ): + return None + interior.append(current) + current = current.args[0] + if not isinstance(current, torch.fx.Node): + return None + + def _absorb_inputs(self, graph_module) -> dict[int, tuple[int, ...]]: + user_inputs = list(self.exported_program.graph_signature.user_inputs) + absorbed: dict[int, tuple[int, ...]] = {} + + for node in list(graph_module.graph.nodes): + if node.op != "placeholder" or node.name not in user_inputs: + continue + found = self._forward_to_copies(node) + if found is None: + continue + interior, copies, dims = found + + for member in (node, *interior): + member.meta["val"] = member.meta["val"].permute(dims) + for copy in copies: + copy.replace_all_uses_with(copy.args[0]) + graph_module.graph.erase_node(copy) + absorbed[user_inputs.index(node.name)] = dims + + return absorbed + + def _absorb_outputs(self, graph_module) -> dict[int, tuple[int, ...]]: + output_node = graph_module.graph.output_node() + results = list(output_node.args[0]) + specs = self.exported_program.graph_signature.output_specs + absorbed: dict[int, tuple[int, ...]] = {} + + for index, result in enumerate(results): + if not isinstance(result, torch.fx.Node): + continue + found = self._backward_to_copy(result) + if found is None: + continue + interior, copy, source, dims = found + + for member in interior: + member.meta["val"] = member.meta["val"].permute(_inverse(dims)) + if interior: + interior[-1].replace_input_with(copy, source) + else: + results[index] = source + # The manager re-derives the signature, but the direct + # exported_program= path does not. + if index < len(specs) and getattr(specs[index].arg, "name", None) == ( + result.name + ): + specs[index].arg.name = source.name + absorbed[index] = dims + + if absorbed: + output_node.args = (results,) + return absorbed diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index e15e089b051..fb7b3f54838 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -562,6 +562,40 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "absorb_boundary_layout_copies", + srcs = [ + "absorb_boundary_layout_copies.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_layout", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_absorb_boundary_layout_copies", + srcs = [ + "test/test_absorb_boundary_layout_copies.py", + # The permute-count matrix the absorption totals are measured on. + "test/test_to_contiguous_channels_last_pass.py", + ], + deps = [ + "//caffe2:torch", + ":absorb_boundary_layout_copies", + ":to_contiguous_channels_last_pass", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) + runtime.python_test( name = "test_replace_ops_with_channels_last_variants", srcs = [ diff --git a/backends/transforms/test/test_absorb_boundary_layout_copies.py b/backends/transforms/test/test_absorb_boundary_layout_copies.py new file mode 100644 index 00000000000..e3675d51d05 --- /dev/null +++ b/backends/transforms/test/test_absorb_boundary_layout_copies.py @@ -0,0 +1,228 @@ +# 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.transforms.absorb_boundary_layout_copies import ( + AbsorbBoundaryLayoutCopies, +) +from executorch.backends.transforms.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + +_LAYOUT_COPY = exir_ops.edge.channels_last.permute_copy.default + + +class Conv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class ResidualConvPool(torch.nn.Module): + """One input feeding two branches, so the region brackets it twice.""" + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.pool = torch.nn.MaxPool2d(3, stride=1, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + self.pool(x) + + +def _count(graph_module, target) -> int: + return sum( + node.op == "call_function" and node.target == target + for node in graph_module.graph.nodes + ) + + +def _lower(module, inputs): + module.eval() + with torch.no_grad(): + exported = torch.export.export(module, inputs) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + edge = edge.transform([ToContiguousChannelsLastPass(edge.exported_program())]) + return edge + + +def _absorb(edge): + layout_pass = AbsorbBoundaryLayoutCopies(edge.exported_program()) + return edge.transform([layout_pass]), layout_pass.contract + + +def _run(edge, contract, inputs): + args = list(inputs) + for index, dims in contract.inputs.items(): + args[index] = args[index].permute(list(dims)).contiguous() + result = edge.exported_program().module()(*args) + results = list(result) if isinstance(result, (tuple, list)) else [result] + for index, dims in contract.outputs.items(): + results[index] = results[index].permute(list(dims)) + return results[0] if len(results) == 1 else results + + +@pytest.mark.parametrize("module", [Conv(), ResidualConvPool()]) +def test_boundary_copies_are_absorbed_and_numerics_hold(module) -> None: + inputs = (torch.randn(1, 4, 8, 8),) + expected = module.eval()(*inputs) + edge = _lower(module, inputs) + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) > 0 + + edge, contract = _absorb(edge) + + assert contract.inputs and contract.outputs + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) == 0 + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_fan_out_collapses_to_one_contract_entry() -> None: + """Both branches of a residual share the input, so one entry covers them.""" + module = ResidualConvPool() + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(module, inputs) + copies_on_input = [ + node + for node in edge.exported_program().graph_module.graph.nodes + if node.op == "placeholder" + and node.name in edge.exported_program().graph_signature.user_inputs + for _ in node.users + ] + assert len(copies_on_input) > 1 + + _, contract = _absorb(edge) + + assert list(contract.inputs) == [0] + + +def test_mixed_users_are_left_alone() -> None: + """An input consumed both by a layout region and directly is not a boundary.""" + + class MixedUse(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + x + + inputs = (torch.randn(1, 4, 8, 8),) + module = MixedUse() + expected = module.eval()(*inputs) + edge = _lower(module, inputs) + + edge, contract = _absorb(edge) + + assert 0 not in contract.inputs + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_absorbing_is_idempotent() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(Conv(), inputs) + edge, first = _absorb(edge) + edge, second = _absorb(edge) + + assert first + assert not second + + +def test_signature_stays_valid() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(Conv(), inputs) + edge, contract = _absorb(edge) + + assert contract + edge.exported_program()._validate() + + +# The layout pass is measured against these 36 models in +# test_to_contiguous_channels_last_pass.py, which pins its own per-case counts +# but stops before absorption. These are the totals across that matrix. +_MATRIX_BASELINE_PERMUTES = 138 +_MATRIX_LAYOUT_ONLY_PERMUTES = 156 +_MATRIX_ABSORBED_PERMUTES = 137 + +# Absorption does not destroy 19 permutes, it moves them across the method +# boundary: 17 become obligations on the caller. Pinning both numbers keeps the +# in-graph count honest, since it could otherwise be driven to zero by handing +# the caller unlimited work. The gap is the genuine saving, and it comes from +# fan-out — one placeholder feeding several branches needs several copies but +# only one contract entry. +_MATRIX_CONTRACT_ENTRIES = 17 + +# Cases still above baseline after absorbing. Every one of these is an internal +# copy left by a region that a non-permutable op (batch_norm, linear) cut in +# two, which is region-merging work rather than boundary work. +_MATRIX_RESIDUAL_REGRESSIONS = { + "conv2d_rank3", + "model_1_conv_maxpool_residual_linear", + "model_8_conv_batchnorm_maxpool_residual", + "model_9_dilated_conv_batchnorm_avgpool_residual", + "views", +} + +_PERMUTE_TARGETS = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.channels_last.permute_copy.default, +} + + +def _permutes(edge) -> int: + return sum( + node.op == "call_function" and node.target in _PERMUTE_TARGETS + for node in edge.exported_program().graph.nodes + ) + + +def test_absorption_pays_for_the_layout_pass_across_the_model_matrix() -> None: + from executorch.backends.transforms.test.test_to_contiguous_channels_last_pass import ( + cases, + ) + + baseline = layout_only = absorbed = contract_entries = 0 + regressions = set() + for name, case in cases.items(): + case.module.eval() + with torch.no_grad(): + exported = torch.export.export(case.module, case.inputs) + config = EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True) + case_baseline = _permutes(to_edge(exported, compile_config=config)) + + edge = to_edge(exported, compile_config=config) + edge = edge.transform( + [ToContiguousChannelsLastPass(edge.exported_program())] + ) + case_layout = _permutes(edge) + + absorb = AbsorbBoundaryLayoutCopies(edge.exported_program()) + edge = edge.transform([absorb]) + case_absorbed = _permutes(edge) + + baseline += case_baseline + layout_only += case_layout + absorbed += case_absorbed + contract_entries += len(absorb.contract.inputs) + len(absorb.contract.outputs) + if case_absorbed > case_baseline: + regressions.add(name) + + assert baseline == _MATRIX_BASELINE_PERMUTES + assert layout_only == _MATRIX_LAYOUT_ONLY_PERMUTES + assert absorbed == _MATRIX_ABSORBED_PERMUTES + assert contract_entries == _MATRIX_CONTRACT_ENTRIES + assert regressions == _MATRIX_RESIDUAL_REGRESSIONS