feat: added support for aten.rsqrt using new Neutron MLIR flow

Author: novak-vaclav

backends/nxp/backend/edge_program_converter.py

@@ -56,6 +56,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

backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py

@@ -89,6 +89,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,
 )
@@ -150,6 +153,7 @@
     "QDQPerTensorDequantizeConverter",
     "QDQQuantizeConverter",
     "ReLUConverter",
+    "RsqrtConverter",
     "SigmoidConverter",
     "SliceTensorConverter",
     "SoftmaxConverter",

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

backends/nxp/neutron_partitioner.py

@@ -229,6 +229,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]):
     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

backends/nxp/quantizer/neutron_quantizer.py

@@ -52,6 +52,7 @@
     ReluInPlacePattern,
     ReluPattern,
     ReshapePattern,
+    RsqrtPattern,
     SharedSpecPattern,
     SigmoidPattern,
     SliceTensorPattern,
@@ -303,6 +304,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),

backends/nxp/quantizer/patterns.py

@@ -1089,6 +1089,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.

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)

backends/nxp/tests/ops_aliases.py

@@ -48,6 +48,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