From fb3712cfa09244495c9181414209dd804531a398 Mon Sep 17 00:00:00 2001 From: Eli Amesefe Date: Mon, 27 Jul 2026 13:23:10 -0700 Subject: [PATCH] Add quantized_div op (#21294) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/21294 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 --- backends/cortex_m/CMakeLists.txt | 1 + backends/cortex_m/ops/op_quantized_div.cpp | 132 ++++++++++++++ backends/cortex_m/ops/operators.py | 65 +++++++ backends/cortex_m/ops/operators.yaml | 6 + backends/cortex_m/ops/targets.bzl | 1 + .../cortex_m/passes/aten_to_cortex_m_pass.py | 30 ++++ .../cortex_m/quantizer/pattern_checkers.py | 42 +++++ .../quantizer/quantization_configs.py | 21 +++ backends/cortex_m/quantizer/quantizer.py | 18 +- .../cortex_m/quantizer/quantizer_support.py | 3 + backends/cortex_m/test/ops/test_div.py | 170 ++++++++++++++++++ 11 files changed, 487 insertions(+), 2 deletions(-) create mode 100644 backends/cortex_m/ops/op_quantized_div.cpp create mode 100644 backends/cortex_m/test/ops/test_div.py diff --git a/backends/cortex_m/CMakeLists.txt b/backends/cortex_m/CMakeLists.txt index 85e52b5782a..2a213640015 100644 --- a/backends/cortex_m/CMakeLists.txt +++ b/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 diff --git a/backends/cortex_m/ops/op_quantized_div.cpp b/backends/cortex_m/ops/op_quantized_div.cpp new file mode 100644 index 00000000000..1702ecbef67 --- /dev/null +++ b/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 +#include + +#include "cortex_m_ops_common.h" + +namespace cortex_m { +namespace native { +namespace { + +template +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(); + const T* input2_ptr = input2.data_ptr(); + T* out_ptr = out.mutable_data_ptr(); + + // 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(std::numeric_limits::min()); + constexpr float kActivationMax = + static_cast(std::numeric_limits::max()); + + const int64_t num_elements = out.numel(); + for (int64_t i = 0; i < num_elements; ++i) { + const int32_t numerator = static_cast(input1_ptr[i]) - zp1; + const int32_t denominator = static_cast(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(numerator) / static_cast(denominator) + : 0.0f; + + const float scaled = + std::round(quotient * effective_scale) + static_cast(out_zp); + const float clamped = + std::max(kActivationMin, std::min(kActivationMax, scaled)); + out_ptr[i] = static_cast(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(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(input1_zero_point); + const int32_t zp2 = static_cast(input2_zero_point); + const int32_t out_zp = static_cast(output_zero_point); + + const float effective_scale = static_cast(output_scale); + + if (dtype == ScalarType::Char) { + quantized_div_typed( + input1, zp1, input2, zp2, out_zp, effective_scale, out); + } else { + quantized_div_typed( + input1, zp1, input2, zp2, out_zp, effective_scale, out); + } + + return out; +} + +} // namespace native +} // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 459b2115913..0c731911e44 100644 --- a/backends/cortex_m/ops/operators.py +++ b/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 # =================================================================== diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index 8f6887a19d9..2c85325f854 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/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: diff --git a/backends/cortex_m/ops/targets.bzl b/backends/cortex_m/ops/targets.bzl index 9ba1d412165..9e41b8ed9d5 100644 --- a/backends/cortex_m/ops/targets.bzl +++ b/backends/cortex_m/ops/targets.bzl @@ -58,6 +58,7 @@ OPERATORS = [ "dequantize_per_tensor", "quantized_add", "quantized_mul", + "quantized_div", "minimum", "maximum", "quantized_linear", diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index 09b0171ad0b..a3af7b5bc1b 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/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 diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index 876b0d855b4..cc89715b537 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/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): diff --git a/backends/cortex_m/quantizer/quantization_configs.py b/backends/cortex_m/quantizer/quantization_configs.py index 0f10bd6afef..48c45215dcf 100644 --- a/backends/cortex_m/quantizer/quantization_configs.py +++ b/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, diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index e331c80ed1c..d3f49114144 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/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, ), diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index 0a696eb96a1..89f631c9f83 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/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 = { diff --git a/backends/cortex_m/test/ops/test_div.py b/backends/cortex_m/test/ops/test_div.py new file mode 100644 index 00000000000..aa476bbbb67 --- /dev/null +++ b/backends/cortex_m/test/ops/test_div.py @@ -0,0 +1,170 @@ +# 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. + + +import torch +from executorch.backends.arm.test.common import parametrize, xfail_type +from executorch.backends.cortex_m.quantizer.quantization_configs import ( + INT16_PER_TENSOR_CONFIG, +) +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.cortex_m.test.tester import ( + CortexMTester, + McuTestCase, + ramp_tensor, +) +from executorch.backends.test.harness.stages import Quantize + + +class _CortexMInt16Quantize(Quantize): + """Quantize stage that drives the Cortex-M quantizer with int16 activations.""" + + def __init__(self) -> None: + super().__init__(CortexMQuantizer(per_tensor_config=INT16_PER_TENSOR_CONFIG)) + + +class CortexMSelfDiv(torch.nn.Module): + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 2, + } + + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_div_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + } + + def forward(self, x): + return x / x + + +class CortexMTensorDiv(torch.nn.Module): + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 3, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3, + } + + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_div_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + } + + def forward(self, x, y): + return x / y + + +# Divisors are kept strictly positive so the quantized denominator never lands +# on its zero point (which the kernel maps to a 0 quotient). +test_cases = { + "self_rank_1": McuTestCase( + CortexMSelfDiv(), + (ramp_tensor(1, 5, (10,)),), + ), + "self_rank_4": McuTestCase( + CortexMSelfDiv(), + (ramp_tensor(1, 5, (2, 2, 2, 2)),), + ), + "tensor_pos": McuTestCase( + CortexMTensorDiv(), + (ramp_tensor(1, 10, (8,)), ramp_tensor(1, 5, (8,))), + ), + "tensor_neg_num": McuTestCase( + CortexMTensorDiv(), + (ramp_tensor(-10, -1, (8,)), ramp_tensor(1, 5, (8,))), + ), + "tensor_rank_4": McuTestCase( + CortexMTensorDiv(), + ( + ramp_tensor(-8, 8, (2, 3, 4, 4)), + ramp_tensor(1, 5, (2, 3, 4, 4)), + ), + ), + # Divisor spans down to ~1 quantization step off its zero point, so the + # smallest denominator produces a very large quotient — exercises the + # kernel's float round/clamp path (a naive int32 cast would overflow). + "tensor_small_denom": McuTestCase( + CortexMTensorDiv(), + (ramp_tensor(1, 10, (8,)), ramp_tensor(0.05, 2.5, (8,))), + ), + "broadcast_1": McuTestCase( + CortexMTensorDiv(), + (ramp_tensor(1, 5, (1,)), ramp_tensor(1, 5, (2, 2, 2, 2))), + ), + "broadcast_2": McuTestCase( + CortexMTensorDiv(), + (ramp_tensor(1, 5, (2, 2, 2, 2)), ramp_tensor(1, 5, (1,))), + ), +} + + +xfail_cases_dialect: dict[str, xfail_type] = { + # The Cortex-M quantizer refuses to quantize divisions that require + # broadcasting, so the div op is not replaced by a cortex_m implementation. + "broadcast_1": "Broadcasting is not supported in Cortex-M backend", + "broadcast_2": "Broadcasting is not supported in Cortex-M backend", +} + + +@parametrize("test_case", test_cases, xfails=xfail_cases_dialect) +def test_dialect_div(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.test_dialect( + test_case.model.ops_before_transforms, + test_case.model.ops_after_transforms, + qtol=1, + ) + + +@parametrize( + "test_case", + test_cases, + xfails=xfail_cases_dialect, +) +def test_implementation_div(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.test_implementation(qtol=1) + + +@parametrize("test_case", test_cases, xfails=xfail_cases_dialect) +def test_dialect_div_int16(test_case, cortex_m_target): + # Same op counts as the int8 flow — the graph topology is identical; only the + # quantize/dequantize dtype (and the div kernel's clamp range) change. + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.quantize(_CortexMInt16Quantize()) + tester.export() + tester.to_edge() + tester.check_count(test_case.model.ops_before_transforms) + tester.run_passes() + tester.check_count(test_case.model.ops_after_transforms) + tester.run_method_and_compare_outputs( + inputs=tester.example_inputs, qtol=1, atol=1e-03 + ) + + +@parametrize("test_case", test_cases, xfails=xfail_cases_dialect) +def test_implementation_div_int16(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.quantize(_CortexMInt16Quantize()) + tester.export() + tester.to_edge() + tester.run_passes() + tester.to_executorch() + tester.serialize() + tester.run_method_and_compare_outputs( + inputs=tester.example_inputs, qtol=1, atol=1e-03 + )