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