Add quantized_div op (#21294)

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,
             ),