Skip to content

Commit ec37498

Browse files
committed
feat: added support for abs using new Neutron flow
1 parent 08b2987 commit ec37498

2 files changed

Lines changed: 104 additions & 25 deletions

File tree

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
1-
# Copyright 2025 NXP
1+
# Copyright 2025-2026 NXP
22
#
33
# This source code is licensed under the BSD-style license found in the
44
# LICENSE file in the root directory of this source tree.
55

66

7+
import torch
8+
79
from executorch.backends.nxp.backend.ir.converter.node_converter import (
810
CustomDelegationOptions,
11+
NeutronTargetSpec,
912
NodeConverter,
1013
)
1114
from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
@@ -25,6 +28,25 @@ def _is_supported_in_IR(
2528
) -> bool:
2629
return True
2730

31+
@staticmethod
32+
def _is_supported_on_target(
33+
node: Node,
34+
neutron_target_spec: NeutronTargetSpec,
35+
parameters_mapping: dict[str, Parameter],
36+
custom_delegation_options: CustomDelegationOptions,
37+
) -> bool:
38+
39+
if custom_delegation_options.use_new_flow_neutron_c:
40+
# Requirements specified by the new Neutron flow documentation.
41+
42+
supported_types = [torch.int8, torch.uint8]
43+
if not NodeConverter.uses_quantization_type_for_io(
44+
node, supported_types, [0], [0]
45+
):
46+
return False
47+
48+
return True
49+
2850
def convert(self, node: Node):
2951
"""Convert 'aten::abs' operator to TFLite 'Abs'."""
3052
self.assert_convertible(node)

backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py

Lines changed: 81 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
1-
# Copyright 2025 NXP
1+
# Copyright 2025-2026 NXP
22
#
33
# This source code is licensed under the BSD-style license found in the
44
# LICENSE file in the root directory of this source tree.
55

66
import numpy as np
77
import pytest
88
import torch
9-
109
from executorch.backends.nxp.backend.edge_program_converter import (
1110
EdgeProgramToIRConverter,
1211
)
@@ -17,6 +16,12 @@
1716
ToChannelFirstPreprocess,
1817
ToChannelLastPreprocess,
1918
)
19+
from executorch.backends.nxp.tests.graph_verifier import (
20+
BaseGraphVerifier,
21+
NonDelegatedNode,
22+
)
23+
24+
from executorch.backends.nxp.tests.nsys_testing import lower_run_compare
2025

2126
from executorch.exir.dialects._ops import ops as exir_ops
2227
from torch.export import ExportedProgram
@@ -29,7 +34,7 @@ def reseed_model_per_test_run():
2934
np.random.seed(23)
3035

3136

32-
class ConvBlocksWithAbs(torch.nn.Module):
37+
class ConvBlocksWithAbsModule(torch.nn.Module):
3338
def __init__(self, conv_in_channels: int = 3):
3439
super().__init__()
3540
self.block1 = torch.nn.Sequential(
@@ -56,36 +61,88 @@ def forward(self, x):
5661
return self.block2(x)
5762

5863

59-
class Abs(torch.nn.Module):
64+
class AbsModule(torch.nn.Module):
6065
def __init__(self):
6166
super().__init__()
6267

6368
def forward(self, x):
6469
return x.abs()
6570

6671

67-
def test_conv_abs(mocker, use_qat, input_shape: tuple[int] = (1, 3, 112, 112)):
68-
model = ConvBlocksWithAbs(conv_in_channels=input_shape[1])
72+
class TestAbsLegacyNeutronFlow:
73+
def test_conv_abs(
74+
self, mocker, use_qat, input_shape: tuple[int] = (1, 3, 112, 112)
75+
):
76+
model = ConvBlocksWithAbsModule(conv_in_channels=input_shape[1])
77+
78+
converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
79+
80+
quantized_program = to_quantized_edge_program(
81+
model,
82+
input_shape,
83+
use_qat=use_qat,
84+
use_neutron_for_format_conversion=False,
85+
use_new_flow_neutron_c=False,
86+
).exported_program()
87+
88+
tflite_flatbuffers_model, io_formats = converter_spy.spy_return
89+
exported_program: ExportedProgram = converter_spy.call_args.args[1]
90+
91+
assert not graph_contains_any_of_ops(
92+
graph=quantized_program.graph, ops=[exir_ops.edge.aten.abs.default]
93+
)
94+
95+
input_data = (np.random.random(input_shape) * 50).astype(np.int8)
96+
convert_run_compare(
97+
exported_program,
98+
tfl_model=tflite_flatbuffers_model,
99+
tflite_input_preprocess=ToChannelLastPreprocess(),
100+
tflite_output_preprocess=ToChannelFirstPreprocess(),
101+
input_data=input_data,
102+
atol=1.0,
103+
)
104+
69105

70-
converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
106+
class TestAbsNewNeutronFlow:
107+
def test__basic_nsys_inference(self):
108+
input_shape = (2, 3, 6, 7)
109+
model = AbsModule()
110+
graph_verifier = BaseGraphVerifier(
111+
exp_num_delegate_call_nodes=1, # Delegated Abs.
112+
exp_non_delegated_nodes=[],
113+
)
71114

72-
quantized_program = to_quantized_edge_program(
73-
model, input_shape, use_qat=use_qat, use_neutron_for_format_conversion=False
74-
).exported_program()
115+
lower_run_compare(
116+
model, input_shape, graph_verifier, use_new_flow_neutron_c=True
117+
)
75118

76-
tflite_flatbuffers_model, io_formats = converter_spy.spy_return
77-
exported_program: ExportedProgram = converter_spy.call_args.args[1]
119+
def test__basic_nsys_inference__big(self):
120+
# some operators have delegation requirement that size must be < 4096
121+
input_shape = (4097, 1)
122+
model = AbsModule()
123+
graph_verifier = BaseGraphVerifier(
124+
exp_num_delegate_call_nodes=1, # Delegated Abs.
125+
exp_non_delegated_nodes=[],
126+
)
78127

79-
assert not graph_contains_any_of_ops(
80-
graph=quantized_program.graph, ops=[exir_ops.edge.aten.abs.default]
81-
)
128+
lower_run_compare(
129+
model, input_shape, graph_verifier, use_new_flow_neutron_c=True
130+
)
131+
132+
def test_basic_nsys_inference__with_conv(self):
133+
input_shape = (2, 3, 6, 7)
134+
in_channels = input_shape[1]
135+
model = ConvBlocksWithAbsModule(conv_in_channels=in_channels)
136+
graph_verifier = BaseGraphVerifier(
137+
exp_num_delegate_call_nodes=1, # Delegated `Abs` + `Relu` in one partition
138+
exp_non_delegated_nodes=[
139+
NonDelegatedNode("aten_convolution_default", 2),
140+
NonDelegatedNode(
141+
"aten_relu_default", 1 # One `Relu` ends up in partition with `Abs`
142+
),
143+
],
144+
)
82145

83-
input_data = (np.random.random(input_shape) * 50).astype(np.int8)
84-
convert_run_compare(
85-
exported_program,
86-
tfl_model=tflite_flatbuffers_model,
87-
tflite_input_preprocess=ToChannelLastPreprocess(),
88-
tflite_output_preprocess=ToChannelFirstPreprocess(),
89-
input_data=input_data,
90-
atol=1.0,
91-
)
146+
lower_run_compare(
147+
model, input_shape, graph_verifier, use_new_flow_neutron_c=True
148+
)

0 commit comments

Comments
 (0)