Skip to content

Commit 2ecfc10

Browse files
NXP backend: Add maximum support using new neutron flow (#20889)
### Summary Enables support for `maximum` by the Neutron backend using the new Neutron MLIR flow, add tests verifying correct support. Only last commit applies. ### Test plan Unit tests provided. cc @robert-kalmar
1 parent c6832c8 commit 2ecfc10

12 files changed

Lines changed: 442 additions & 15 deletions

File tree

backends/nxp/backend/edge_program_converter.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
exir_ops.edge.aten.leaky_relu.default: LeakyReluConverter, # noqa F405
4646
exir_ops.edge.aten.log.default: LogConverter, # noqa F405
4747
exir_ops.edge.aten.max_pool2d_with_indices.default: MaxPool2DWithIndicesConverter, # noqa F405
48+
exir_ops.edge.aten.maximum.default: MaximumConverter, # noqa F405
4849
exir_ops.edge.aten.mean.dim: MeanDimConverter, # noqa F405
4950
exir_ops.edge.aten.mm.default: MMConverter, # noqa F405
5051
exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405

backends/nxp/backend/ir/converter/node_converter.py

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,12 @@ def supports_partitioning_result(
226226

227227
@staticmethod
228228
def _has_shared_q_params_if_quantized(node: Node) -> bool:
229-
"""Check if node has shared quantization parameters if it's quantized."""
229+
"""Check if node has shared quantization parameters for first input and output if it's quantized.
230+
231+
:param node: PyTorch fx Node with 'all_input_nodes' initialized.
232+
:return: True, if first input and output parameters have the same quantization parameters.
233+
False otherwise.
234+
"""
230235
if len(node.users) < 1 or len(node.all_input_nodes) < 1:
231236
# Some exotic operator (only consumer or only produces)
232237
return True
@@ -236,17 +241,17 @@ def _has_shared_q_params_if_quantized(node: Node) -> bool:
236241

237242
if _is_dequant_node(pre_node) and _is_quant_node(post_node):
238243
# Node is quantized
239-
pre_zp = pre_node.args[1]
240-
pre_scale = pre_node.args[2]
244+
pre_scale = pre_node.args[1]
245+
pre_zp = pre_node.args[2]
241246
pre_type = pre_node.args[5]
242247

243-
post_zp = post_node.args[1]
244-
post_scale = post_node.args[2]
248+
post_scale = post_node.args[1]
249+
post_zp = post_node.args[2]
245250
post_type = pre_node.args[5]
246251

247252
# Q-params match?
248253
return (
249-
pre_zp == post_zp and pre_scale == post_scale and pre_type == post_type
254+
pre_scale == post_scale and pre_zp == post_zp and pre_type == post_type
250255
)
251256

252257
# Node not quantized
@@ -493,3 +498,38 @@ def inputs_satisfy_broadcast_condition(node: Node) -> bool:
493498
prod(input_tensor.meta["val"].shape) == num_elements
494499
for input_tensor in node.all_input_nodes
495500
)
501+
502+
@staticmethod
503+
def _all_io_shares_quantization_parameters(node: Node) -> bool:
504+
"""Determine if given PyTorch fx Node has all inputs and output with same quantization parameters.
505+
506+
:param node: PyTorch fx Node with 'all_input_nodes' initialized.
507+
:return: True, if all input and output parameters have the same quantization parameters.
508+
False otherwise.
509+
"""
510+
post_node = list(node.users.keys())[0]
511+
if not _is_quant_node(post_node):
512+
return False
513+
output_scale, output_zp, output_type = (
514+
post_node.args[1],
515+
post_node.args[2],
516+
post_node.args[5],
517+
)
518+
519+
for input_node in node.all_input_nodes:
520+
if not _is_dequant_node(input_node):
521+
return False
522+
523+
input_scale, input_zp, input_type = (
524+
input_node.args[1],
525+
input_node.args[2],
526+
input_node.args[5],
527+
)
528+
if (input_scale, input_zp, input_type) != (
529+
output_scale,
530+
output_zp,
531+
output_type,
532+
):
533+
return False
534+
535+
return True

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@
5252
from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.max_pool2d_with_indices_converter import (
5353
MaxPool2DWithIndicesConverter,
5454
)
55+
from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.maximum_converter import (
56+
MaximumConverter,
57+
)
5558
from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.mean_dim_converter import (
5659
MeanDimConverter,
5760
)
@@ -127,6 +130,7 @@
127130
"LeakyReluConverter",
128131
"LogConverter",
129132
"MaxPool2DWithIndicesConverter",
133+
"MaximumConverter",
130134
"MeanDimConverter",
131135
"MMConverter",
132136
"MulTensorConverter",

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ def _all_io_shares_quantization_parameters(node: Node) -> bool:
4545
post_node = list(node.users.keys())[0]
4646
if not _is_quant_node(post_node):
4747
return False
48-
output_zp, output_scale, output_type = (
48+
output_scale, output_zp, output_type = (
4949
post_node.args[1],
5050
post_node.args[2],
5151
post_node.args[5],
@@ -55,14 +55,14 @@ def _all_io_shares_quantization_parameters(node: Node) -> bool:
5555
if not _is_dequant_node(input_node):
5656
return False
5757

58-
input_zp, input_scale, input_type = (
58+
input_scale, input_zp, input_type = (
5959
input_node.args[1],
6060
input_node.args[2],
6161
input_node.args[5],
6262
)
63-
if (input_zp, input_scale, input_type) != (
64-
output_zp,
63+
if (input_scale, input_zp, input_type) != (
6564
output_scale,
65+
output_zp,
6666
output_type,
6767
):
6868
return False
Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright 2026 NXP
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import torch
7+
8+
from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
9+
from executorch.backends.nxp.backend.ir.converter.node_converter import (
10+
CustomDelegationOptions,
11+
NodeConverter,
12+
)
13+
from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
14+
maximum_options,
15+
)
16+
from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
17+
from torch.fx import Node
18+
from torch.nn import Parameter
19+
20+
21+
class MaximumConverter(NodeConverter):
22+
@staticmethod
23+
def _is_supported_on_target(
24+
node: Node,
25+
neutron_target_spec: NeutronTargetSpec,
26+
parameters_mapping: dict[str, Parameter],
27+
custom_delegation_options: CustomDelegationOptions,
28+
) -> bool:
29+
if not NodeConverter.inputs_satisfy_broadcast_condition(node):
30+
return False
31+
32+
supported_types = [torch.int8, torch.uint8]
33+
if not NodeConverter.uses_quantization_type_for_io(
34+
node, supported_types, [0, 1], [0]
35+
):
36+
return False
37+
38+
return True
39+
40+
@staticmethod
41+
def _is_supported_in_IR(
42+
node: Node,
43+
parameters_mapping: dict[str, Parameter],
44+
custom_delegation_options: CustomDelegationOptions,
45+
) -> bool:
46+
if len(node.args) != 2:
47+
return False
48+
49+
if not NodeConverter._all_io_shares_quantization_parameters(node):
50+
# The IR requires all inputs to have the same quantization parameters as the output.
51+
# The quantizer should quantize the operator so that this case does not happen.
52+
return False
53+
54+
return True
55+
56+
def convert(self, node: Node):
57+
"""Convert 'maximum' operator to NeutronIR 'Maximum'.
58+
The ExecuTorch schema is:
59+
maximum(Tensor input, Tensor other, *, Tensor out=None)
60+
"""
61+
self.assert_convertible(node)
62+
t_op = self._create_tflite_op_with_io_tensors(node)
63+
t_op.builtin_options = maximum_options.Maximum()
64+
65+
ops = OpsList(middle_op=t_op)
66+
# Create additional ops in case of shape broadcasting
67+
ops.add_pre(self.builder.ensure_correct_broadcasting(t_op, t_op.tmp_outputs[0]))
68+
self.builder.append_operators(ops.flatten())

backends/nxp/neutron_partitioner.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]):
218218
exir_ops.edge.aten.leaky_relu.default: LeakyReluConverter, # noqa F405
219219
exir_ops.edge.aten.log.default: LogConverter, # noqa F405
220220
exir_ops.edge.aten.max_pool2d_with_indices.default: MaxPool2DWithIndicesConverter, # noqa F405
221+
exir_ops.edge.aten.maximum.default: MaximumConverter, # noqa F405
221222
exir_ops.edge.aten.mean.dim: MeanDimConverter, # noqa F405
222223
exir_ops.edge.aten.mm.default: MMConverter, # noqa F405
223224
exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405

backends/nxp/quantizer/neutron_quantizer.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
LeakyReluPattern,
3535
LinearPattern,
3636
LogPattern,
37+
MaximumPattern,
3738
MaxPool1DPattern,
3839
MaxPool2DPattern,
3940
MeanDimPattern,
@@ -282,6 +283,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False)
282283
OpQuantizer(LeakyReluInPlacePattern(is_qat=is_qat), static_fc_qconfig),
283284
OpQuantizer(LinearPattern(self, is_qat=is_qat), static_fc_qconfig),
284285
OpQuantizer(LogPattern(is_qat=is_qat), static_qconfig),
286+
OpQuantizer(MaximumPattern(is_qat=is_qat), static_qconfig),
285287
OpQuantizer(MaxPool1DPattern(is_qat=is_qat), static_qconfig),
286288
OpQuantizer(MaxPool2DPattern(is_qat=is_qat), static_qconfig),
287289
OpQuantizer(MeanDimPattern(is_qat=is_qat), static_qconfig),

backends/nxp/quantizer/patterns.py

Lines changed: 62 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ class SharedSpecPattern(QuantizationPattern):
120120
"""
121121

122122
@abstractmethod
123-
def partition_types(self) -> list[torch.nn.Module]:
123+
def partition_types(self) -> list[OpOverload]:
124124
pass
125125

126126
@staticmethod
@@ -152,6 +152,52 @@ def get_anchors(
152152
return self.get_shared_spec_anchors(gm, fused_partition)
153153

154154

155+
class SharedSpecMultipleInputPattern(QuantizationPattern):
156+
"""
157+
Quantization pattern for shared quantization with multiple inputs.
158+
159+
The quantization is derived from the previous node quantization and inputs and the output shares the same
160+
quantization parameters (scale and zero-point).
161+
"""
162+
163+
@abstractmethod
164+
def partition_types(self) -> list[OpOverload]:
165+
pass
166+
167+
@staticmethod
168+
def get_shared_spec_anchors(
169+
gm: fx.GraphModule, fused_partition: list[fx.GraphModule]
170+
) -> PartitionAnchors | None:
171+
node = fused_partition[0].nodes[-1]
172+
173+
quantized_input = None
174+
for prev_node in node.args:
175+
if Q_ANNOTATION_KEY in prev_node.meta:
176+
quantized_input = prev_node
177+
break
178+
179+
if quantized_input is not None:
180+
inputs = []
181+
for idx, _ in enumerate(node.args):
182+
inputs.append(
183+
(node, NodeArgsIdx(idx), SharedQuantizationSpec(quantized_input))
184+
)
185+
else:
186+
return None
187+
188+
return PartitionAnchors(
189+
inputs=inputs,
190+
weights=[],
191+
biases=[],
192+
output=[(node, SharedQuantizationSpec(quantized_input))],
193+
)
194+
195+
def get_anchors(
196+
self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule]
197+
) -> PartitionAnchors | None:
198+
return self.get_shared_spec_anchors(gm, fused_partition)
199+
200+
155201
class SingleInputBasicPattern(QuantizationPattern):
156202
@abstractmethod
157203
def partition_types(self) -> list[OpOverload]:
@@ -300,7 +346,7 @@ class AddTensorPattern(QuantizationPattern):
300346
Basic quantization for all inputs and output.
301347
"""
302348

303-
def partition_types(self) -> list[torch.nn.Module]:
349+
def partition_types(self) -> list[OpOverload]:
304350
return [torch.ops.aten.add.Tensor]
305351

306352
def get_anchors(
@@ -333,7 +379,7 @@ class BMMPattern(QuantizationPattern):
333379
Quantizer for BatchMatMul operator.
334380
"""
335381

336-
def partition_types(self) -> list[torch.nn.Module]:
382+
def partition_types(self) -> list[OpOverload]:
337383
return [torch.ops.aten.bmm.default]
338384

339385
def get_anchors(
@@ -372,7 +418,7 @@ class SubTensorPattern(QuantizationPattern):
372418
Basic quantization for all inputs and output.
373419
"""
374420

375-
def partition_types(self) -> list[torch.nn.Module]:
421+
def partition_types(self) -> list[OpOverload]:
376422
return [torch.ops.aten.sub.Tensor]
377423

378424
def get_anchors(
@@ -426,7 +472,7 @@ def get_anchors(
426472

427473
quantized_input = None
428474
for prev_node in node.args[0]:
429-
if "quantization_annotation" in prev_node.meta:
475+
if Q_ANNOTATION_KEY in prev_node.meta:
430476
quantized_input = prev_node
431477
break
432478

@@ -876,6 +922,17 @@ def partition_types(self):
876922
return [torch.ops.aten.log.default]
877923

878924

925+
class MaximumPattern(SharedSpecMultipleInputPattern):
926+
"""
927+
Quantization pattern for Maximum quantization. Accepts 1 or 2 input nodes.
928+
929+
Basic quantization for all inputs and output.
930+
"""
931+
932+
def partition_types(self) -> list[OpOverload]:
933+
return [torch.ops.aten.maximum.default]
934+
935+
879936
class MaxPool1DPattern(SharedSpecPattern):
880937
"""Quantizer for the MaxPool1D operator."""
881938

0 commit comments

Comments
 (0)