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