Cortex-M: annotate and preserve in-place activations

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

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

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

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],