diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index 2a76a558329..c028bfde7a4 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -52,6 +52,7 @@ exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter, # noqa F405 exir_ops.edge.aten.prelu.default: PReLUConverter, # noqa F405 exir_ops.edge.aten.relu.default: ReLUConverter, # noqa F405 + exir_ops.edge.aten.rsqrt.default: RsqrtConverter, # noqa F405 exir_ops.edge.aten.sigmoid.default: SigmoidConverter, # noqa F405 exir_ops.edge.aten.slice_copy.Tensor: SliceTensorConverter, # noqa F405 exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py index ece9d5aa672..76eedb59c72 100755 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py @@ -80,6 +80,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.relu_converter import ( ReLUConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.rsqrt_converter import ( + RsqrtConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.sigmoid_converter import ( SigmoidConverter, ) @@ -137,6 +140,7 @@ "QDQPerTensorDequantizeConverter", "QDQQuantizeConverter", "ReLUConverter", + "RsqrtConverter", "SigmoidConverter", "SliceTensorConverter", "SoftmaxConverter", diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/rsqrt_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/rsqrt_converter.py new file mode 100644 index 00000000000..ff041652bbb --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/rsqrt_converter.py @@ -0,0 +1,62 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, + NodeConverter, +) +from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( + BuiltinOperator, +) + +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from torch.fx import Node +from torch.nn import Parameter + + +class RsqrtConverter(NodeConverter): + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + return True + + @staticmethod + def _is_supported_on_target( + node: Node, + neutron_target_spec: NeutronTargetSpec, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if not NodeConverter.uses_quantization_type_for_io( + node, + supported_types=[torch.int8, torch.uint8], + input_indices=[0], + output_indices=[0], + ): + return False + + return True + + def convert(self, node: Node): + """Convert the `aten.rsqrt.default` node to NeutronIR `RSQRT` operator. + The ExecuTorch schema is: + rsqrt( + Tensor self + ) -> Tensor + """ + self.assert_convertible(node) + + t_op = self._create_tflite_op_with_io_tensors(node) + t_op.opcode_index = self.builder.op_code_index_for_op_type( + BuiltinOperator.RSQRT + ) + + self.builder.append_operators([t_op]) diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 9262d1a5814..b55016e551d 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -48,6 +48,7 @@ ReluInPlacePattern, ReluPattern, ReshapePattern, + RsqrtPattern, SharedSpecPattern, SigmoidPattern, SliceTensorPattern, @@ -294,6 +295,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) OpQuantizer(ReluPattern(is_qat=is_qat), static_qconfig), OpQuantizer(ReluInPlacePattern(is_qat=is_qat), static_qconfig), OpQuantizer(ReshapePattern(is_qat=is_qat), static_qconfig), + OpQuantizer(RsqrtPattern(is_qat=is_qat), static_qconfig), OpQuantizer(SigmoidPattern(is_qat=is_qat), static_qconfig), OpQuantizer(SliceTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(SoftMaxPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index 5a4d86b65d7..a38bbd2ebc0 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -1060,6 +1060,13 @@ def partition_types(self): return [torch.ops.aten.reshape.default] +class RsqrtPattern(SingleInputBasicPattern): + """Quantizer for the `aten.rsqrt.default` operator.""" + + def partition_types(self): + return [torch.ops.aten.rsqrt.default] + + class ViewPattern(SharedSpecPattern): """ Quantizer for View operator. diff --git a/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py new file mode 100644 index 00000000000..67101410d9d --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py @@ -0,0 +1,81 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np + +# noinspection PyUnusedImports +import pytest +import torch + +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + AllCloseOutputComparator, +) +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import Rsqrt +from executorch.backends.nxp.tests.use_qat import * # noqa F403 + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(23) + np.random.seed(23) + + +class RsqrtModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.rsqrt(x) + + +class TestRsqrt: + def assert_delegated(self, model, input_shape, mocker, request, use_qat=False): + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={Rsqrt: 1}, + expected_non_delegated_ops={}, + ) + + # Use positive-only values because rsqrt is only defined for x > 0. + dataset_creator = RandomDatasetCreator(low=0.1, high=2.0) + + # Allow a single quantization bit error in the output. + comparator = AllCloseOutputComparator(atol=1) + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + dataset_creator, + comparator, + use_qat=use_qat, + ) + + def test__basic_nsys_inference(self, mocker, request): + input_shape = (2, 13, 7, 9) + model = RsqrtModule() + self.assert_delegated(model, input_shape, mocker, request) + + def test__basic_nsys_inference__qat(self, mocker, request, use_qat): + input_shape = (3, 5, 7, 11) + model = RsqrtModule() + self.assert_delegated(model, input_shape, mocker, request, use_qat=use_qat) + + @pytest.mark.parametrize( + "input_shape", + [ + pytest.param((2,), id="1D"), + pytest.param((2, 3), id="2D"), + pytest.param((2, 3, 5), id="3D"), + pytest.param((2, 3, 5, 7), id="4D"), + ], + ) + def test__input_shapes(self, mocker, request, input_shape): + model = RsqrtModule() + self.assert_delegated(model, input_shape, mocker, request) diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index 17833edb14b..bca7de24f58 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -45,6 +45,7 @@ QuantizePerChannel = exir_ops.edge.quantized_decomposed.quantize_per_channel.default QuantizePerTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default Relu = exir_ops.edge.aten.relu.default +Rsqrt = exir_ops.edge.aten.rsqrt.default Sigmoid = exir_ops.edge.aten.sigmoid.default Slice = exir_ops.edge.aten.slice.Tensor SliceCopy = exir_ops.edge.aten.slice_copy.Tensor