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