Skip to content

NXP backend: Enable amin with new Neutron flow #20549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
roman-janik-nxp wants to merge 1 commit into
base: main
Choose a base branch
from
Open

NXP backend: Enable amin with new Neutron flow #20549

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
exir_ops.edge.aten._adaptive_avg_pool2d.default: AdaptiveAvgPool2dConverter, # noqa F405
exir_ops.edge.aten.addmm.default: AddMMConverter, # noqa F405
exir_ops.edge.aten.add.Tensor: AddTensorConverter, # noqa F405
exir_ops.edge.aten.amin.default: AminConverter, # noqa F405
exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter, # noqa F405
exir_ops.edge.aten.bmm.default: BMMConverter, # noqa F405
exir_ops.edge.aten.cat.default: CatConverter, # noqa F405
Expand Down
4 changes: 4 additions & 0 deletions backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,9 @@
from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.addmm_converter import (
AddMMConverter,
)
from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.amin_converter import (
AminConverter,
)
from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.avg_pool_2d_converter import (
AvgPool2dConverter,
)
Expand Down Expand Up @@ -107,6 +110,7 @@
"AdaptiveAvgPool2dConverter",
"AddMMConverter",
"AddTensorConverter",
"AminConverter",
"AvgPool2dConverter",
"BMMConverter",
"CatConverter",
Expand Down
101 changes: 101 additions & 0 deletions backends/nxp/backend/ir/converter/node_converters/ops_converters/amin_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
# 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.conversion.common import OpsList
from executorch.backends.nxp.backend.ir.converter.node_converter import (
CustomDelegationOptions,
is_not_qdq_node,
NodeConverter,
)
from executorch.backends.nxp.backend.ir.converter.node_converters.shared.reduce_utils import (
convert_axes_from_attribute,
get_dim_and_handle_io_formats,
get_reduce_node_attrs,
)
from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
reduce_min_options,
)
from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
from torch.fx import Node
from torch.fx.passes.infra.partitioner import Partition
from torch.nn import Parameter


class AminConverter(NodeConverter):

@classmethod
def supports_partitioning_result(
cls,
node: Node,
partition_list: list[Partition],
custom_delegation_options: CustomDelegationOptions,
neutron_target_spec: NeutronTargetSpec,
parameters_mapping: dict[str, Parameter],
) -> bool:
dim, keepdim = get_reduce_node_attrs(node)
input_shape = node.args[0].meta["val"].shape

is_alone_in_partition = cls.is_node_alone_in_partition(
node, partition_list, filter_fn=is_not_qdq_node
)

if is_alone_in_partition and keepdim and all(input_shape[d] == 1 for d in dim):
# The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the
# partition, the graph would end up empty.
return False
Comment on lines +46 to +49

@MartinPavella MartinPavella Jun 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Well done for considering this case 👍🏻

Just FYI, NeutronPartitioner automatically tries to identify all no-ops by running every operator selected for delegation with random data. If it produces outputs that are exactly equal to the inputs (data, shape, type), it considers the node to be a no-op, and it will not delegate it as long as the operator is in a list of allowed no-ops. In the case of amin, I believe this system would have always correctly identified the no-ops, so simply adding the aten.amin into the list of allowed no-ops would have done the job.

But your solution is valid as well. No change needed.

@roman-janik-nxp roman-janik-nxp Jun 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I learn from the best (you 😄 ). However I don't see any logic for said identifying in NeutronPartitioner and when I omit these checks in supports_partitioning_result() the tests fail. So I'm keeping it.


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

@staticmethod
def _is_supported_in_IR(
node: Node,
parameters_mapping: dict[str, Parameter],
custom_delegation_options: CustomDelegationOptions,
) -> bool:
if not NodeConverter._has_shared_q_params_if_quantized(node):
return False
Comment on lines +76 to +77

@MartinPavella MartinPavella Jun 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where did you find this requirement?

@roman-janik-nxp roman-janik-nxp Jun 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In your mean dim converter, I kept it, but I don't think it's necessary.

@MartinPavella MartinPavella Jun 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It depends on the requirements of Neutron/Neutron IR. Often we can see requirements for equal input and output quantization here: https://developers.google.com/edge/litert/conversion/tensorflow/quantization/quantization_spec
The spec is, however, incomplete. It doesn't mention the ReduceMin operator. The Neutron documentation doesn't mention any such restrictions, so I think we should remove it. Also note that you used the shared quantization spec in our quantizer for the amin operator, so please update that too.

P.S.: We can handle mean later.


return True

def convert(self, node: Node):
"""Convert the 'amin' operator to NeutronIR 'ReduceMin'.
The ExecuTorch schema is:
amin(
Tensor self,
int[1]? dim,
bool keepdim=False,
) -> Tensor
"""
self.assert_convertible(node)

dim, keepdim = get_reduce_node_attrs(node)

t_op = self._create_tflite_op_with_io_tensors(node)
t_op.builtin_options = reduce_min_options.ReduceMin(keepdim)

ops = OpsList(middle_op=t_op)
dim = get_dim_and_handle_io_formats(self.builder, ops, dim, keepdim)

convert_axes_from_attribute(t_op, self.builder, dim)
self.builder.append_operators(ops.flatten())
147 changes: 5 additions & 142 deletions backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,15 @@

import torch

from executorch.backends.nxp.backend.data_format import DataFormat
from executorch.backends.nxp.backend.ir.converter.conversion import translator
from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
from executorch.backends.nxp.backend.ir.converter.conversion.translator import (
create_channels_last_to_channels_first_permutation,
)
from executorch.backends.nxp.backend.ir.converter.node_converter import (
CustomDelegationOptions,
NodeConverter,
)
from executorch.backends.nxp.backend.ir.converter.node_converters.shared.reduce_utils import (
convert_axes_from_attribute,
get_dim_and_handle_io_formats,
get_reduce_node_attrs,
)
from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
mean_options,
Expand All @@ -38,7 +35,7 @@ def supports_partitioning_result(
neutron_target_spec: NeutronTargetSpec,
parameters_mapping: dict[str, Parameter],
) -> bool:
dim, keepdim = MeanDimConverter._get_attrs(node)
dim, keepdim = get_reduce_node_attrs(node)
input_shape = node.args[0].meta["val"].shape

is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list)
Expand Down Expand Up @@ -85,140 +82,6 @@ def _is_supported_in_IR(

return True

@staticmethod
def _to_pos_dim(d: int, rank: int):
return d + rank if d < 0 else d

@staticmethod
def _normalize_dim(dim: list[int], rank: int) -> list[int]:
# convert negative index to positive
return [MeanDimConverter._to_pos_dim(d, rank) for d in dim]

@staticmethod
def _normalize_and_to_channel_last_dim(dim: list[int], rank: int) -> list[int]:
# convert negative index to positive
dim = MeanDimConverter._normalize_dim(dim, rank)

perm = create_channels_last_to_channels_first_permutation(rank, True)
dim = [perm[d] for d in dim]

# noinspection PyTypeChecker
return dim

@staticmethod
def _get_attrs(node: Node) -> tuple[list[int], bool]:
dim = node.args[1]
keepdim = node.args[2] if len(node.args) >= 3 else False
return dim, keepdim

def _get_dim_and_handle_io_formats(
self, ops: OpsList, dim: list[int], keep_dim: bool
):
t_op = ops.middle_op
x = t_op.tmp_inputs[0]
y = t_op.tmp_outputs[0]

channels_last_input = x.tensor_format.is_channels_last()
channels_last_output = y.tensor_format.is_channels_last()
formatless_input = not channels_last_input
formatless_output = not channels_last_output

dim = self._normalize_dim(dim, x.rank)

if keep_dim:
# The rank is preserved and the io formats should always be equal.
assert (
x.tensor_format == y.tensor_format
), "NXP backend: There is a bug in `mean.dim` format inference."

# Just adjust the dim to match the input format.
if channels_last_input:
dim = self._normalize_and_to_channel_last_dim(dim, x.rank)

else:
# `keep_dim = False`, so the output rank != input rank, and the operator changes the tensor format.

if channels_last_input and formatless_output:
if 1 in dim:
# If we are reducing over the channels, the channels dimension gets removed and the output ends up
# exactly equal in channels last and channels first, regardless of which other dimensions are
# removed. Therefore, we can just adjust the `dim` and we don't need to insert any `Transpose` ops.
dim = self._normalize_and_to_channel_last_dim(dim, x.rank)
elif all(spatial_dim in dim for spatial_dim in range(2, x.rank)):
# All spatial dims are reduced, leaving only batch and channels (both optionally). So the result is
# equal in channels first and channels last as long as we adjust the `dim` to match a channels last
# input (similarly to the case above).
dim = self._normalize_and_to_channel_last_dim(dim, x.rank)
else:
# If the channels dimension is preserved, we must transpose the input to channels first (to match
# the edge model) and we must keep the `dim` unchanged (referencing channels first dimensions).
# Otherwise, the output would not match the input.
to_channels_first_perm = (
translator.create_channels_last_to_channels_first_permutation(
x.rank
)
)
ops.add_pre(
self.builder.create_transpose_operator_before(
t_op, 0, to_channels_first_perm
)
)
t_op.tmp_inputs[0].tensor_format = DataFormat.CHANNELS_FIRST

elif formatless_input and channels_last_output:
# We need apply the `mean` with the original `dim`, which will produce a channels first output. Then,
# we need to append a `Transpose` operator to make the output channels last.
to_channels_last_perm = (
translator.create_channels_first_to_channels_last_permutation(
y.rank, True
)
)
ops.add_post(
self.builder.create_transpose_operator_after(
t_op, 0, to_channels_last_perm
)
)
t_op.tmp_outputs[0].tensor_format = DataFormat.CHANNELS_FIRST

elif formatless_input and formatless_output:
# No action needed.
pass

else: # channels_last_input and channels_last_output
# This case cannot currently occur, as it would require the case:
# channels last 4D -> mean -> channels_last 3D
# which cannot currently happen as the 3D conv/pooling/... is supported by adding `view_copy` nodes in
# the edge dialect and converting the node to 4D, and the `view_copy` nodes prevent the propagation of
# the format to the `mean.dim` output.
# Therefore, the implementation cannot be tested. But from experience with other operators, it should
# work correctly. We just need to add 2 `Transpose` ops to make the IO channels first, and keep the
# `dim` unchanged.
to_channels_first_perm = (
translator.create_channels_last_to_channels_first_permutation(
x.rank
)
)
ops.add_pre(
self.builder.create_transpose_operator_before(
t_op, 0, to_channels_first_perm
)
)
t_op.tmp_inputs[0].tensor_format = DataFormat.CHANNELS_FIRST

to_channels_last_perm = (
translator.create_channels_first_to_channels_last_permutation(
y.rank, True
)
)
ops.add_post(
self.builder.create_transpose_operator_after(
t_op, 0, to_channels_last_perm
)
)
t_op.tmp_outputs[0].tensor_format = DataFormat.CHANNELS_FIRST

return dim

def convert(self, node: Node):
"""Convert the 'mean.dim' operator to NeutronIR 'Mean'.
The ExecuTorch schema is:
Expand All @@ -232,13 +95,13 @@ def convert(self, node: Node):
"""
self.assert_convertible(node)

dim, keepdim = self._get_attrs(node)
dim, keepdim = get_reduce_node_attrs(node)

t_op = self._create_tflite_op_with_io_tensors(node)
t_op.builtin_options = mean_options.Mean(keepdim)

ops = OpsList(middle_op=t_op)
dim = self._get_dim_and_handle_io_formats(ops, dim, keepdim)
dim = get_dim_and_handle_io_formats(self.builder, ops, dim, keepdim)

convert_axes_from_attribute(t_op, self.builder, dim)
self.builder.append_operators(ops.flatten())
Loading
Loading