feat: added support for abs using new Neutron flow

novak-vaclav · novak-vaclav · commit 420ab94b49fe · 2026-05-06T15:38:33.000+02:00
diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py
@@ -1,11 +1,14 @@
-# Copyright 2025 NXP
+# Copyright 2025-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,
+    NeutronTargetSpec,
     NodeConverter,
 )
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
@@ -25,6 +28,25 @@ def _is_supported_in_IR(
     ) -> 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 custom_delegation_options.use_new_flow_neutron_c:
+            # Requirements specified by the new Neutron flow documentation.
+
+            supported_types = [torch.int8, torch.uint8]
+            if not NodeConverter.uses_quantization_type_for_io(
+                node, supported_types, [0], None
+            ):
+                return False
+
+        return True
+
     def convert(self, node: Node):
         """Convert 'aten::abs' operator to TFLite 'Abs'."""
         self.assert_convertible(node)
diff --git a/backends/nxp/tests/dataset_creator.py b/backends/nxp/tests/dataset_creator.py
@@ -45,8 +45,10 @@ def generate_samples(
 class RandomDatasetCreator(DatasetCreator):
     """Dataset creator that generates random input samples."""
 
-    def __init__(self, num_samples=2):
+    def __init__(self, num_samples=2, low=0.0, high=1.0):
         self._num_samples = num_samples
+        self.low = low
+        self.high = high
 
     def generate_samples(
         self, dataset_dir: str, input_spec: list[ModelInputSpec]
@@ -103,9 +105,11 @@ def _gen_samples(
                     case _:
                         raise ValueError(f"Unsupported dim_order: {spec.dim_order}")
 
-                sample_vector = rng.random(
-                    np.prod(shape), torch_type_to_numpy_type(spec.dtype)
-                ).reshape(shape)
+                sample_vector = (
+                    rng.uniform(self.low, self.high, size=np.prod(shape))
+                    .astype(torch_type_to_numpy_type(spec.dtype))
+                    .reshape(shape)
+                )
                 file_name = (
                     f"{str(spec_idx).zfill(2)}.bin"
                     if len(input_spec) > 1
diff --git a/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py
@@ -1,12 +1,11 @@
-# Copyright 2025 NXP
+# Copyright 2025-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
 import pytest
 import torch
-
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
 )
@@ -17,6 +16,13 @@
     ToChannelFirstPreprocess,
     ToChannelLastPreprocess,
 )
+from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier
+
+from executorch.backends.nxp.tests.nsys_testing import (
+    lower_run_compare,
+    RandomDatasetCreator,
+)
+from executorch.backends.nxp.tests.ops_aliases import Abs, Convolution, Relu
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
@@ -29,7 +35,7 @@ def reseed_model_per_test_run():
     np.random.seed(23)
 
 
-class ConvBlocksWithAbs(torch.nn.Module):
+class ConvBlocksWithAbsModule(torch.nn.Module):
     def __init__(self, conv_in_channels: int = 3):
         super().__init__()
         self.block1 = torch.nn.Sequential(
@@ -56,36 +62,108 @@ def forward(self, x):
         return self.block2(x)
 
 
-class Abs(torch.nn.Module):
+class AbsModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
 
     def forward(self, x):
         return x.abs()
 
 
-def test_conv_abs(mocker, use_qat, input_shape: tuple[int] = (1, 3, 112, 112)):
-    model = ConvBlocksWithAbs(conv_in_channels=input_shape[1])
+class TestAbsLegacyNeutronFlow:
+    def test_conv_abs(
+        self, mocker, use_qat, input_shape: tuple[int, ...] = (1, 3, 112, 112)
+    ):
+        model = ConvBlocksWithAbsModule(conv_in_channels=input_shape[1])
 
-    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+        converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
 
-    quantized_program = to_quantized_edge_program(
-        model, input_shape, use_qat=use_qat, use_neutron_for_format_conversion=False
-    ).exported_program()
+        quantized_program = to_quantized_edge_program(
+            model,
+            input_shape,
+            use_qat=use_qat,
+            use_neutron_for_format_conversion=False,
+            use_new_flow_neutron_c=False,
+        ).exported_program()
 
-    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
-    exported_program: ExportedProgram = converter_spy.call_args.args[1]
+        tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+        exported_program: ExportedProgram = converter_spy.call_args.args[1]
 
-    assert not graph_contains_any_of_ops(
-        graph=quantized_program.graph, ops=[exir_ops.edge.aten.abs.default]
-    )
+        assert not graph_contains_any_of_ops(
+            graph=quantized_program.graph, ops=[exir_ops.edge.aten.abs.default]
+        )
+
+        input_data = (np.random.random(input_shape) * 50).astype(np.int8)
+        convert_run_compare(
+            exported_program,
+            tfl_model=tflite_flatbuffers_model,
+            tflite_input_preprocess=ToChannelLastPreprocess(),
+            tflite_output_preprocess=ToChannelFirstPreprocess(),
+            input_data=input_data,
+            atol=1.0,
+        )
+
+
+class TestAbsNewNeutronFlow:
+    @staticmethod
+    def _get_dataset_creator():
+        # to test `abs` reliably, we need to include negative values
+        low = -255.0
+        high = 255.0
+
+        dataset = RandomDatasetCreator(low=low, high=high)
+        return dataset
+
+    def test__basic_nsys_inference(self, mocker):
+        input_shape = (2, 3, 6, 7)
+        model = AbsModule()
+        graph_verifier = DetailedGraphVerifier(
+            mocker, expected_delegated_ops={Abs: 1}, expected_non_delegated_ops={}
+        )
 
-    input_data = (np.random.random(input_shape) * 50).astype(np.int8)
-    convert_run_compare(
-        exported_program,
-        tfl_model=tflite_flatbuffers_model,
-        tflite_input_preprocess=ToChannelLastPreprocess(),
-        tflite_output_preprocess=ToChannelFirstPreprocess(),
-        input_data=input_data,
-        atol=1.0,
-    )
+        dataset_creator = self._get_dataset_creator()
+        lower_run_compare(
+            model,
+            input_shape,
+            graph_verifier,
+            dataset_creator,
+            use_new_flow_neutron_c=True,
+        )
+
+    def test__basic_nsys_inference__big(self, mocker):
+        # some operators have delegation requirement that size must be < 4096
+        input_shape = (4097, 1)
+        model = AbsModule()
+        graph_verifier = DetailedGraphVerifier(
+            mocker, expected_delegated_ops={Abs: 1}, expected_non_delegated_ops={}
+        )
+
+        dataset_creator = self._get_dataset_creator()
+        lower_run_compare(
+            model,
+            input_shape,
+            graph_verifier,
+            dataset_creator,
+            use_new_flow_neutron_c=True,
+        )
+
+    def test_basic_nsys_inference__with_conv(self, mocker):
+        input_shape = (2, 3, 6, 7)
+        in_channels = input_shape[1]
+        model = ConvBlocksWithAbsModule(conv_in_channels=in_channels)
+
+        # one `relu` ends up in the same delegated partition as `abs`
+        graph_verifier = DetailedGraphVerifier(
+            mocker,
+            expected_delegated_ops={Abs: 1, Relu: 1},
+            expected_non_delegated_ops={Relu: 1, Convolution: 2},
+        )
+
+        dataset_creator = self._get_dataset_creator()
+        lower_run_compare(
+            model,
+            input_shape,
+            graph_verifier,
+            dataset_creator,
+            use_new_flow_neutron_c=True,
+        )
diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py
@@ -11,8 +11,10 @@
 import torch
 from executorch.exir.dialects._ops import ops as exir_ops
 
+Abs = exir_ops.edge.aten.abs.default
 AvgPool2D = exir_ops.edge.aten.avg_pool2d.default
 Bmm = exir_ops.edge.aten.bmm.default
+Convolution = exir_ops.edge.aten.convolution.default
 DequantizePerChannel = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default
 DequantizePerTensor = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
 ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate
@@ -28,6 +30,7 @@
 Squeeze = exir_ops.edge.aten.squeeze.default
 SqueezeDim = exir_ops.edge.aten.squeeze.dim
 SqueezeDims = exir_ops.edge.aten.squeeze.dims
+Relu = exir_ops.edge.aten.relu.default
 Unsqueeze = exir_ops.edge.aten.unsqueeze.default
 UpsampleBilinear2D = exir_ops.edge.aten.upsample_bilinear2d.vec
 UpsampleNearest2D = exir_ops.edge.aten.upsample_nearest2d.vec
