Per-channel int8 quantized linear (pytorch#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,12 @@ namespace cortex_m {
 namespace native {
 using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
 
-// cppcheck-suppress unusedFunction
-Tensor& quantized_linear_out(
+namespace {
+// Runs `fully_connected`, a CMSIS-NN fully connected call bound to its
+// requantization parameters, on the op's tensors.
+template <typename FullyConnected>
+Tensor& quantized_linear_impl(
+    const char* op_name,
     KernelRuntimeContext& context,
     const Tensor& input,
     const Tensor& weights,
@@ -23,12 +27,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,
+    FullyConnected fully_connected,
     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 +53,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 +76,9 @@ 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 = fully_connected(
       &ctx,
       &fc_params,
-      &per_tensor_quant_params,
       &input_dims,
       input_data,
       &filter_dims,
@@ -92,16 +89,99 @@ 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) {
+  const cmsis_nn_per_tensor_quant_params quant_params{
+      static_cast<int32_t>(requantize_multipliers.at(0)),
+      static_cast<int32_t>(requantize_shifts.at(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](auto ctx, auto fc_params, auto... args) {
+        return arm_fully_connected_s8(ctx, fc_params, &quant_params, args...);
+      },
+      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;
+  }
+  const cmsis_nn_per_channel_quant_params quant_params{
+      requantize_multipliers.mutable_data_ptr<int32_t>(),
+      requantize_shifts.mutable_data_ptr<int32_t>()};
+  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](auto ctx, auto fc_params, auto... args) {
+        return arm_fully_connected_per_channel_s8(
+            ctx, fc_params, &quant_params, args...);
+      },
+      out);
+}
 
 } // namespace native
 } // namespace cortex_m

backends/cortex_m/ops/operators.py

@@ -580,16 +580,46 @@ def quantized_linear_impl(
         output_shape = (*input.shape[:-1], output.shape[-1])
         output_reshaped = output.reshape(output_shape)
 
-    output = requantize_cmsis(
-        output_reshaped,
-        int(requantize_multipliers[0]),
-        int(requantize_shifts[0]),
-    )
+    if len(requantize_multipliers) == 1:
+        output = requantize_cmsis(
+            output_reshaped,
+            int(requantize_multipliers[0]),
+            int(requantize_shifts[0]),
+        )
+    else:
+        output = torch.stack(
+            [
+                requantize_cmsis(output_reshaped[..., channel], int(mult), int(shift))
+                for channel, (mult, shift) in enumerate(
+                    zip(requantize_multipliers, requantize_shifts)
+                )
+            ],
+            dim=-1,
+        )
     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,29 @@ 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:
+        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
@@ -386,14 +394,23 @@ def _get_linear_replacement(
     linear_args = node.args
     weights = cast(Node, linear_args[1])
     weights_tensor = get_param_tensor(exported_program, weights)
+    if weights_tensor is None:
+        raise RuntimeError(
+            f"Expected linear weight parameter tensor for node {node.name}."
+        )
+    if len(weight_scales) not in (1, weights_tensor.shape[0]):
+        raise ValueError(
+            f"cortex_m::quantized_linear expects one weight scale per output "
+            f"feature ({weights_tensor.shape[0]}); got {len(weight_scales)}"
+        )
     bias_node = cast(Node | None, linear_args[2]) if len(linear_args) > 2 else None
     bias_tensor = (
         get_param_tensor(exported_program, bias_node) if bias_node is not None else None
     )
 
     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 +425,44 @@ 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, shifts_arg = (
+                create_constant_placeholder(
+                    exported_program,
+                    node.graph,
+                    node.name + suffix,
+                    InputKind.PARAMETER,
+                    torch.tensor(values, dtype=torch.int32),
+                )
+                for suffix, values in (
+                    ("_quantized_multiplier", quantized_multipliers),
+                    ("_quantized_shift", quantized_shifts),
+                )
+            )
+
     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 in (0, None)
+            )
         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