Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions backends/nxp/backend/edge_program_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -137,6 +140,7 @@
"QDQPerTensorDequantizeConverter",
"QDQQuantizeConverter",
"ReLUConverter",
"RsqrtConverter",
"SigmoidConverter",
"SliceTensorConverter",
"SoftmaxConverter",
Expand Down
Original file line number Diff line number Diff line change
@@ -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])
2 changes: 2 additions & 0 deletions backends/nxp/quantizer/neutron_quantizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
ReluInPlacePattern,
ReluPattern,
ReshapePattern,
RsqrtPattern,
SharedSpecPattern,
SigmoidPattern,
SliceTensorPattern,
Expand Down Expand Up @@ -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),
Expand Down
7 changes: 7 additions & 0 deletions backends/nxp/quantizer/patterns.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
1 change: 1 addition & 0 deletions backends/nxp/tests/ops_aliases.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading