Cortex-M: lower transcendental unary operators through the activation LUT

backends/cortex_m/passes/BUCK

@@ -65,5 +65,6 @@ fbcode_target(_kind = runtime.python_library,
     ],
     deps=[
         "fbcode//caffe2:torch",
+        "//executorch/exir/dialects:lib",
     ],
 )

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.

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
 
 

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.

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 = {

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,15 +64,15 @@
     "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,
     "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_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,
 }
 

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

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 = [