Per-channel int8 quantized linear (#24397)

Summary:

`cortex_m::quantized_linear` supports only per-tensor weights: it calls `arm_fully_connected_s8` with a single requantization multiplier and shift. This adds per-channel linear weights.

With `CortexMQuantizer.set_node_target` (previous diff), `set_node_target(torch.ops.aten.linear.default, INT8_PER_CHANNEL_CONFIG)` gives per-channel int8 linear weights; without it, linear layers keep `per_tensor_config`, so existing exports are unchanged. `CortexMLinearCheck` rejects per-channel weights that are not quantized per output feature (axis 0, or unset as the convolution check allows), and the lowering checks that it gets one weight scale per output feature, since the kernel indexes its multipliers by output feature.

Per-channel linear weights lower to a new `cortex_m::quantized_linear_per_channel`. It matches `quantized_linear` but takes int32 multiplier and shift tensors, as the convolution ops do, and calls `arm_fully_connected_per_channel_s8`. CMSIS-NN takes int32 arrays, so tensors avoid any runtime conversion. `int[]` lists would also cost about 20 bytes per output feature in the program, where a tensor costs 4. The two C++ entry points share one implementation that takes the CMSIS-NN call; `quantized_linear` is unchanged. The Python reference requantizes each output feature with its own parameters.

This diff was authored with Claude Code.

Differential Revision: D126929702

Author: pssrawat

backends/cortex_m/ops/op_quantized_linear.cpp

@@ -13,8 +13,9 @@ namespace cortex_m {
 namespace native {
 using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
 
-// cppcheck-suppress unusedFunction
-Tensor& quantized_linear_out(
+namespace {
+Tensor& quantized_linear_impl(
+    const char* op_name,
     KernelRuntimeContext& context,
     const Tensor& input,
     const Tensor& weights,
@@ -23,12 +24,11 @@ Tensor& quantized_linear_out(
     const int64_t input_offset,
     const int64_t filter_offset,
     const int64_t output_offset,
-    const Int64ArrayRef requantize_multipliers,
-    const Int64ArrayRef requantize_shifts,
     const int64_t activation_max,
     const int64_t activation_min,
+    const cmsis_nn_quant_params& quant_params,
     Tensor& out) {
-  ET_LOG(Debug, "quantized_linear_out: called");
+  ET_LOG(Debug, "%s: called", op_name);
 
   const int8_t* input_data = input.const_data_ptr<int8_t>();
   const int8_t* weight_data = weights.const_data_ptr<int8_t>();
@@ -50,11 +50,6 @@ Tensor& quantized_linear_out(
   fc_params.activation.min = static_cast<int32_t>(activation_min);
   fc_params.activation.max = static_cast<int32_t>(activation_max);
 
-  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;
@@ -78,10 +73,10 @@ Tensor& quantized_linear_out(
   cmsis_nn_dims bias_dims = {1, 1, 1, out_feat};
   cmsis_nn_dims output_dims = {batches, 1, 1, out_feat};
 
-  arm_cmsis_nn_status status = arm_fully_connected_s8(
+  const arm_cmsis_nn_status status = arm_fully_connected_wrapper_s8(
       &ctx,
       &fc_params,
-      &per_tensor_quant_params,
+      &quant_params,
       &input_dims,
       input_data,
       &filter_dims,
@@ -92,16 +87,105 @@ Tensor& quantized_linear_out(
       output_data);
 
   if (status != ARM_CMSIS_NN_SUCCESS) {
-    ET_LOG(
-        Error,
-        "quantized_linear_out: CMSIS-NN failed with status [%d]",
-        status);
+    ET_LOG(Error, "%s: CMSIS-NN failed with status [%d]", op_name, status);
     context.fail(Error::Internal);
     return out;
   }
 
   return out;
 }
+} // namespace
+
+// cppcheck-suppress unusedFunction
+Tensor& quantized_linear_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    const Tensor& weights,
+    const std::optional<Tensor>& bias,
+    const std::optional<Tensor>& kernel_sum,
+    const int64_t input_offset,
+    const int64_t filter_offset,
+    const int64_t output_offset,
+    const Int64ArrayRef requantize_multipliers,
+    const Int64ArrayRef requantize_shifts,
+    const int64_t activation_max,
+    const int64_t activation_min,
+    Tensor& out) {
+  int32_t multiplier = static_cast<int32_t>(requantize_multipliers.at(0));
+  int32_t shift = static_cast<int32_t>(requantize_shifts.at(0));
+  const cmsis_nn_quant_params quant_params{
+      &multiplier, &shift, /*is_per_channel=*/0};
+  return quantized_linear_impl(
+      "quantized_linear_out",
+      context,
+      input,
+      weights,
+      bias,
+      kernel_sum,
+      input_offset,
+      filter_offset,
+      output_offset,
+      activation_max,
+      activation_min,
+      quant_params,
+      out);
+}
+
+// cppcheck-suppress unusedFunction
+Tensor& quantized_linear_per_channel_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    const Tensor& weights,
+    const std::optional<Tensor>& bias,
+    const std::optional<Tensor>& kernel_sum,
+    const int64_t input_offset,
+    const int64_t filter_offset,
+    const int64_t output_offset,
+    const Tensor& requantize_multipliers,
+    const Tensor& requantize_shifts,
+    const int64_t activation_max,
+    const int64_t activation_min,
+    Tensor& out) {
+  const int64_t out_features = out.size(out.dim() - 1);
+  if (requantize_multipliers.scalar_type() != ScalarType::Int ||
+      requantize_shifts.scalar_type() != ScalarType::Int ||
+      requantize_multipliers.numel() != out_features ||
+      requantize_shifts.numel() != out_features) {
+    ET_LOG(
+        Error,
+        "quantized_linear_per_channel_out: requantization parameters must be int32 with one per output feature (%d)",
+        static_cast<int>(out_features));
+    context.fail(Error::InvalidArgument);
+    return out;
+  }
+  // CMSIS-NN's MVE per-channel kernel applies one multiplier to four rows
+  // when filter_offset is non-zero.
+  if (filter_offset != 0) {
+    ET_LOG(
+        Error,
+        "quantized_linear_per_channel_out: filter_offset must be 0 (symmetric weights)");
+    context.fail(Error::InvalidArgument);
+    return out;
+  }
+  const cmsis_nn_quant_params quant_params{
+      requantize_multipliers.mutable_data_ptr<int32_t>(),
+      requantize_shifts.mutable_data_ptr<int32_t>(),
+      /*is_per_channel=*/1};
+  return quantized_linear_impl(
+      "quantized_linear_per_channel_out",
+      context,
+      input,
+      weights,
+      bias,
+      kernel_sum,
+      input_offset,
+      filter_offset,
+      output_offset,
+      activation_max,
+      activation_min,
+      quant_params,
+      out);
+}
 
 } // namespace native
 } // namespace cortex_m

backends/cortex_m/ops/operators.py

@@ -581,15 +581,32 @@ def quantized_linear_impl(
         output_reshaped = output.reshape(output_shape)
 
     output = requantize_cmsis(
-        output_reshaped,
-        int(requantize_multipliers[0]),
-        int(requantize_shifts[0]),
+        output_reshaped, requantize_multipliers, requantize_shifts
     )
     output += output_offset
     output = torch.clamp(output, activation_min, activation_max).to(torch.int8)
     return output
 
 
+# Same as quantized_linear, with one int32 requantization multiplier and shift
+# per output feature.
+_LINEAR_PER_CHANNEL_ARGS = (
+    "Tensor input, Tensor weights, Tensor? bias, Tensor? kernel_sum, "
+    "int input_offset, int filter_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_max, int activation_min"
+)
+lib.define(f"quantized_linear_per_channel({_LINEAR_PER_CHANNEL_ARGS}) -> Tensor")
+lib.define(
+    f"quantized_linear_per_channel.out({_LINEAR_PER_CHANNEL_ARGS}, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+register_fake("cortex_m::quantized_linear_per_channel")(quantized_linear_meta)
+impl(lib, "quantized_linear_per_channel", "CompositeExplicitAutograd")(
+    quantized_linear_impl
+)
+
+
 # ===================================================================
 # SOFTMAX OPERATION DEFINITION
 # ===================================================================

backends/cortex_m/ops/operators.yaml

@@ -59,6 +59,12 @@
     - arg_meta: null
       kernel_name: cortex_m::quantized_linear_out
 
+- func: cortex_m::quantized_linear_per_channel.out(Tensor input, Tensor weights, Tensor? bias, Tensor? kernel_sum, int input_offset, int filter_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_max, int activation_min, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: cortex_m::quantized_linear_per_channel_out
+
 - func: cortex_m::softmax.out(Tensor input, int dim, int input_zero_point, int output_zero_point, int input_multiplier, int input_shift, int diff_min, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -360,21 +360,32 @@ def _get_linear_replacement(
 
     input_scale = node.meta["input_qparams"][0].scale
     input_zp = node.meta["input_qparams"][0].zp
-    weight_scale = node.meta["input_qparams"][1].scale
-    weight_zp = node.meta["input_qparams"][1].zp
+    weight_qparams = node.meta["input_qparams"][1]
     output_scale = node.meta["output_qparams"][0].scale
     output_zp = node.meta["output_qparams"][0].zp
     output_min = node.meta["output_qparams"][0].qmin
     output_max = node.meta["output_qparams"][0].qmax
 
-    if weight_zp != 0:
+    # Per-channel weights carry one scale and zero point per output feature.
+    if weight_qparams.per_channel:
+        assert (
+            weight_qparams.axis == 0
+        ), f"Per-channel linear weights must be quantized on axis 0; got {weight_qparams.axis}"
+        weight_scales = weight_qparams.get_scale_per_channel()
+        weight_zps = weight_qparams.get_zp_per_channel()
+    else:
+        weight_scales, weight_zps = [weight_qparams.scale], [weight_qparams.zp]
+    if any(zp != 0 for zp in weight_zps):
         raise NotImplementedError(
             f"cortex_m::quantized_linear assumes symmetric weight "
-            f"quantization (weight_zp == 0); got weight_zp={weight_zp}"
+            f"quantization (weight_zp == 0); got weight_zp={weight_qparams.zp}"
         )
 
-    quantized_multiplier, quantized_shift = quantize_multiplier_aot(
-        (input_scale * weight_scale) / output_scale
+    quantized_multipliers, quantized_shifts = zip(
+        *(
+            quantize_multiplier_aot((input_scale * weight_scale) / output_scale)
+            for weight_scale in weight_scales
+        )
     )
 
     # CMSIS-NN's MVE `arm_fully_connected_s8` path reads a precomputed
@@ -393,7 +404,7 @@ def _get_linear_replacement(
 
     if target_config.backend == cmsis_nn.Backend.MVE:
         kernel_sum_tensor = _compute_kernel_sum(
-            weights_tensor, bias_tensor, -input_zp, -weight_zp
+            weights_tensor, bias_tensor, -input_zp, -weight_zps[0]
         )
         with node.graph.inserting_after(weights):
             kernel_sum_arg = create_constant_placeholder(
@@ -408,21 +419,45 @@ def _get_linear_replacement(
         kernel_sum_arg = None
         bias_arg = bias_node
 
+    multipliers_arg: Argument
+    shifts_arg: Argument
+    if len(weight_scales) == 1:
+        target = exir_ops.edge.cortex_m.quantized_linear.default
+        multipliers_arg = list(quantized_multipliers)
+        shifts_arg = list(quantized_shifts)
+    else:
+        target = exir_ops.edge.cortex_m.quantized_linear_per_channel.default
+        with node.graph.inserting_after(weights):
+            multipliers_arg = create_constant_placeholder(
+                exported_program,
+                node.graph,
+                node.name + "_quantized_multiplier",
+                InputKind.PARAMETER,
+                torch.tensor(quantized_multipliers, dtype=torch.int32),
+            )
+            shifts_arg = create_constant_placeholder(
+                exported_program,
+                node.graph,
+                node.name + "_quantized_shift",
+                InputKind.PARAMETER,
+                torch.tensor(quantized_shifts, dtype=torch.int32),
+            )
+
     args = (
         linear_args[0],
         weights,
         bias_arg,
         kernel_sum_arg,
         -input_zp,
-        -weight_zp,
+        -weight_zps[0],
         output_zp,
-        [quantized_multiplier],
-        [quantized_shift],
+        multipliers_arg,
+        shifts_arg,
         output_max,
         output_min,
     )
 
-    return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_linear.default, args)
+    return DialectNodeSpec(target, args)
 
 
 @AtenToCortexMPass.register_dialect_substitution(

backends/cortex_m/quantizer/pattern_checkers.py

@@ -149,9 +149,19 @@ def check_quantization_config(
         cls, pattern: list[Node], quantization_config: QuantizationConfig
     ) -> bool:
         """
-        Checks that the quantization config uses per-tensor int8 quantization.
+        Checks that the quantization config uses int8 activations, and that
+        per-channel weights are quantized per output feature (axis 0).
         """
         is_int8 = cls.is_int8_activations(quantization_config)
+        weight_qspec = quantization_config.get_weight_qspec(
+            pattern[0] if pattern else None
+        )
+        if cls.is_per_channel(weight_qspec):
+            return (
+                is_int8
+                and isinstance(weight_qspec, QuantizationSpec)
+                and weight_qspec.ch_axis == 0
+            )
         return is_int8
 
 

backends/cortex_m/test/build_test_runner.sh

@@ -67,6 +67,7 @@ ops_list=(
     cortex_m::minimum.out
     cortex_m::maximum.out
     cortex_m::quantized_linear.out
+    cortex_m::quantized_linear_per_channel.out
     cortex_m::softmax.out
     cortex_m::transpose.out
     cortex_m::pad.out