Skip to content

Cortex-M: add explicit-NHWC channel-broadcast ops #21966

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/11/head
Choose a base branch
from
Draft

Cortex-M: add explicit-NHWC channel-broadcast ops #21966

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
8690ea3
Update
rascani Aug 20, 2026
97fa49b
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
17 changes: 13 additions & 4 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 @@ -149,7 +149,10 @@ inline bool is_channels_last_tensor(const Tensor& tensor) {
return tensor.dim_order() == channels_last_order;
}

inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) {
inline bool is_channel_broadcast(
const Tensor& tensor1,
const Tensor& tensor2,
int64_t channel_dim) {
if (tensor1.dim() != tensor2.dim()) {
return false;
}
Expand All @@ -158,16 +161,22 @@ inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) {
return false;
}

if (tensor1.size(1) != tensor2.size(1)) {
if (tensor1.size(channel_dim) != tensor2.size(channel_dim)) {
return false;
}

const bool tensor1_channels_only = tensor1.numel() == tensor1.size(1);
const bool tensor2_channels_only = tensor2.numel() == tensor2.size(1);
const bool tensor1_channels_only =
tensor1.numel() == tensor1.size(channel_dim);
const bool tensor2_channels_only =
tensor2.numel() == tensor2.size(channel_dim);

return tensor1_channels_only || tensor2_channels_only;
}

inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) {
return is_channel_broadcast(tensor1, tensor2, 1);
}

inline bool check_int32_within_range(
KernelRuntimeContext& context,
const char* op_name,
Expand Down
129 changes: 120 additions & 9 deletions backends/cortex_m/ops/op_quantized_add.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@ namespace cortex_m {
namespace native {
using KernelRuntimeContext = torch::executor::KernelRuntimeContext;

// cppcheck-suppress unusedFunction
Tensor& quantized_add_out(
static Tensor& quantized_add_out_impl(
KernelRuntimeContext& context,
const Tensor& input1_int8,
const int64_t input1_zero_point,
Expand All @@ -29,16 +28,51 @@ Tensor& quantized_add_out(
const int64_t output_shift,
const int64_t activation_min,
const int64_t activation_max,
ActivationLayout layout,
const char* op_name,
Tensor& out) {
// Validate tensor types and dim order
bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8);
const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1;
bool channel_broadcast =
is_channel_broadcast(input1_int8, input2_int8, channel_dim);
validate_cmsis_nn_tensor_requirements(
input1_int8,
input2_int8,
out,
ScalarType::Char,
/*require_channels_last=*/channel_broadcast,
/*require_channels_last=*/
channel_broadcast && layout == ActivationLayout::NCHWLogical,
/*require_same_sizes=*/!channel_broadcast);
if (layout == ActivationLayout::NHWCLogical) {
ET_CHECK_MSG(
input1_int8.dim() == 4 && input2_int8.dim() == 4 && out.dim() == 4,
"%s: tensors must be 4-D",
op_name);
ET_CHECK_MSG(
executorch::runtime::is_contiguous_dim_order(
input1_int8.dim_order().data(), input1_int8.dim_order().size()) &&
executorch::runtime::is_contiguous_dim_order(
input2_int8.dim_order().data(),
input2_int8.dim_order().size()) &&
executorch::runtime::is_contiguous_dim_order(
out.dim_order().data(), out.dim_order().size()),
"%s: tensors must use contiguous dimension order",
op_name);
} else if (channel_broadcast) {
ET_CHECK_MSG(
is_channels_last_tensor(input1_int8) &&
is_channels_last_tensor(input2_int8) &&
is_channels_last_tensor(out),
"%s: channel-broadcast tensors must use channels-last dimension order",
op_name);
}
if (channel_broadcast) {
const Tensor& full_input =
input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8;
ET_CHECK_MSG(
out.sizes() == full_input.sizes(),
"%s: output must have the broadcast result shape",
op_name);
}

// Validate quantization parameters
validate_quantization_params(
Expand All @@ -54,7 +88,8 @@ Tensor& quantized_add_out(

ET_LOG(
Debug,
"quantized_add_out: input1_int8.sizes() = %zu",
"%s: input1_int8.sizes() = %zu",
op_name,
input1_int8.sizes().size());

int32_t zp1 = static_cast<int32_t>(input1_zero_point);
Expand Down Expand Up @@ -101,7 +136,7 @@ Tensor& quantized_add_out(
std::swap<int>(input1_shift_val, input2_shift_val);
std::swap<int8_t*>(input1_ptr, input2_ptr);
}
adds_per_loop = input1_int8.size(1);
adds_per_loop = input1_int8.size(channel_dim);
} else {
adds_per_loop = out.numel();
}
Expand Down Expand Up @@ -130,7 +165,8 @@ Tensor& quantized_add_out(
if (status != ARM_CMSIS_NN_SUCCESS) {
ET_LOG(
Error,
"quantized_add_out: arm_elementwise_add_s8 failed with status [%d]",
"%s: arm_elementwise_add_s8 failed with status [%d]",
op_name,
status);

context.fail(Error::Internal); // Fail the execution context
Expand All @@ -139,10 +175,85 @@ Tensor& quantized_add_out(
}
ET_LOG(
Debug,
"quantized_add_out: Successfully completed with AoT-computed parameters!");
"%s: Successfully completed with AoT-computed parameters!",
op_name);

return out;
}

// cppcheck-suppress unusedFunction
Tensor& quantized_add_out(
KernelRuntimeContext& context,
const Tensor& input1_int8,
const int64_t input1_zero_point,
const int64_t input1_multiplier,
const int64_t input1_shift,
const Tensor& input2_int8,
const int64_t input2_zero_point,
const int64_t input2_multiplier,
const int64_t input2_shift,
const int64_t output_zero_point,
const int64_t output_multiplier,
const int64_t output_shift,
const int64_t activation_min,
const int64_t activation_max,
Tensor& out) {
return quantized_add_out_impl(
context,
input1_int8,
input1_zero_point,
input1_multiplier,
input1_shift,
input2_int8,
input2_zero_point,
input2_multiplier,
input2_shift,
output_zero_point,
output_multiplier,
output_shift,
activation_min,
activation_max,
ActivationLayout::NCHWLogical,
"quantized_add_out",
out);
}

// cppcheck-suppress unusedFunction
Tensor& quantized_add_nhwc_out(
KernelRuntimeContext& context,
const Tensor& input1_int8,
const int64_t input1_zero_point,
const int64_t input1_multiplier,
const int64_t input1_shift,
const Tensor& input2_int8,
const int64_t input2_zero_point,
const int64_t input2_multiplier,
const int64_t input2_shift,
const int64_t output_zero_point,
const int64_t output_multiplier,
const int64_t output_shift,
const int64_t activation_min,
const int64_t activation_max,
Tensor& out) {
return quantized_add_out_impl(
context,
input1_int8,
input1_zero_point,
input1_multiplier,
input1_shift,
input2_int8,
input2_zero_point,
input2_multiplier,
input2_shift,
output_zero_point,
output_multiplier,
output_shift,
activation_min,
activation_max,
ActivationLayout::NHWCLogical,
"quantized_add_nhwc_out",
out);
}

} // namespace native
} // namespace cortex_m
100 changes: 92 additions & 8 deletions backends/cortex_m/ops/op_quantized_mul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,7 @@ constexpr int32_t kInt8ActivationMax = std::numeric_limits<int8_t>::max();

using KernelRuntimeContext = torch::executor::KernelRuntimeContext;

// cppcheck-suppress unusedFunction
Tensor& quantized_mul_out(
static Tensor& quantized_mul_out_impl(
KernelRuntimeContext& context,
const Tensor& input1_int8,
const int64_t input1_zero_point,
Expand All @@ -28,17 +27,51 @@ Tensor& quantized_mul_out(
const int64_t output_zero_point,
const int64_t output_multiplier,
const int64_t output_shift,
ActivationLayout layout,
const char* op_name,
Tensor& out) {
// Validate tensor types and quantization parameters

bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8);
const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1;
bool channel_broadcast =
is_channel_broadcast(input1_int8, input2_int8, channel_dim);
validate_cmsis_nn_tensor_requirements(
input1_int8,
input2_int8,
out,
ScalarType::Char,
/*require_channels_last=*/channel_broadcast,
/*require_channels_last=*/
channel_broadcast && layout == ActivationLayout::NCHWLogical,
/*require_same_sizes=*/!channel_broadcast);
if (layout == ActivationLayout::NHWCLogical) {
ET_CHECK_MSG(
input1_int8.dim() == 4 && input2_int8.dim() == 4 && out.dim() == 4,
"%s: tensors must be 4-D",
op_name);
ET_CHECK_MSG(
executorch::runtime::is_contiguous_dim_order(
input1_int8.dim_order().data(), input1_int8.dim_order().size()) &&
executorch::runtime::is_contiguous_dim_order(
input2_int8.dim_order().data(),
input2_int8.dim_order().size()) &&
executorch::runtime::is_contiguous_dim_order(
out.dim_order().data(), out.dim_order().size()),
"%s: tensors must use contiguous dimension order",
op_name);
} else if (channel_broadcast) {
ET_CHECK_MSG(
is_channels_last_tensor(input1_int8) &&
is_channels_last_tensor(input2_int8) &&
is_channels_last_tensor(out),
"%s: channel-broadcast tensors must use channels-last dimension order",
op_name);
}
if (channel_broadcast) {
const Tensor& full_input =
input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8;
ET_CHECK_MSG(
out.sizes() == full_input.sizes(),
"%s: output must have the broadcast result shape",
op_name);
}

const int32_t kIdentityMultiplier(/*value=*/1);
const int32_t kZeroShift(/*value=*/0);
Expand Down Expand Up @@ -70,7 +103,7 @@ Tensor& quantized_mul_out(
std::swap<int8_t*>(input1_ptr, input2_ptr);
}

muls_per_loop = input1_int8.size(1);
muls_per_loop = input1_int8.size(channel_dim);
} else {
muls_per_loop = out.numel();
}
Expand Down Expand Up @@ -108,7 +141,8 @@ Tensor& quantized_mul_out(
if (status != ARM_CMSIS_NN_SUCCESS) {
ET_LOG(
Error,
"quantized_mul_out: arm_elementwise_mul_s8 failed with status [%d]",
"%s: arm_elementwise_mul_s8 failed with status [%d]",
op_name,
status);
context.fail(Error::Internal);
return out;
Expand All @@ -117,5 +151,55 @@ Tensor& quantized_mul_out(
return out;
}

// cppcheck-suppress unusedFunction
Tensor& quantized_mul_out(
KernelRuntimeContext& context,
const Tensor& input1_int8,
const int64_t input1_zero_point,
const Tensor& input2_int8,
const int64_t input2_zero_point,
const int64_t output_zero_point,
const int64_t output_multiplier,
const int64_t output_shift,
Tensor& out) {
return quantized_mul_out_impl(
context,
input1_int8,
input1_zero_point,
input2_int8,
input2_zero_point,
output_zero_point,
output_multiplier,
output_shift,
ActivationLayout::NCHWLogical,
"quantized_mul_out",
out);
}

// cppcheck-suppress unusedFunction
Tensor& quantized_mul_nhwc_out(
KernelRuntimeContext& context,
const Tensor& input1_int8,
const int64_t input1_zero_point,
const Tensor& input2_int8,
const int64_t input2_zero_point,
const int64_t output_zero_point,
const int64_t output_multiplier,
const int64_t output_shift,
Tensor& out) {
return quantized_mul_out_impl(
context,
input1_int8,
input1_zero_point,
input2_int8,
input2_zero_point,
output_zero_point,
output_multiplier,
output_shift,
ActivationLayout::NHWCLogical,
"quantized_mul_nhwc_out",
out);
}

} // namespace native
} // namespace cortex_m
Loading
Loading