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