From 55aab808ff8734c1ade1b85ae6e203cbf8f08508 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 12 Aug 2026 10:46:40 -0700 Subject: [PATCH] Cortex-M: annotate and preserve in-place activations ### Summary A model that runs its activation in place emits aten.silu_ rather than aten.silu in the pre-dispatch graph. This is not a corner case: Ultralytics gives every Conv block the same class-level nn.SiLU instance, which initialize_weights() then flips to inplace=True, so one attribute decides the whole network. ACTIVATION_OP_PATTERNS listed only the functional overloads, so the quantizer never annotated those nodes; FoldAndAnnotateQParamsPass then declined to fold them, since it folds only into nodes carrying ArmAnnotationInfo, and each activation was left as an fp32 island between two quantized convolutions. On yolo11n that is 76 of them, visible only as a count of unannotated nodes in the quantizer report. The other pattern dicts in this file already enumerate the in-place variants of relu, hardtanh, clamp and hardsigmoid, so this follows that convention rather than adding a normalization pass. aten.gelu_ is left out because no idiomatic model reaches it: there is no Tensor.gelu_, and neither nn.GELU nor F.gelu takes an inplace argument. Annotating alone is not enough on the compiler path users actually run. aot_arm_compiler.py preserved hardsigmoid and hardswish through to_edge but not silu, so silu decomposed into sigmoid * mul, the mul carried qparams on only one input, and AtenToCortexMPass raised a KeyError. That made an annotated conv+SiLU model fail to compile where it had previously produced a working fp32 island. Adding silu to that preserve list, and to the same list in the overview doc, makes it lower to quantized_conv2d + quantized_activation instead. The test harness had silu preserved already, which is why no existing test saw this. That list is Cortex-M specific, and the Arm backend maps aten.silu to a TOSA TABLE op, so preserving it does not disturb the Ethos-U path. ### Test plan pytest backends/cortex_m/test/ops/test_activation_quant.py -- 48 cases, dialect and Corstone-300, green; 24 dialect cases also green on cortex-m0plus and cortex-m7. Full backends/cortex_m/test dialect run: 407 passed. Removing the three dict entries fails the three new in-place cases; removing the preserve entry reproduces the compile failure above on conv + nn.SiLU(inplace=True). backends/cortex_m/test/models/test_yolo11.py exercises the motivating model, but it importorskips ultralytics, which is not installed in the cortex-m CI image, so it does not run in CI. Authored with assistance from Claude Code. --- backends/arm/scripts/aot_arm_compiler.py | 1 + .../cortex_m/quantizer/quantizer_support.py | 3 + .../test/ops/test_activation_quant.py | 90 +++++++++++++++++++ .../arm-cortex-m/arm-cortex-m-overview.md | 1 + 4 files changed, 95 insertions(+) diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py index 405b4a66efd..921573c68f3 100644 --- a/backends/arm/scripts/aot_arm_compiler.py +++ b/backends/arm/scripts/aot_arm_compiler.py @@ -973,6 +973,7 @@ def _to_channels_last(x): torch.ops.aten.hardsigmoid_.default, torch.ops.aten.hardswish.default, torch.ops.aten.hardswish_.default, + torch.ops.aten.silu.default, ], _check_ir_validity=False, ), diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index 89f631c9f83..aaaf6414d06 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -125,8 +125,11 @@ ACTIVATION_OP_PATTERNS = { (torch.ops.aten.sigmoid.default,): CortexMActivationCheck, + (torch.ops.aten.sigmoid_.default,): CortexMActivationCheck, (torch.ops.aten.tanh.default,): CortexMActivationCheck, + (torch.ops.aten.tanh_.default,): CortexMActivationCheck, (torch.ops.aten.silu.default,): CortexMActivationCheck, + (torch.ops.aten.silu_.default,): CortexMActivationCheck, (torch.ops.aten.gelu.default,): CortexMActivationCheck, } diff --git a/backends/cortex_m/test/ops/test_activation_quant.py b/backends/cortex_m/test/ops/test_activation_quant.py index 2d68fddbbdf..265c0af3512 100644 --- a/backends/cortex_m/test/ops/test_activation_quant.py +++ b/backends/cortex_m/test/ops/test_activation_quant.py @@ -39,6 +39,19 @@ def forward(self, x): return torch.sigmoid(x) +# nn.Sigmoid and nn.Tanh take no `inplace` argument, so the tensor method is +# the only way to reach aten.sigmoid_ / aten.tanh_ from Python. +class _SigmoidInplace(torch.nn.Module): + ops_before_transforms = { + **_OPS_BEFORE, + "executorch_exir_dialects_edge__ops_aten_sigmoid_default": 1, + } + ops_after_transforms = _OPS_AFTER + + def forward(self, x): + return x.sigmoid_() + + class _Tanh(torch.nn.Module): ops_before_transforms = { **_OPS_BEFORE, @@ -50,6 +63,17 @@ def forward(self, x): return torch.tanh(x) +class _TanhInplace(torch.nn.Module): + ops_before_transforms = { + **_OPS_BEFORE, + "executorch_exir_dialects_edge__ops_aten_tanh_default": 1, + } + ops_after_transforms = _OPS_AFTER + + def forward(self, x): + return x.tanh_() + + class _SiLU(torch.nn.Module): ops_before_transforms = { **_OPS_BEFORE, @@ -61,6 +85,48 @@ def forward(self, x): return torch.nn.functional.silu(x) +class _SiLUInplace(torch.nn.Module): + ops_before_transforms = { + **_OPS_BEFORE, + "executorch_exir_dialects_edge__ops_aten_silu_default": 1, + } + ops_after_transforms = _OPS_AFTER + + def __init__(self): + super().__init__() + self.silu = torch.nn.SiLU(inplace=True) + + def forward(self, x): + return self.silu(x) + + +class _ConvSiLUInplace(torch.nn.Module): + """The shape a real model has: the activation consumes a convolution + output, so the conv is matched first by the per-channel quantizer and the + activation only afterwards. + """ + + # No _OPS_BEFORE here: the convolution brings its own weight quant/dequant, + # so the boundary counts the other cases share do not apply. + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_silu_default": 1, + "executorch_exir_dialects_edge__ops_aten_convolution_default": 1, + } + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_activation_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 1, + "executorch_exir_dialects_edge__ops_aten_silu_default": 0, + } + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(4, 8, 3, padding=1) + self.silu = torch.nn.SiLU(inplace=True) + + def forward(self, x): + return self.silu(self.conv(x)) + + class _GELU(torch.nn.Module): ops_before_transforms = { **_OPS_BEFORE, @@ -111,6 +177,16 @@ def _zero_input(shape): model=_Sigmoid(), example_inputs=(_zero_input((16,)),), ), + # These three activate the placeholder itself, so calibration rewrites the + # input tensor. Building it per call keeps one case from feeding the next; + # within a case both sides still see the rewritten tensor, which narrows the + # compared range. That is fine here -- they exist to prove the in-place + # spelling gets annotated, and the functional siblings above already cover + # the LUT over its full range -- but do not read them as range coverage. + "sigmoid_inplace": McuTestCase( + model=_SigmoidInplace(), + example_inputs=lambda: (ramp_tensor(-4, 4, (1, 8, 4, 4)),), + ), "tanh_rank1": McuTestCase( model=_Tanh(), example_inputs=(ramp_tensor(-3, 3, (16,)),), @@ -131,6 +207,10 @@ def _zero_input(shape): model=_Tanh(), example_inputs=(_zero_input((16,)),), ), + "tanh_inplace": McuTestCase( + model=_TanhInplace(), + example_inputs=lambda: (ramp_tensor(-2, 2, (1, 8, 4, 4)),), + ), "silu_rank1": McuTestCase( model=_SiLU(), example_inputs=(ramp_tensor(-6, 6, (16,)),), @@ -151,6 +231,16 @@ def _zero_input(shape): model=_SiLU(), example_inputs=(_zero_input((16,)),), ), + "silu_inplace": McuTestCase( + model=_SiLUInplace(), + example_inputs=lambda: (ramp_tensor(-4, 4, (1, 8, 4, 4)),), + ), + "conv_silu_inplace": McuTestCase( + model=_ConvSiLUInplace(), + example_inputs=lambda: ( + ramp_tensor(-4, 4, (1, 4, 8, 8)).to(memory_format=torch.channels_last), + ), + ), "gelu_rank1": McuTestCase( model=_GELU(), example_inputs=(ramp_tensor(-6, 6, (16,)),), diff --git a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md index bf4e41a73bf..b595363b39d 100644 --- a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md +++ b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md @@ -113,6 +113,7 @@ config = EdgeCompileConfig( torch.ops.aten.hardsigmoid_.default, torch.ops.aten.hardswish.default, torch.ops.aten.hardswish_.default, + torch.ops.aten.silu.default, ], _check_ir_validity=False, _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default],