Add quantized_div op (#21294)

Summary:

Adds a `quantized_div` operator to the ExecuTorch Cortex-M (CMSIS-NN) backend, giving int8 and int16 elementwise division alongside the existing `quantized_add`/`quantized_mul` ops.

CMSIS-NN has no integer elementwise-division primitive, so the quotient of the zero-point-corrected operands is evaluated in float and rescaled by the effective scale `scale_a / (scale_b * scale_out)`. Because there is no fixed-point path to feed (unlike `quantized_mul`/`quantized_add`), that scale is computed AoT and carried directly as a `float output_scale` in the op schema rather than as a `multiplier`/`shift` pair. The kernel does the rounding, zero-point add, and int8/int16 saturation in float, casting to the output type only after clamping, so a near-zero denominator cannot overflow int32. A zero denominator maps to a 0 quotient in both the kernel and the reference.

Division is not commutative, so the operand-swap trick that lets `quantized_mul` support channel broadcasting does not apply here; the op and its quantizer pattern check (`CortexMDivCheck`) currently require identically shaped per-tensor inputs and reject any broadcasting.

int8 and int16 activations are both supported. The kernel dispatches on the output dtype (`Char`/`Short`) into a templated loop, and the composite reference clamps to the dtype range. A new `INT16_PER_TENSOR_CONFIG` (symmetric int16 activations) is added, `CortexMDivCheck` accepts int8 or int16, and `CortexMQuantizer` takes an optional, documented `per_tensor_config` so callers can opt into int16.

Wiring: `operators.py` (schema + reference impl), `op_quantized_div.cpp` (kernel), `operators.yaml`/`targets.bzl`/`CMakeLists.txt` (registration and build), `aten_to_cortex_m_pass.py` (`aten.div.Tensor` lowering), `quantization_configs.py` (int16 config), `pattern_checkers.py` + `quantizer.py` (int16 support). The Buck target that runs these tests is added in the stacked diff D113492101. All changes are mirrored across the fbcode and xplat trees.

Authored with Claude Code.

Reviewed By: rascani

Differential Revision: D113440315

Author: 3l1

backends/cortex_m/CMakeLists.txt

@@ -106,6 +106,7 @@ if(EXECUTORCH_BUILD_CORTEX_M)
       ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_batch_matmul.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_conv2d.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_depthwise_conv2d.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_div.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_linear.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_max_pool2d.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_mul.cpp

backends/cortex_m/ops/op_quantized_div.cpp

@@ -0,0 +1,132 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <algorithm>
+#include <cmath>
+
+#include "cortex_m_ops_common.h"
+
+namespace cortex_m {
+namespace native {
+namespace {
+
+template <typename T>
+void quantized_div_typed(
+    const Tensor& input1,
+    const int32_t zp1,
+    const Tensor& input2,
+    const int32_t zp2,
+    const int32_t out_zp,
+    const float effective_scale,
+    Tensor& out) {
+  const T* input1_ptr = input1.data_ptr<T>();
+  const T* input2_ptr = input2.data_ptr<T>();
+  T* out_ptr = out.mutable_data_ptr<T>();
+
+  // Saturation bounds kept in float: a denominator quantized to a single step
+  // off its zero point yields a very large quotient, so rounding and clamping
+  // in float avoids overflowing int32 before the saturating cast below.
+  constexpr float kActivationMin =
+      static_cast<float>(std::numeric_limits<T>::min());
+  constexpr float kActivationMax =
+      static_cast<float>(std::numeric_limits<T>::max());
+
+  const int64_t num_elements = out.numel();
+  for (int64_t i = 0; i < num_elements; ++i) {
+    const int32_t numerator = static_cast<int32_t>(input1_ptr[i]) - zp1;
+    const int32_t denominator = static_cast<int32_t>(input2_ptr[i]) - zp2;
+
+    // A zero-point-corrected denominator of 0 has no representable reciprocal;
+    // emit a 0 quotient so the op stays total (callers keep divisors off the
+    // zero point).
+    const float quotient = (denominator != 0)
+        ? static_cast<float>(numerator) / static_cast<float>(denominator)
+        : 0.0f;
+
+    const float scaled =
+        std::round(quotient * effective_scale) + static_cast<float>(out_zp);
+    const float clamped =
+        std::max(kActivationMin, std::min(kActivationMax, scaled));
+    out_ptr[i] = static_cast<T>(clamped);
+  }
+}
+
+} // namespace
+
+using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
+
+// CMSIS-NN has no integer elementwise-division primitive, so the quotient is
+// evaluated in float. Unlike quantized_mul/add there is no fixed-point path to
+// feed, so the effective scale (scale_in1 / (scale_in2 * scale_out)) is
+// computed AoT and carried directly as a float rather than as a
+// multiplier/shift pair. Both int8 and int16 activations are supported.
+// cppcheck-suppress unusedFunction
+Tensor& quantized_div_out(
+    KernelRuntimeContext& context,
+    const Tensor& input1,
+    const int64_t input1_zero_point,
+    const Tensor& input2,
+    const int64_t input2_zero_point,
+    const int64_t output_zero_point,
+    const double output_scale,
+    Tensor& out) {
+  const ScalarType dtype = out.scalar_type();
+  if (dtype != ScalarType::Char && dtype != ScalarType::Short) {
+    ET_LOG(
+        Error,
+        "quantized_div: only int8 and int16 are supported, got %d",
+        static_cast<int>(dtype));
+    context.fail(Error::InvalidArgument);
+    return out;
+  }
+
+  // Division is not commutative, so channel broadcasting (which relies on
+  // operand swapping in quantized_mul) is unsupported: require equal shapes.
+  validate_cmsis_nn_tensor_requirements(
+      input1,
+      input2,
+      out,
+      dtype,
+      /*require_channels_last=*/false,
+      /*require_same_sizes=*/true);
+
+  // The rescale is carried entirely by effective_scale (float), so the shared
+  // validator only needs to sanity-check the three zero points; pass identity
+  // multiplier/shift for each operand.
+  const int32_t kIdentityMultiplier(/*value=*/1);
+  const int32_t kZeroShift(/*value=*/0);
+  validate_quantization_params(
+      input1_zero_point,
+      kIdentityMultiplier,
+      kZeroShift,
+      input2_zero_point,
+      kIdentityMultiplier,
+      kZeroShift,
+      output_zero_point,
+      kIdentityMultiplier,
+      kZeroShift);
+
+  const int32_t zp1 = static_cast<int32_t>(input1_zero_point);
+  const int32_t zp2 = static_cast<int32_t>(input2_zero_point);
+  const int32_t out_zp = static_cast<int32_t>(output_zero_point);
+
+  const float effective_scale = static_cast<float>(output_scale);
+
+  if (dtype == ScalarType::Char) {
+    quantized_div_typed<int8_t>(
+        input1, zp1, input2, zp2, out_zp, effective_scale, out);
+  } else {
+    quantized_div_typed<int16_t>(
+        input1, zp1, input2, zp2, out_zp, effective_scale, out);
+  }
+
+  return out;
+}
+
+} // namespace native
+} // namespace cortex_m

backends/cortex_m/ops/operators.py

@@ -264,6 +264,71 @@ def quantized_mul_impl(
     return result
 
 
+# ===================================================================
+# QUANTIZED DIV OPERATION DEFINITION
+# ===================================================================
+lib.define(
+    "quantized_div("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, float output_scale) -> Tensor"
+)
+lib.define(
+    "quantized_div.out("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, float output_scale, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_div")  # type: ignore[misc]
+def quantized_div_meta(
+    self: torch.Tensor,
+    self_zero_point: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    output_zero_point: int,
+    output_scale: float,
+) -> torch.Tensor:
+    # Division is not commutative, so broadcasting (handled via operand swaps in
+    # quantized_mul) is not supported: require identical shapes.
+    assert self.shape == other.shape, (
+        "Cortex-M quantized_div: broadcasting is not supported — "
+        f"got self.shape={self.shape}, other.shape={other.shape}"
+    )
+    return torch.empty_like(self)
+
+
+@impl(lib, "quantized_div", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_div_impl(
+    self: torch.Tensor,
+    self_zero_point: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    output_zero_point: int,
+    output_scale: float,
+) -> torch.Tensor:
+    # Mirror the kernel: the quotient of the zero-point-corrected int8/int16
+    # operands is evaluated in float and rescaled by the effective scale
+    # (scale_in1 / (scale_in2 * scale_out)) that the AoT pass carries directly.
+    assert self.shape == other.shape, (
+        "Cortex-M quantized_div: broadcasting is not supported — "
+        f"got self.shape={self.shape}, other.shape={other.shape}"
+    )
+    if self.dtype not in (torch.int8, torch.int16):
+        raise TypeError(
+            f"cortex_m.quantized_div: expected int8 or int16 inputs, got {self.dtype}"
+        )
+    self_fp = (self.to(torch.int32) - self_zero_point).to(torch.float32)
+    other_fp = (other.to(torch.int32) - other_zero_point).to(torch.float32)
+
+    quotient = torch.where(other_fp != 0, self_fp / other_fp, torch.zeros_like(self_fp))
+    result = torch.round(quotient * output_scale) + output_zero_point
+    dtype_info = torch.iinfo(self.dtype)
+    return torch.clamp(result, dtype_info.min, dtype_info.max).to(self.dtype)
+
+
 # ===================================================================
 # QUANTIZED ACTIVATION (LUT) OPERATION DEFINITION
 # ===================================================================

backends/cortex_m/ops/operators.yaml

@@ -29,6 +29,12 @@
     - arg_meta: null
       kernel_name: cortex_m::quantized_mul_out
 
+- func: cortex_m::quantized_div.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, float output_scale, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: cortex_m::quantized_div_out
+
 - func: cortex_m::quantized_activation.out(Tensor input, Tensor lut, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:

backends/cortex_m/ops/targets.bzl

@@ -58,6 +58,7 @@ OPERATORS = [
     "dequantize_per_tensor",
     "quantized_add",
     "quantized_mul",
+    "quantized_div",
     "minimum",
     "maximum",
     "quantized_linear",

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -960,6 +960,36 @@ def _get_mul_replacement(
     return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_mul.default, args)
 
 
+@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.div.Tensor)
+def _get_div_replacement(
+    node: Node, dialect_pass: AtenToDialectPass
+) -> DialectNodeSpec | None:
+    del dialect_pass
+    if not _has_qparams(node):
+        return None
+
+    scale1 = node.meta["input_qparams"][0].scale
+    zero_point1 = node.meta["input_qparams"][0].zp
+    scale2 = node.meta["input_qparams"][1].scale
+    zero_point2 = node.meta["input_qparams"][1].zp
+    output_scale = node.meta["output_qparams"][0].scale
+    output_zero_point = node.meta["output_qparams"][0].zp
+
+    # No CMSIS fixed-point div path exists, so the kernel divides in float;
+    # carry the effective scale directly rather than encoding it as a
+    # multiplier/shift pair (as quantized_mul/add do).
+    effective_scale = float(scale1 / (scale2 * output_scale))
+    args = (
+        node.args[0],
+        zero_point1,
+        node.args[1],
+        zero_point2,
+        output_zero_point,
+        effective_scale,
+    )
+    return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_div.default, args)
+
+
 @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten._softmax.default)
 def _get_softmax_replacement(
     node: Node, dialect_pass: AtenToDialectPass

backends/cortex_m/quantizer/pattern_checkers.py

@@ -56,6 +56,48 @@ def check_quantization_config(
         return is_per_tensor and is_int8
 
 
+class CortexMDivCheck(PatternCheck):
+
+    @classmethod
+    def check_pattern(cls, pattern):
+        """
+        Reject any broadcasting. Division is not commutative, so the operand
+        swapping used to support channel broadcast in add/mul does not apply;
+        only identically shaped inputs are supported.
+        """
+        for node in pattern:
+            if len(node.all_input_nodes) == 2:
+                t1 = get_first_fake_tensor(node.all_input_nodes[0])
+                t2 = get_first_fake_tensor(node.all_input_nodes[1])
+                if t1.shape != t2.shape:
+                    return False
+
+        return True
+
+    @classmethod
+    def check_quantization_config(
+        cls, pattern: list[Node], quantization_config: QuantizationConfig
+    ) -> bool:
+        """
+        Checks that the quantization config uses per-tensor int8 or int16
+        quantization (the div kernel supports both).
+        """
+        input_qspec = quantization_config.get_input_act_qspec()
+        output_qspec = quantization_config.get_output_act_qspec()
+        is_per_tensor = PatternCheck.is_per_tensor(
+            input_qspec
+        ) and PatternCheck.is_per_tensor(output_qspec)
+        allowed_dtypes = (torch.int8, torch.int16)
+        is_valid_dtype = (
+            isinstance(input_qspec, QuantizationSpec)
+            and isinstance(output_qspec, QuantizationSpec)
+            and input_qspec.dtype in allowed_dtypes
+            and output_qspec.dtype in allowed_dtypes
+            and input_qspec.dtype == output_qspec.dtype
+        )
+        return is_per_tensor and is_valid_dtype
+
+
 class CortexMConv2DCheck(PatternCheck):
     @classmethod
     def check_pattern(cls, pattern):

backends/cortex_m/quantizer/quantization_configs.py

@@ -59,6 +59,16 @@
     ch_axis=0,
 )
 
+# 16-bit activation spec. Symmetric (zero point 0) with the -32767 quant_min
+# used by the Arm a16w8 config, so the range is symmetric about zero.
+INT16_ACTIVATION_PER_TENSOR_QSPEC = QuantizationSpec(
+    dtype=torch.int16,
+    observer_or_fake_quant_ctr=MinMaxObserver,
+    qscheme=torch.per_tensor_symmetric,
+    quant_min=-32767,
+    quant_max=32767,
+)
+
 # Constants shared by Cortex-M quantized operators.
 CMSIS_SOFTMAX_SCALE: float = 1.0 / 256.0
 CMSIS_SOFTMAX_ZERO_POINT: int = -128
@@ -196,6 +206,17 @@ def get_bias_qspec(
 )
 
 
+# int16 activations (weight/bias qspecs are unused by weightless elementwise
+# ops such as quantized_div; they carry the int8/int32 defaults).
+INT16_PER_TENSOR_CONFIG = CortexMQuantizationConfig(
+    INT16_ACTIVATION_PER_TENSOR_QSPEC,
+    INT16_ACTIVATION_PER_TENSOR_QSPEC,
+    INT8_WEIGHT_PER_TENSOR_QSPEC,
+    _get_int32_bias_qspec,
+    f"{__name__}.INT16_PER_TENSOR_CONFIG",
+)
+
+
 INT8_PER_CHANNEL_CONFIG = CortexMQuantizationConfig(
     INT8_ACTIVATION_PER_TENSOR_QSPEC,
     INT8_ACTIVATION_PER_TENSOR_QSPEC,

backends/cortex_m/quantizer/quantizer.py

@@ -12,6 +12,7 @@
     PatternQuantizer,
     SharedQspecQuantizer,
 )
+from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
 from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
 from executorch.backends.cortex_m.quantizer.node_finders import (
     GlobalNodeFinder,
@@ -45,7 +46,20 @@ def mark_node_as_annotated(
 
 class CortexMQuantizer(ComposableQuantizer):
 
-    def __init__(self) -> None:
+    def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> None:
+        """Cortex-M PT2E quantizer.
+
+        Args:
+            per_tensor_config: Per-tensor activation config applied to the
+                non-conv elementwise ops (div/mul/add/...) that
+                ``GlobalNodeFinder`` matches anywhere in the graph. Convolutions
+                are always quantized with the per-channel config. Defaults to
+                ``INT8_PER_TENSOR_CONFIG``; pass ``INT16_PER_TENSOR_CONFIG`` to
+                quantize the ops that support it (e.g. ``quantized_div``) with
+                int16 activations.
+        """
+        per_tensor_config = per_tensor_config or INT8_PER_TENSOR_CONFIG
+
         conv_targets: set[OpOverload] = set()
         for key in CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys():
             conv_targets.update(key)
@@ -67,7 +81,7 @@ def __init__(self) -> None:
                 pattern_matcher=pattern_matcher,
             ),
             PatternQuantizer(
-                INT8_PER_TENSOR_CONFIG,
+                per_tensor_config,
                 node_finder=GlobalNodeFinder(),
                 pattern_matcher=pattern_matcher,
             ),