Skip to content

Commit 7628aa6

Browse files
rascaniclaude
andauthored
Cortex-M backend: lower quantized GELU via the activation LUT (pytorch#21025)
### Summary 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. ### Test plan 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). Authored with Claude Code. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 53180df commit 7628aa6

4 files changed

Lines changed: 45 additions & 0 deletions

File tree

backends/cortex_m/passes/aten_to_cortex_m_pass.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,7 @@ def _has_qparams(node: Node) -> bool:
304304
@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.sigmoid.default)
305305
@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.tanh.default)
306306
@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.silu.default)
307+
@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.gelu.default)
307308
def _get_activation_replacement(
308309
node: Node, dialect_pass: AtenToDialectPass
309310
) -> DialectNodeSpec | None:

backends/cortex_m/passes/passes_utils.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,10 +204,15 @@ def _stable_silu(x: float) -> float:
204204
return x * _stable_sigmoid(x)
205205

206206

207+
def _gelu(x: float) -> float:
208+
return 0.5 * x * (1.0 + math.erf(x / math.sqrt(2.0)))
209+
210+
207211
_ACTIVATION_FNS = {
208212
exir_ops.edge.aten.sigmoid.default: _stable_sigmoid,
209213
exir_ops.edge.aten.tanh.default: math.tanh,
210214
exir_ops.edge.aten.silu.default: _stable_silu,
215+
exir_ops.edge.aten.gelu.default: _gelu,
211216
}
212217

213218

backends/cortex_m/quantizer/quantizer_support.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@
124124
(torch.ops.aten.sigmoid.default,): CortexMActivationCheck,
125125
(torch.ops.aten.tanh.default,): CortexMActivationCheck,
126126
(torch.ops.aten.silu.default,): CortexMActivationCheck,
127+
(torch.ops.aten.gelu.default,): CortexMActivationCheck,
127128
}
128129

129130
POOL_OP_PATTERNS = {

backends/cortex_m/test/ops/test_activation_quant.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,24 @@ def forward(self, x):
6161
return torch.nn.functional.silu(x)
6262

6363

64+
class _GELU(torch.nn.Module):
65+
ops_before_transforms = {
66+
**_OPS_BEFORE,
67+
"executorch_exir_dialects_edge__ops_aten_gelu_default": 1,
68+
}
69+
ops_after_transforms = {
70+
**_OPS_AFTER,
71+
"executorch_exir_dialects_edge__ops_aten_gelu_default": 0,
72+
}
73+
74+
def __init__(self):
75+
super().__init__()
76+
self.gelu = torch.nn.GELU() # default: exact / erf
77+
78+
def forward(self, x):
79+
return self.gelu(x)
80+
81+
6482
import torch as _torch
6583

6684

@@ -133,6 +151,26 @@ def _zero_input(shape):
133151
model=_SiLU(),
134152
example_inputs=(_zero_input((16,)),),
135153
),
154+
"gelu_rank1": McuTestCase(
155+
model=_GELU(),
156+
example_inputs=(ramp_tensor(-6, 6, (16,)),),
157+
),
158+
"gelu_rank4": McuTestCase(
159+
model=_GELU(),
160+
example_inputs=(ramp_tensor(-4, 4, (1, 8, 4, 4)),),
161+
),
162+
"gelu_saturating": McuTestCase(
163+
model=_GELU(),
164+
example_inputs=(ramp_tensor(-50, 50, (32,)),),
165+
),
166+
"gelu_asymmetric_zp": McuTestCase(
167+
model=_GELU(),
168+
example_inputs=(ramp_tensor(-1, 9, (16,)),),
169+
),
170+
"gelu_zero": McuTestCase(
171+
model=_GELU(),
172+
example_inputs=(_zero_input((16,)),),
173+
),
136174
}
137175

138176

0 commit comments

Comments
 (0)