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,123 @@
+/*
+ * 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>();
+
+  constexpr int32_t kActivationMin = std::numeric_limits<T>::min();
+  constexpr int32_t kActivationMax = 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;
+
+    const float quotient = (denominator != 0)
+        ? static_cast<float>(numerator) / static_cast<float>(denominator)
+        : 0.0f;
+
+    int32_t result =
+        static_cast<int32_t>(std::round(quotient * effective_scale)) + out_zp;
+    result = std::max(kActivationMin, std::min(kActivationMax, result));
+    out_ptr[i] = static_cast<T>(result);
+  }
+}
+
+} // namespace
+
+using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
+
+// CMSIS-NN has no integer elementwise-division primitive, so the quotient is
+// evaluated in float. The effective scale (scale_in1 / (scale_in2 * scale_out))
+// is carried in the AoT-computed output_multiplier/output_shift and
+// reconstructed here, mirroring the softmax kernel's fixed-point-to-float
+// reconstruction. 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 int64_t output_multiplier,
+    const int64_t output_shift,
+    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);
+
+  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,
+      output_multiplier,
+      output_shift);
+
+  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 = std::ldexp(
+      static_cast<float>(output_multiplier) / static_cast<float>(1LL << 31),
+      static_cast<int>(output_shift));
+
+  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,77 @@ 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, int output_multiplier, int output_shift) -> Tensor"
+)
+lib.define(
+    "quantized_div.out("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, int output_multiplier, int output_shift, "
+    "*, 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_multiplier: int,
+    output_shift: int,
+) -> 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_multiplier: int,
+    output_shift: int,
+) -> torch.Tensor:
+    # Mirror the CMSIS-NN kernel: the quotient of the zero-point-corrected int8
+    # or int16 operands is evaluated in float and rescaled by the effective
+    # scale, which the AoT pass folds into output_multiplier/output_shift.
+    # Reconstruct that float scale the same way the softmax kernel does.
+    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)
+
+    effective_scale = math.ldexp(
+        float(output_multiplier) / float(1 << 31), output_shift
+    )
+    quotient = torch.where(other_fp != 0, self_fp / other_fp, torch.zeros_like(self_fp))
+    result = torch.round(quotient * effective_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, int output_multiplier, int output_shift, *, 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
+
+    output_mult, output_shift = quantize_multiplier_aot(
+        scale1 / (scale2 * output_scale)
+    )
+    args = (
+        node.args[0],
+        zero_point1,
+        node.args[1],
+        zero_point2,
+        output_zero_point,
+        output_mult,
+        output_shift,
+    )
+    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,12 @@ def mark_node_as_annotated(
 
 class CortexMQuantizer(ComposableQuantizer):
 
-    def __init__(self) -> None:
+    def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> None:
+        # Per-tensor activation config used for the "global" (non-conv) ops such
+        # as elementwise div/mul/add. Defaults to int8; pass INT16_PER_TENSOR_CONFIG
+        # to quantize supported ops (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 +73,7 @@ def __init__(self) -> None:
                 pattern_matcher=pattern_matcher,
             ),
             PatternQuantizer(
-                INT8_PER_TENSOR_CONFIG,
+                per_tensor_config,
                 node_finder=GlobalNodeFinder(),
                 pattern_matcher=pattern_matcher,
             ),

backends/cortex_m/quantizer/quantizer_support.py

@@ -11,6 +11,7 @@
     CortexMBmmCheck,
     CortexMConv2DCheck,
     CortexMConvTranspose2DCheck,
+    CortexMDivCheck,
     CortexMLinearCheck,
     CortexMMaxPool2DCheck,
     CortexMSoftmaxCheck,
@@ -29,6 +30,8 @@
     (torch.ops.aten.mul_.Tensor,): CortexMAddMulCheck,
     (torch.ops.aten.hardswish.default,): CortexMAddMulCheck,  # lowers to mul
     (torch.ops.aten.hardswish_.default,): CortexMAddMulCheck,  # lowers to mul
+    (torch.ops.aten.div.Tensor,): CortexMDivCheck,
+    (torch.ops.aten.div_.Tensor,): CortexMDivCheck,
 }
 
 LINEAR_OP_PATTERNS = {