Skip to content

Cortex-M: share layout-aware CMSIS-NN kernels #21964

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
rascani wants to merge 2 commits into
base: gh/rascani/9/head
Choose a base branch
from
Draft

Cortex-M: share layout-aware CMSIS-NN kernels #21964

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ec4b327
Update
rascani Aug 20, 2026
643cd15
Update
rascani Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 37 additions & 11 deletions backends/cortex_m/ops/cortex_m_ops_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

#include <executorch/kernels/portable/cpu/util/elementwise_util.h>
#include <executorch/kernels/portable/cpu/util/kernel_ops_util.h>
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
#include <executorch/runtime/platform/assert.h>

#include <cinttypes>
Expand All @@ -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,
Expand Down Expand Up @@ -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);
Expand All @@ -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",
Expand All @@ -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,
Expand Down Expand Up @@ -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;
}

Expand Down
35 changes: 33 additions & 2 deletions backends/cortex_m/ops/op_quantized_avg_pool2d.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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<int8_t>::min();
constexpr int32_t activation_max = std::numeric_limits<int8_t>::max();
Expand All @@ -97,7 +100,7 @@ Tensor& quantized_avg_pool2d_out(
activation_min,
activation_max,
pool_config,
true,
layout,
true)) {
return out;
}
Expand Down Expand Up @@ -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
96 changes: 74 additions & 22 deletions backends/cortex_m/ops/op_quantized_conv2d.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 <executorch/runtime/core/exec_aten/util/dim_order_util.h>

#include "cortex_m_ops_common.h"

namespace cortex_m {
Expand All @@ -25,28 +29,31 @@ 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");
context.fail(Error::InvalidArgument);
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;
}
Expand Down Expand Up @@ -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(
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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<int32_t>(input.size(0));
const int32_t input_channels = static_cast<int32_t>(input.size(1));
const int32_t input_height = static_cast<int32_t>(input.size(2));
const int32_t input_width = static_cast<int32_t>(input.size(3));
const int32_t input_channels = static_cast<int32_t>(
input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1));
const int32_t input_height = static_cast<int32_t>(
input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2));
const int32_t input_width = static_cast<int32_t>(
input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3));

const int32_t kernel_output_channels = static_cast<int32_t>(weight.size(0));
const int32_t kernel_height = static_cast<int32_t>(weight.size(1));
const int32_t kernel_width = static_cast<int32_t>(weight.size(2));
const int32_t kernel_input_channels = static_cast<int32_t>(weight.size(3));

const int32_t output_channels = static_cast<int32_t>(out.size(1));
const int32_t output_height = static_cast<int32_t>(out.size(2));
const int32_t output_width = static_cast<int32_t>(out.size(3));
const int32_t output_channels = static_cast<int32_t>(
out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1));
const int32_t output_height = static_cast<int32_t>(
out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2));
const int32_t output_width = static_cast<int32_t>(
out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3));

const int32_t input_offset_val = static_cast<int32_t>(input_offset);
const int32_t output_offset_val = static_cast<int32_t>(output_offset);
Expand Down Expand Up @@ -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<Tensor>& 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
Loading
Loading