From f9d4f418f1f22f2ed8eb4706ffdad6fea3ddb1bd Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Thu, 16 Jul 2026 13:52:27 -0700 Subject: [PATCH] Cortex-M backend: lower quantized GELU via the activation LUT Extend the existing sigmoid/tanh/silu activation-LUT path to GELU: add an exact (erf) _gelu to _ACTIVATION_FNS, register aten.gelu in the quantizer's ACTIVATION_OP_PATTERNS (the generic CortexMActivationCheck needs no change), and add aten.gelu to the quantized_activation dialect substitution. Unlike silu, aten.gelu is a core-aten op that survives to_edge unchanged, so no to_edge preservation is required. build_activation_lut only sees the op target, so the table is always the exact/erf GELU regardless of the approximate kwarg -- fine for an int8 LUT. Verified test-first (RED->GREEN) with five nn.GELU() dialect cases (ramp / saturating / asymmetric-zp / zero) whose numerics are validated against the reference (qtol=1); the full cortex_m op dialect suite still passes (177 passed). This removes the last non-normalization fp32 activation from a SAM mask decoder's two-way transformer. Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../cortex_m/passes/aten_to_cortex_m_pass.py | 1 + backends/cortex_m/passes/passes_utils.py | 5 +++ .../cortex_m/quantizer/quantizer_support.py | 1 + .../test/ops/test_activation_quant.py | 38 +++++++++++++++++++ 4 files changed, 45 insertions(+) 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 3f5a6055331..21827ec31ec 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -227,6 +227,7 @@ def _has_qparams(node: Node) -> bool: @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.sigmoid.default) @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) def _get_activation_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py index a8033662662..bcb828c5928 100644 --- a/backends/cortex_m/passes/passes_utils.py +++ b/backends/cortex_m/passes/passes_utils.py @@ -204,10 +204,15 @@ def _stable_silu(x: float) -> float: return x * _stable_sigmoid(x) +def _gelu(x: float) -> float: + return 0.5 * x * (1.0 + math.erf(x / math.sqrt(2.0))) + + _ACTIVATION_FNS = { 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, } diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index 317189a5f3e..0a696eb96a1 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -124,6 +124,7 @@ (torch.ops.aten.sigmoid.default,): CortexMActivationCheck, (torch.ops.aten.tanh.default,): CortexMActivationCheck, (torch.ops.aten.silu.default,): CortexMActivationCheck, + (torch.ops.aten.gelu.default,): CortexMActivationCheck, } POOL_OP_PATTERNS = { diff --git a/backends/cortex_m/test/ops/test_activation_quant.py b/backends/cortex_m/test/ops/test_activation_quant.py index 2cc40d7aaee..2d68fddbbdf 100644 --- a/backends/cortex_m/test/ops/test_activation_quant.py +++ b/backends/cortex_m/test/ops/test_activation_quant.py @@ -61,6 +61,24 @@ def forward(self, x): return torch.nn.functional.silu(x) +class _GELU(torch.nn.Module): + ops_before_transforms = { + **_OPS_BEFORE, + "executorch_exir_dialects_edge__ops_aten_gelu_default": 1, + } + ops_after_transforms = { + **_OPS_AFTER, + "executorch_exir_dialects_edge__ops_aten_gelu_default": 0, + } + + def __init__(self): + super().__init__() + self.gelu = torch.nn.GELU() # default: exact / erf + + def forward(self, x): + return self.gelu(x) + + import torch as _torch @@ -133,6 +151,26 @@ def _zero_input(shape): model=_SiLU(), example_inputs=(_zero_input((16,)),), ), + "gelu_rank1": McuTestCase( + model=_GELU(), + example_inputs=(ramp_tensor(-6, 6, (16,)),), + ), + "gelu_rank4": McuTestCase( + model=_GELU(), + example_inputs=(ramp_tensor(-4, 4, (1, 8, 4, 4)),), + ), + "gelu_saturating": McuTestCase( + model=_GELU(), + example_inputs=(ramp_tensor(-50, 50, (32,)),), + ), + "gelu_asymmetric_zp": McuTestCase( + model=_GELU(), + example_inputs=(ramp_tensor(-1, 9, (16,)),), + ), + "gelu_zero": McuTestCase( + model=_GELU(), + example_inputs=(_zero_input((16,)),), + ), }