Cortex_m backend: Simplify add + linear fusion passes (#15526)

Reuses the FoldAndAnnotateQParamsPass from the Arm backend to greatly
simplify the logic for fusing the ops.

Additionally updates the linear kernel to be numerically correct and
computes the kernel_sum aot in the quantized_linear_fusion pass. Note
that since this replaces the bias node it typically causes no extra
memory usage.

Updates the Linear tests to mirror this, including removing the various
matmul tests. Since the linear is handled as a separate op rather than a
particular type of matmul these tests are not related anymore.

Removes unnecessary stub definitions in operators.py, operators.yaml and
op_quantized_linear.cpp

Leaving a few TODO:s since the patch is large already.


Signed-off-by: Adrian Lundell <adrian.lundell@arm.com>

Author: AdrianLundell

backends/cortex_m/ops/cortex_m_ops_common.h

@@ -22,6 +22,7 @@ using Tensor = torch::executor::Tensor;
 using ScalarType = executorch::aten::ScalarType;
 using Scalar = torch::executor::Scalar;
 using Error = executorch::runtime::Error;
+using IntArrayRef = executorch::aten::ArrayRef<int64_t>;
 
 // From arm_nn_math_types.h
 #define ARM_NN_Q31_MAX ((int32_t)(0x7FFFFFFFL))

backends/cortex_m/ops/op_quantized_linear.cpp

@@ -1,12 +1,12 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
+ * Copyright 2025 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 "cmsis_scratch_buffer_context.h"
 #include "cortex_m_ops_common.h"
 
 extern "C" {
@@ -20,151 +20,90 @@ using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
 Tensor& quantized_linear_out(
     KernelRuntimeContext& context,
     const Tensor& input,
-    const Scalar& input_zero_point,
-    const Scalar& input_multiplier,
-    const Scalar& input_shift,
     const Tensor& weights,
-    const Tensor& weight_zero_point,
-    const Tensor& weight_multiplier,
-    const Tensor& weight_shift,
     const torch::executor::optional<Tensor>& bias,
-    const Tensor& bias_multiplier,
-    const Tensor& bias_shift,
-    const Tensor& scratch_buffer,
-    const Scalar& output_zero_point,
-    const Scalar& in_features,
-    const Scalar& out_features,
+    const torch::executor::optional<Tensor>& kernel_sum,
+    const Scalar& input_offset,
+    const Scalar& filter_offset,
+    const Scalar& output_offset,
+    const IntArrayRef requantize_multipliers,
+    const IntArrayRef requantize_shifts,
+    const Scalar& activation_max,
+    const Scalar& activation_min,
     Tensor& out) {
   ET_LOG(Info, "quantized_linear_out: called");
-  validate_cmsis_nn_tensor_requirements(input, weights, out);
-
-  ET_CHECK_MSG(
-      scratch_buffer.scalar_type() == ScalarType::Char,
-      "Scratch buffer must be int8");
-
-  const int32_t batch_size = input.size(0);
-  const int32_t in_feat = static_cast<int32_t>(in_features.to<int64_t>());
-  const int32_t out_feat = static_cast<int32_t>(out_features.to<int64_t>());
-  const int32_t input_zp = static_cast<int32_t>(input_zero_point.to<int64_t>());
-  const int32_t output_zp =
-      static_cast<int32_t>(output_zero_point.to<int64_t>());
-  const bool is_per_channel = (weight_zero_point.numel() > 1);
 
   const int8_t* input_data = input.const_data_ptr<int8_t>();
   const int8_t* weight_data = weights.const_data_ptr<int8_t>();
   const int32_t* bias_data =
       bias.has_value() ? bias.value().const_data_ptr<int32_t>() : nullptr;
+  int32_t* kernel_sum_data =
+      kernel_sum.has_value() ? kernel_sum.value().data_ptr<int32_t>() : nullptr;
   int8_t* output_data = out.mutable_data_ptr<int8_t>();
-  const int32_t* weight_zp_data = weight_zero_point.const_data_ptr<int32_t>();
-  const int32_t* weight_mult_data = weight_multiplier.const_data_ptr<int32_t>();
-  const int32_t* weight_shift_data = weight_shift.const_data_ptr<int32_t>();
-
-  if (!validate_per_channel_quant_params(
-          weight_mult_data, weight_shift_data, out_feat)) {
-    context.fail(Error::InvalidArgument);
-    return out;
-  }
-
-  // Initialize scratch buffer context (validates early)
-  CMSISScratchBufferContext scratch_ctx(
-      const_cast<Tensor&>(scratch_buffer), weights, weight_zero_point, bias);
 
-  scratch_ctx.compute_kernel_sums_if_needed();
-  cmsis_nn_context ctx = scratch_ctx.get_cmsis_ctx();
+  cmsis_nn_context ctx;
+  ctx.size = 0; // Not used in CMSIS-NN
+  ctx.buf = kernel_sum_data;
 
   // Setup CMSIS-NN parameters
   cmsis_nn_fc_params fc_params;
-  fc_params.input_offset = -input_zp;
-  fc_params.output_offset = output_zp;
-  fc_params.activation.min = std::numeric_limits<int8_t>::min();
-  fc_params.activation.max = std::numeric_limits<int8_t>::max();
-
-  cmsis_nn_dims input_dims = {1, 1, 1, in_feat};
+  fc_params.input_offset = static_cast<int32_t>(input_offset.to<int64_t>());
+  fc_params.filter_offset = static_cast<int32_t>(filter_offset.to<int64_t>());
+  fc_params.output_offset = static_cast<int32_t>(output_offset.to<int64_t>());
+  fc_params.activation.min = static_cast<int32_t>(activation_min.to<int64_t>());
+  fc_params.activation.max = static_cast<int32_t>(activation_max.to<int64_t>());
+
+  cmsis_nn_per_tensor_quant_params per_tensor_quant_params;
+  per_tensor_quant_params.multiplier =
+      static_cast<int32_t>(requantize_multipliers.at(0));
+  per_tensor_quant_params.shift = static_cast<int32_t>(requantize_shifts.at(0));
+
+  auto in_feat = input.size(input.dim() - 1);
+  auto out_feat = out.size(out.dim() - 1);
+  auto batches = 1;
+  for (size_t i = 0; i < input.dim() - 1; i++) {
+    batches *= input.size(i);
+  }
+  ET_LOG(
+      Info,
+      "in features: %d, out_features: %d, batches: %d, kernel_sum_size: %d",
+      in_feat,
+      out_feat,
+      batches,
+      kernel_sum.has_value() ? kernel_sum.value().numel() : 0);
+  ET_LOG(
+      Info,
+      "kernel_sum[0]: %d, kernel_sum[1]: %d",
+      kernel_sum_data != nullptr ? kernel_sum_data[0] : -1,
+      kernel_sum_data != nullptr ? kernel_sum_data[1] : -1);
+  cmsis_nn_dims input_dims = {batches, 1, 1, in_feat};
   cmsis_nn_dims filter_dims = {in_feat, 1, 1, out_feat};
   cmsis_nn_dims bias_dims = {1, 1, 1, out_feat};
-  cmsis_nn_dims output_dims = {1, 1, 1, out_feat};
-
-  arm_cmsis_nn_status status;
-  for (int32_t b = 0; b < batch_size; b++) {
-    const int8_t* batch_input = input_data + b * in_feat;
-    int8_t* batch_output = output_data + b * out_feat;
-
-    ET_CHECK_MSG(
-        batch_input != nullptr && weight_data != nullptr,
-        "Null input pointers");
-    ET_CHECK_MSG(in_feat > 0 && out_feat > 0, "Invalid dimensions");
-
-    if (is_per_channel) {
-      cmsis_nn_per_channel_quant_params per_channel_quant_params;
-      per_channel_quant_params.multiplier =
-          const_cast<int32_t*>(weight_mult_data);
-      per_channel_quant_params.shift = const_cast<int32_t*>(weight_shift_data);
-
-      status = arm_fully_connected_per_channel_s8(
-          &ctx,
-          &fc_params,
-          &per_channel_quant_params,
-          &input_dims,
-          batch_input,
-          &filter_dims,
-          weight_data,
-          &bias_dims,
-          bias_data,
-          &output_dims,
-          batch_output);
-    } else {
-      fc_params.filter_offset = -weight_zp_data[0];
-      cmsis_nn_per_tensor_quant_params per_tensor_quant_params;
-      per_tensor_quant_params.multiplier = weight_mult_data[0];
-      per_tensor_quant_params.shift = weight_shift_data[0];
-
-      status = arm_fully_connected_s8(
-          &ctx,
-          &fc_params,
-          &per_tensor_quant_params,
-          &input_dims,
-          batch_input,
-          &filter_dims,
-          weight_data,
-          &bias_dims,
-          bias_data,
-          &output_dims,
-          batch_output);
-    }
-
-    if (status != ARM_CMSIS_NN_SUCCESS) {
-      ET_LOG(
-          Error,
-          "quantized_linear_out: CMSIS-NN failed with status [%d]",
-          status);
-      context.fail(Error::Internal);
-      return out;
-    }
+  cmsis_nn_dims output_dims = {batches, 1, 1, out_feat};
+
+  arm_cmsis_nn_status status = arm_fully_connected_s8(
+      &ctx,
+      &fc_params,
+      &per_tensor_quant_params,
+      &input_dims,
+      input_data,
+      &filter_dims,
+      weight_data,
+      &bias_dims,
+      bias_data,
+      &output_dims,
+      output_data);
+
+  if (status != ARM_CMSIS_NN_SUCCESS) {
+    ET_LOG(
+        Error,
+        "quantized_linear_out: CMSIS-NN failed with status [%d]",
+        status);
+    context.fail(Error::Internal);
+    return out;
   }
-  return out;
-}
 
-// Functional variant (stub, not used at runtime)
-Tensor quantized_linear(
-    KernelRuntimeContext& context,
-    const Tensor& input,
-    const Scalar& input_zero_point,
-    const Scalar& input_multiplier,
-    const Scalar& input_shift,
-    const Tensor& weights,
-    const Tensor& weight_zero_point,
-    const Tensor& weight_multiplier,
-    const Tensor& weight_shift,
-    const torch::executor::optional<Tensor>& bias,
-    const Tensor& bias_multiplier,
-    const Tensor& bias_shift,
-    const Tensor& scratch_buffer,
-    const Scalar& output_zero_point,
-    const Scalar& in_features,
-    const Scalar& out_features) {
-  ET_LOG(Info, "quantized_linear: called");
-  assert(false);
-  return const_cast<Tensor&>(input);
+  return out;
 }
 
 } // namespace native