From 3dd9bb300033fdd7f66d0d2897d10f681389b135 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 14 Aug 2026 16:04:20 -0700 Subject: [PATCH] Cortex-M: lower transcendental unary operators through the activation LUT cortex_m::quantized_activation is a 256-entry int8 table with the qparams folded in. Only sigmoid, tanh, silu and gelu were registered; this adds log, log2, log10, log1p, sqrt and rsqrt, functional and in-place. Not exp: its codomain is unbounded, so an output scale covering the calibrated maximum leaves almost no resolution elsewhere -- over [-8, 8] the table holds two distinct levels across x in [-2, 2]. Poles saturate, and the undefined region beyond a pole continues the boundary value, so log below zero reads -128 and rsqrt below zero reads 127. Emitting the output zero point there would sit a mid-range value below the rail and break monotonicity. That case is reachable rather than theoretical: a shared quantization spec can widen an operand's grid across zero even where the tensor never goes. The tables are evaluated through torch rather than math, which raises where these functions are undefined instead of returning the -inf or nan the table needs. Silero VAD's magnitude sqrt now lowers, so its expected counts move with this. Authored with Claude Code. --- backends/cortex_m/passes/BUCK | 1 + .../cortex_m/passes/aten_to_cortex_m_pass.py | 8 +- backends/cortex_m/passes/passes_utils.py | 60 ++++++- .../cortex_m/quantizer/pattern_checkers.py | 4 +- .../cortex_m/quantizer/quantizer_support.py | 12 ++ .../cortex_m/test/models/test_silero_vad.py | 12 +- .../test/ops/test_activation_quant.py | 78 ++++++++- backends/cortex_m/test/targets.bzl | 14 ++ backends/cortex_m/test/test_activation_lut.py | 155 ++++++++++++++++++ 9 files changed, 328 insertions(+), 16 deletions(-) create mode 100644 backends/cortex_m/test/test_activation_lut.py diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index d301e14823c..f33ddbf9cf3 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -65,5 +65,6 @@ fbcode_target(_kind = runtime.python_library, ], deps=[ "fbcode//caffe2:torch", + "//executorch/exir/dialects:lib", ], ) 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 08484ce3b38..d125dd1a895 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -305,10 +305,16 @@ def _has_qparams(node: Node) -> bool: @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.tanh.default) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.silu.default) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.gelu.default) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.log.default) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.log2.default) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.log10.default) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.log1p.default) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.sqrt.default) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.rsqrt.default) def _get_activation_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: - """Lower a standalone quantized sigmoid / tanh / silu to a single + """Lower a standalone quantized unary activation to a single cortex_m.quantized_activation call backed by an AoT-built 256-entry int8 LUT. The kernel is shape-agnostic; the LUT encodes both the activation function and the input/output qparams. diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py index bcb828c5928..dc8d10acbfe 100644 --- a/backends/cortex_m/passes/passes_utils.py +++ b/backends/cortex_m/passes/passes_utils.py @@ -6,11 +6,12 @@ # LICENSE file in the root directory of this source tree. import math -from typing import Any +from typing import Any, Callable import torch from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload from torch.fx import Node @@ -208,11 +209,34 @@ def _gelu(x: float) -> float: return 0.5 * x * (1.0 + math.erf(x / math.sqrt(2.0))) -_ACTIVATION_FNS = { +def _via_torch(fn: Callable[[torch.Tensor], torch.Tensor]) -> Callable[[float], float]: + """Evaluate a torch unary at one point, in double precision.""" + + def evaluate(x: float) -> float: + return fn(torch.tensor(x, dtype=torch.float64)).item() + + return evaluate + + +_ACTIVATION_FNS: dict[EdgeOpOverload, Callable[[float], float]] = { exir_ops.edge.aten.sigmoid.default: _stable_sigmoid, exir_ops.edge.aten.tanh.default: math.tanh, exir_ops.edge.aten.silu.default: _stable_silu, exir_ops.edge.aten.gelu.default: _gelu, + # Only functions with no cheap closed form in the quantized domain belong + # here; anything expressible as a rescale should not spend a table. exp is + # deliberately absent despite qualifying: its codomain is unbounded, so an + # int8 output scale leaves it almost no resolution. + # + # Evaluated through torch rather than math so the table inherits IEEE + # semantics at the edges of each domain: math.log(0) raises where torch + # returns the -inf that saturates. + exir_ops.edge.aten.log.default: _via_torch(torch.log), + exir_ops.edge.aten.log2.default: _via_torch(torch.log2), + exir_ops.edge.aten.log10.default: _via_torch(torch.log10), + exir_ops.edge.aten.log1p.default: _via_torch(torch.log1p), + exir_ops.edge.aten.sqrt.default: _via_torch(torch.sqrt), + exir_ops.edge.aten.rsqrt.default: _via_torch(torch.rsqrt), } @@ -247,12 +271,38 @@ def build_activation_lut( f"(supported: {sorted(t.__name__ for t in _ACTIVATION_FNS)})" ) f = _ACTIVATION_FNS[target] - lut = torch.empty(256, dtype=torch.int8) + defined: dict[int, int] = {} + undefined: list[int] = [] for q in range(-128, 128): x = (q - input_zp) * input_scale y = f(x) - q_out = _round_half_away_from_zero(y / output_scale + output_zp) - lut[q + 128] = max(-128, min(127, q_out)) + scaled = y / output_scale + output_zp if math.isfinite(y) else y + if math.isnan(scaled): + # log of a negative, rsqrt of a negative. Filled in below. + undefined.append(q + 128) + continue + if not math.isfinite(scaled): + # A pole. The rail is the closest int8 has to it. + q_out = 127 if scaled > 0 else -128 + else: + q_out = _round_half_away_from_zero(scaled) + defined[q + 128] = max(-128, min(127, q_out)) + + if not defined: + raise ValueError( + f"build_activation_lut: {target} is undefined across the whole " + f"input range (scale {input_scale}, zero point {input_zp})" + ) + + lut = torch.empty(256, dtype=torch.int8) + for index, value in defined.items(): + lut[index] = value + # Each of these functions is undefined on one side of a boundary, so an + # undefined entry continues the value at that boundary. That keeps the table + # monotone; emitting the output zero point instead would put a mid-range + # value below the pole's rail. + for index in undefined: + lut[index] = defined[min(defined, key=lambda d: abs(d - index))] return lut diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index cc89715b537..c7dec6f8bb7 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/backends/cortex_m/quantizer/pattern_checkers.py @@ -140,8 +140,8 @@ def check_quantization_config( class CortexMActivationCheck(PatternCheck): - """Accept standalone elementwise activations (sigmoid / tanh / silu) - that the LUT-based cortex_m.quantized_activation op handles uniformly. + """Accept the standalone elementwise activations that the LUT-based + cortex_m.quantized_activation op handles uniformly. The kernel is shape-agnostic and the LUT is computed AoT from per-tensor qparams, so the only thing to enforce is int8 per-tensor quantization. diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index aaaf6414d06..0a579912e70 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -131,6 +131,18 @@ (torch.ops.aten.silu.default,): CortexMActivationCheck, (torch.ops.aten.silu_.default,): CortexMActivationCheck, (torch.ops.aten.gelu.default,): CortexMActivationCheck, + (torch.ops.aten.log.default,): CortexMActivationCheck, + (torch.ops.aten.log_.default,): CortexMActivationCheck, + (torch.ops.aten.log2.default,): CortexMActivationCheck, + (torch.ops.aten.log2_.default,): CortexMActivationCheck, + (torch.ops.aten.log10.default,): CortexMActivationCheck, + (torch.ops.aten.log10_.default,): CortexMActivationCheck, + (torch.ops.aten.log1p.default,): CortexMActivationCheck, + (torch.ops.aten.log1p_.default,): CortexMActivationCheck, + (torch.ops.aten.sqrt.default,): CortexMActivationCheck, + (torch.ops.aten.sqrt_.default,): CortexMActivationCheck, + (torch.ops.aten.rsqrt.default,): CortexMActivationCheck, + (torch.ops.aten.rsqrt_.default,): CortexMActivationCheck, } POOL_OP_PATTERNS = { diff --git a/backends/cortex_m/test/models/test_silero_vad.py b/backends/cortex_m/test/models/test_silero_vad.py index 9793f94f2c6..11bef2212a4 100644 --- a/backends/cortex_m/test/models/test_silero_vad.py +++ b/backends/cortex_m/test/models/test_silero_vad.py @@ -36,11 +36,11 @@ "executorch_exir_dialects_edge__ops_aten_tanh_default": 2, "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 2, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 15, - "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 14, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 16, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 15, } -# The final `sigmoid(final_conv(x))` now lowers to cortex_m.quantized_activation. -# The 3 remaining sigmoids and 2 tanhs are LSTMCell gates: PyTorch export +# The final `sigmoid(final_conv(x))` and the STFT magnitude's sqrt now lower to +# cortex_m.quantized_activation. The 3 remaining sigmoids and 2 tanhs are LSTMCell gates: PyTorch export # captures nn.LSTMCell as a single high-level op, so the quantizer never sees # the gate activations and can't annotate them. They're decomposed only at # to_edge -- which runs after the quantizer, so by then the gates have no @@ -64,7 +64,7 @@ "executorch_exir_dialects_edge__ops_aten_sigmoid_default": 3, "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 2, "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default": 1, - "executorch_exir_dialects_edge__ops_aten_sqrt_default": 1, + "executorch_exir_dialects_edge__ops_aten_sqrt_default": 0, "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 2, "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 2, "executorch_exir_dialects_edge__ops_aten_tanh_default": 2, @@ -72,7 +72,7 @@ "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 7, "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 7, - "executorch_exir_dialects_edge__ops_cortex_m_quantized_activation_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantized_activation_default": 2, "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 1, } diff --git a/backends/cortex_m/test/ops/test_activation_quant.py b/backends/cortex_m/test/ops/test_activation_quant.py index 265c0af3512..eda3f484f89 100644 --- a/backends/cortex_m/test/ops/test_activation_quant.py +++ b/backends/cortex_m/test/ops/test_activation_quant.py @@ -145,11 +145,28 @@ def forward(self, x): return self.gelu(x) -import torch as _torch +class _Transcendental(torch.nn.Module): + """The transcendental set differs only in the function and the domain the + input has to stay inside, so one class covers all of them.""" + + def __init__(self, fn, edge_name): + super().__init__() + self.fn = fn + self.ops_before_transforms = { + **_OPS_BEFORE, + f"executorch_exir_dialects_edge__ops_aten_{edge_name}_default": 1, + } + self.ops_after_transforms = { + **_OPS_AFTER, + f"executorch_exir_dialects_edge__ops_aten_{edge_name}_default": 0, + } + + def forward(self, x): + return self.fn(x) def _zero_input(shape): - return _torch.zeros(shape, dtype=_torch.float32) + return torch.zeros(shape, dtype=torch.float32) # Wide-magnitude inputs exercise the `max(-128, min(127, q_out))` clamp inside @@ -261,6 +278,63 @@ def _zero_input(shape): model=_GELU(), example_inputs=(_zero_input((16,)),), ), + # Each of these stays inside its function's domain. What the table does + # outside it is pinned by test_activation_lut instead, since the quantized + # reference here saturates to the same rail whatever the table holds. + "log": McuTestCase( + model=_Transcendental(torch.log, "log"), + example_inputs=(ramp_tensor(0.5, 8, (16,)),), + ), + "log2": McuTestCase( + model=_Transcendental(torch.log2, "log2"), + example_inputs=(ramp_tensor(0.5, 8, (16,)),), + ), + "log10": McuTestCase( + model=_Transcendental(torch.log10, "log10"), + example_inputs=(ramp_tensor(0.5, 8, (16,)),), + ), + "log1p": McuTestCase( + model=_Transcendental(torch.log1p, "log1p"), + example_inputs=(ramp_tensor(-0.5, 8, (16,)),), + ), + "sqrt": McuTestCase( + model=_Transcendental(torch.sqrt, "sqrt"), + example_inputs=(ramp_tensor(0, 9, (16,)),), + ), + "rsqrt": McuTestCase( + model=_Transcendental(torch.rsqrt, "rsqrt"), + example_inputs=(ramp_tensor(0.5, 9, (16,)),), + ), + "sqrt_rank4": McuTestCase( + model=_Transcendental(torch.sqrt, "sqrt"), + example_inputs=(ramp_tensor(0, 9, (1, 8, 4, 4)),), + ), + # An in-place activation rewrites the placeholder, so each input range is + # chosen to keep the result inside its own function's domain. + "log_inplace": McuTestCase( + model=_Transcendental(torch.Tensor.log_, "log"), + example_inputs=lambda: (ramp_tensor(1.5, 8, (16,)),), + ), + "log2_inplace": McuTestCase( + model=_Transcendental(torch.Tensor.log2_, "log2"), + example_inputs=lambda: (ramp_tensor(1.5, 8, (16,)),), + ), + "log10_inplace": McuTestCase( + model=_Transcendental(torch.Tensor.log10_, "log10"), + example_inputs=lambda: (ramp_tensor(1.5, 8, (16,)),), + ), + "sqrt_inplace": McuTestCase( + model=_Transcendental(torch.Tensor.sqrt_, "sqrt"), + example_inputs=lambda: (ramp_tensor(0, 9, (16,)),), + ), + "rsqrt_inplace": McuTestCase( + model=_Transcendental(torch.Tensor.rsqrt_, "rsqrt"), + example_inputs=lambda: (ramp_tensor(0.5, 9, (16,)),), + ), + "log1p_inplace": McuTestCase( + model=_Transcendental(torch.Tensor.log1p_, "log1p"), + example_inputs=lambda: (ramp_tensor(-0.5, 8, (16,)),), + ), } diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index b639aaebed1..2ea3a5b3b99 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -36,6 +36,20 @@ def define_common_targets(is_fbcode = False): define_operator_test_target(op) if is_fbcode: + python_unittest( + name = "test_activation_lut", + srcs = [ + "test_activation_lut.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m/passes:passes_utils", + "//executorch/backends/cortex_m/quantizer:quantizer", + "//executorch/exir/dialects:lib", + ], + ) + python_unittest( name = "test_replace_quant_nodes", srcs = [ diff --git a/backends/cortex_m/test/test_activation_lut.py b/backends/cortex_m/test/test_activation_lut.py new file mode 100644 index 00000000000..d7973cbb433 --- /dev/null +++ b/backends/cortex_m/test/test_activation_lut.py @@ -0,0 +1,155 @@ +# 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. + +"""Edge behaviour of the activation lookup table. + +The end-to-end tests compare against a quantized reference, which clamps a pole +to the same rail whatever the table holds, so they agree there by construction. +These pin the entries directly, which is where the domain handling lives. +""" + +import unittest + +import torch +from executorch.backends.cortex_m.passes.aten_to_cortex_m_pass import ( + _get_activation_replacement, + AtenToCortexMPass, +) +from executorch.backends.cortex_m.passes.passes_utils import ( + _ACTIVATION_FNS, + _round_half_away_from_zero, + build_activation_lut, +) +from executorch.backends.cortex_m.quantizer.quantizer_support import ( + ACTIVATION_OP_PATTERNS, +) +from executorch.exir.dialects._ops import ops as exir_ops + +aten = exir_ops.edge.aten + + +def _functional_edge_target(aten_op): + """The edge op an aten activation lowers to. The in-place spelling is + functionalized onto the same one, so both map here.""" + name = aten_op._schema.name.split("::")[1].rstrip("_") + return getattr(aten, name).default + + +class ActivationLutEdgeTests(unittest.TestCase): + def _entry(self, target, x, input_scale, input_zp, output_scale, output_zp): + """The table entry an input of `x` would index.""" + q = round(x / input_scale) + input_zp + self.assertTrue(-128 <= q <= 127, f"{x} is outside the int8 range") + lut = build_activation_lut( + target, input_scale, input_zp, output_scale, output_zp + ) + return int(lut[q + 128]) + + def test_log_family_saturates_at_the_pole(self): + # Zero is always exactly representable in an affine int8 quantizer, so + # this entry exists in every log table the backend builds. + for target in (aten.log.default, aten.log2.default, aten.log10.default): + with self.subTest(target=str(target)): + self.assertEqual( + self._entry(target, 0.0, 0.01568, -128, 0.01624, 42), -128 + ) + + def test_log1p_saturates_at_its_own_pole(self): + # log1p's pole is at -1, not at 0. + self.assertEqual(self._entry(aten.log1p.default, -1.0, 0.02, 0, 0.05, 0), -128) + + def test_log1p_continues_the_last_grid_point_before_its_pole(self): + # Unlike the log family's pole at zero, -1 lands on the grid only for + # some input scales. Where it does not, the undefined side continues an + # ordinary finite entry rather than the rail. + self.assertEqual( + self._entry(aten.log1p.default, -1.2, 0.4, 0, 0.05, 0), + self._entry(aten.log1p.default, -0.8, 0.4, 0, 0.05, 0), + ) + + def test_rsqrt_saturates_at_zero(self): + self.assertEqual(self._entry(aten.rsqrt.default, 0.0, 0.02, -128, 0.05, 0), 127) + + def test_undefined_points_take_the_value_at_the_domain_boundary(self): + # log of a negative is nan in eager, and the table cannot hold one. It + # continues the pole instead of emitting the output zero point, which + # would decode to 0.0 -- a value log can legitimately return, so a wrong + # answer would be indistinguishable from a real one. + self.assertEqual(self._entry(aten.log.default, -1.0, 0.02, 0, 0.05, 7), -128) + self.assertEqual(self._entry(aten.rsqrt.default, -1.0, 0.02, 0, 0.05, 7), 127) + # sqrt has no pole: it reaches zero at the boundary, so that is what the + # undefined side continues. + self.assertEqual( + self._entry(aten.sqrt.default, -1.0, 0.02, 0, 0.05, 7), + self._entry(aten.sqrt.default, 0.0, 0.02, 0, 0.05, 7), + ) + + def test_the_table_stays_monotone_across_the_pole(self): + # The bug this replaces: an entry below the pole decoded to 0.0 and so + # read larger than the entry just above it. + lut = build_activation_lut(aten.log.default, 0.02, 0, 0.05, 7) + self.assertEqual(list(lut), sorted(lut)) + + def test_every_registry_lists_the_same_activations(self): + """An activation needs a table, a quantizer pattern and a substitution. + Missing the table raises, but missing either of the other two only makes + the op quietly stay in float, which no end-to-end test would notice for + an op nothing exercises yet.""" + # Functional and in-place spellings are matched separately by the + # quantizer, so neither stands in for the other. + quantizer = { + _functional_edge_target(pattern[0]) + for pattern in ACTIVATION_OP_PATTERNS + if len(pattern) == 1 and not pattern[0]._schema.name.endswith("_") + } + in_place = { + _functional_edge_target(pattern[0]) + for pattern in ACTIVATION_OP_PATTERNS + if len(pattern) == 1 and pattern[0]._schema.name.endswith("_") + } + # gelu is the exception: torch exposes no in-place spelling for it. + self.assertEqual(in_place, quantizer - {aten.gelu.default}) + substituted = { + target + for target, fn in AtenToCortexMPass._DIALECT_SUBSTITUTIONS.items() + if fn is _get_activation_replacement + } + self.assertEqual(quantizer, set(_ACTIVATION_FNS)) + self.assertEqual(quantizer, substituted) + + def test_in_domain_entries_match_eager(self): + cases = ( + (aten.sqrt.default, torch.sqrt, 0.02, -128), + (aten.rsqrt.default, torch.rsqrt, 0.02, -128), + (aten.log.default, torch.log, 0.02, -128), + (aten.log2.default, torch.log2, 0.02, -128), + (aten.log10.default, torch.log10, 0.02, -128), + (aten.log1p.default, torch.log1p, 0.02, 0), + (aten.sigmoid.default, torch.sigmoid, 0.05, 0), + ) + for target, eager, input_scale, input_zp in cases: + with self.subTest(target=str(target)): + output_scale, output_zp = 0.05, 3 + lut = build_activation_lut( + target, input_scale, input_zp, output_scale, output_zp + ) + for q in range(-128, 128): + x = (q - input_zp) * input_scale + y = eager(torch.tensor(x, dtype=torch.float64)).item() + if not torch.isfinite(torch.tensor(y)): + continue + expected = max( + -128, + min( + 127, + _round_half_away_from_zero(y / output_scale + output_zp), + ), + ) + self.assertEqual(int(lut[q + 128]), expected, f"{target} at x={x}") + + +if __name__ == "__main__": + unittest.main()