diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index 9df8290e85d..1ac7d7d7f47 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -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 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py index 93ba24e61bd..cc648b9fef8 100755 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py @@ -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, ) @@ -107,6 +110,7 @@ "AdaptiveAvgPool2dConverter", "AddMMConverter", "AddTensorConverter", + "AminConverter", "AvgPool2dConverter", "BMMConverter", "CatConverter", diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/amin_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/amin_converter.py new file mode 100644 index 00000000000..d8042ff01cc --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/amin_converter.py @@ -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 + + 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 + + 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()) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py index 8674bf697c7..4da910b4027 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py @@ -5,12 +5,7 @@ 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, is_not_qdq_node, @@ -18,6 +13,8 @@ ) 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, @@ -39,7 +36,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( @@ -88,140 +85,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: @@ -235,13 +98,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()) diff --git a/backends/nxp/backend/ir/converter/node_converters/shared/reduce_utils.py b/backends/nxp/backend/ir/converter/node_converters/shared/reduce_utils.py index 65528bb30f8..1e51107f499 100755 --- a/backends/nxp/backend/ir/converter/node_converters/shared/reduce_utils.py +++ b/backends/nxp/backend/ir/converter/node_converters/shared/reduce_utils.py @@ -4,6 +4,7 @@ # LICENSE file in the root directory of this source tree. import numpy as np + from executorch.backends.nxp.backend.data_format import DataFormat from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( @@ -12,6 +13,7 @@ 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.tflite_generator import tflite_model +from torch.fx import Node def convert_axes_from_attribute( @@ -95,3 +97,132 @@ def ensure_reduce_transposition(builder, ops: OpsList): transpose.tmp_inputs[0].tensor_format = DataFormat.CHANNELS_FIRST ops.post_ops.insert(0, transpose) + + +def _to_pos_dim(d: int, rank: int): + return d + rank if d < 0 else d + + +def _normalize_dim(dim: list[int], rank: int) -> list[int]: + # convert negative index to positive + return [_to_pos_dim(d, rank) for d in dim] + + +def _normalize_and_to_channel_last_dim(dim: list[int], rank: int) -> list[int]: + # convert negative index to positive + dim = _normalize_dim(dim, rank) + + perm = translator.create_channels_last_to_channels_first_permutation(rank, True) + dim = [perm[d] for d in dim] + + # noinspection PyTypeChecker + return dim + + +def get_reduce_node_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( + builder, 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 = _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 node format inference." + + # Just adjust the dim to match the input format. + if channels_last_input: + dim = _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 = _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 = _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( + 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 `reduce-type node` 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( + 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 -> reduce-type node -> 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 `reduce-type node` 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( + 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( + builder.create_transpose_operator_after(t_op, 0, to_channels_last_perm) + ) + t_op.tmp_outputs[0].tensor_format = DataFormat.CHANNELS_FIRST + + return dim diff --git a/backends/nxp/backend/node_format_inference.py b/backends/nxp/backend/node_format_inference.py index 030873c88ab..288831b05bc 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -16,6 +16,7 @@ from executorch.backends.nxp.backend.edge_program_converter import functions_converters from executorch.backends.nxp.tests.ops_aliases import ( AdaptiveAvgPool2D, + Amin, AvgPool2D, Convolution, DequantizePerChannel, @@ -58,6 +59,7 @@ class NodeFormatInference: ViewCopy, PermuteCopy, MeanDim, + Amin, } _type_changed_during_last_run: bool @@ -134,7 +136,7 @@ def _infer_format_of_nodes(self, node: Node): self._node_inputs[node][0], DataFormat.FORMATLESS ) - elif op_type == MeanDim: + elif op_type in [MeanDim, Amin]: # The operator schema is: # mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor keep_dim = try_get_arg(node, 2) or False diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index 9cc174b97e0..8b81eb00505 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -204,6 +204,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]): 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 diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 94ee8e8656a..bf3badb7d65 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -16,6 +16,7 @@ AdaptiveAvgPoolPattern, AddmmPattern, AddTensorPattern, + AminPattern, AvgPool1DPattern, AvgPool2DPattern, BatchNormPattern, @@ -260,6 +261,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) OpQuantizer(AdaptiveAvgPoolPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddmmPattern(self, is_qat=is_qat), static_fc_qconfig), + OpQuantizer(AminPattern(is_qat=is_qat), static_fc_qconfig), OpQuantizer(AvgPool1DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AvgPool2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(BatchNormPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index 6bd8c73ac97..6b6395d887c 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -316,6 +316,15 @@ def get_anchors( ) +class AminPattern(SharedSpecPattern): + """ + Quantizer for Amin operator. + """ + + def partition_types(self): + return [torch.ops.aten.amin.default] + + class BMMPattern(QuantizationPattern): """ Quantizer for BatchMatMul operator. diff --git a/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py new file mode 100644 index 00000000000..6fa7038fc97 --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py @@ -0,0 +1,412 @@ +# 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.backend.ir.converter.builder.model_builder import ( + ModelBuilder, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.max_pool_2d_options import ( + MaxPool2D, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.reduce_min_options import ( + ReduceMin, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( + Transpose, +) +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +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 ( + AddTensor, + Amin, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) +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 AminModule(torch.nn.Module): + def __init__( + self, dim: int | torch.Size | list[int] | tuple[int, ...], keepdim: bool + ): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + return torch.amin(x, dim=self.dim, keepdim=self.keepdim) + + +class AminAddModule(AminModule): + def forward(self, x): + x = super().forward(x) + return x + x + + +class MaxPoolAminModule(torch.nn.Module): + @staticmethod + def noop_max_pool_2d(x): + """Call `torch.max_pool2d` that is a NoOp, but it enforces the ChannelsFirst format in the `NodeFormatInference`.""" + return torch.max_pool2d(x, kernel_size=1) + + def __init__( + self, dim: int | torch.Size | list[int] | tuple[int, ...], keepdim: bool + ): + super().__init__() + self.dim, self.keepdim = dim, keepdim + + def forward(self, x): + x = self.noop_max_pool_2d(x) + x = torch.amin(x, dim=self.dim, keepdim=self.keepdim) + return x + + +class AminMaxPoolModule(MaxPoolAminModule): + def forward(self, x): + x = torch.amin(x, dim=self.dim, keepdim=self.keepdim) + x = self.noop_max_pool_2d(x) + return x + + +def assert_delegated( + model, + input_shape, + mocker, + request, + use_qat=False, + expected_delegated_ops=None, +): + if expected_delegated_ops is None: + expected_delegated_ops = {Amin: 1} + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops=expected_delegated_ops, + expected_non_delegated_ops={}, + ) + + # Cover also negative values to thoroughly test the operator. + dataset_creator = RandomDatasetCreator(low=-2, high=2) + + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + dataset_creator, + output_comparator, + use_qat=use_qat, + remove_quant_io_ops=remove_quant_io_ops, + ) + + +def assert_not_delegated(model, input_shape): + delegated_ep = to_quantized_edge_program(model, input_shape).exported_program() + + # Make sure the `amin` was NOT delegated. + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [Amin]) + + +class TestAminConverter: + # noinspection PyMethodMayBeStatic + @pytest.fixture(params=[True, False], ids=lambda keep_dim: f"keep_dim = {keep_dim}") + def keep_dim(self, request): + return request.param + + def test__basic_nsys_inference__qat(self, mocker, request, use_qat, keep_dim): + input_shape = (23,) + model = AminModule(0, keep_dim) + assert_delegated(model, input_shape, mocker, request, use_qat=use_qat) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((5,), 0, id="1D, dim = 0."), + pytest.param((4, 2), 0, id="2D, dim = 0."), + pytest.param((4, 2), -1, id="2D, dim = -1."), + pytest.param((3, 1, 4), 2, id="3D, dim = 2."), + pytest.param((1, 3, 3, 7), 3, id="4D, dim = 3."), + pytest.param((3, 1, 4, 1, 5), -1, id="5D, dim = -1."), + pytest.param((3, 1, 4, 1, 5), 0, id="5D, dim = 0."), + ], + ) + def test__single_dims(self, mocker, request, input_shape, dim, keep_dim): + model = AminModule(dim, keep_dim) + assert_delegated(model, input_shape, mocker, request) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((4, 2), (-2,), id="2D, dim = (-2,)."), + pytest.param((2, 3, 4), (0, 2), id="3D, dim = (0, 2,)."), + pytest.param((1, 3, 3, 7), (2, -3), id="4D, dim = (2, -3)."), + pytest.param((1, 3, 3, 7), -2, id="4D, dim = -2."), + pytest.param((3, 1, 4, 1, 5), (3, -5, -4), id="5D, dim = (3, -5 ,-4)."), + ], + ) + def test__tuple_dims(self, mocker, request, input_shape, dim, keep_dim): + model = AminModule(dim, keep_dim) + assert_delegated(model, input_shape, mocker, request) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((3, 1, 4), 1, id="3D, dim = 1."), + pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + ], + ) + def test__noop__only_node__not_delegated(self, input_shape, dim): + keep_dim = True # Reduction over a dimension of size `1` with `keep_dim=True` is a no-op. + model = AminModule(dim, keep_dim) + assert_not_delegated(model, input_shape) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((3, 1, 4), 1, id="3D, dim = 1."), + pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + ], + ) + def test__noop__not_only_node__delegated(self, mocker, request, input_shape, dim): + keep_dim = True # Reduction over a dimension of size `1` with `keep_dim=True` is a no-op. + model = AminAddModule(dim, keep_dim) + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={Amin: 1, AddTensor: 1}, + ) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((3, 1, 4), 1, id="3D, dim = 1."), + pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + pytest.param((1, 7, 3, 3), [0], id="4D, dim = [0]."), + ], + ) + def test__no_reduction__keepdim_false__delegated( + self, mocker, request, input_shape, dim + ): + # These cases reduce over a dimension of size 1. + # When `keep_dim=True` the node is a noop, and it's not delegated (see `test__noop__only_node__not_delegated`), + # but with `keep_dim=False` it changes the shape so it's not a noop and is therefore delegated successfully. + keep_dim = False + model = AminModule(dim, keep_dim) + assert_delegated(model, input_shape, mocker, request) + + def test__channels_first__keep_dim__true(self, mocker, request): + # Just 1 test case to verify correct handling of the `dim`. + # Most cases fall into the single bit error case, and since this test uses 2 operators, the error accumulates + # and the final error is larger. We cannot with 100% certainty say that the error is only caused by the single + # bit errors and not related to the format. That's why only this 1 case with no errors is used. + input_shape, dim = (1, 7, 3, 3), 1 + model = MaxPoolAminModule(dim, True) + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={MaxPool2DWithIndices: 1, GetItem: 1, Amin: 1}, + ) + + class TestKeepDimFalseFormatHandling: + """When `keep_dim = False`, the `amin` operator changes the rank, so the format have to be explicitly + handled. The tests in this class focus on the related edge cases. + """ + + def _assert_neutron_ir_model_has_ops( + self, model_builder_finish_spy, expected_ops + ): + assert ( + model_builder_finish_spy.call_count == 1 + ), "Conversion to Neutron IR happened multiple times." + + neutron_ir_ops = model_builder_finish_spy.spy_return.sub_graphs[ + 0 + ].operators.vector + assert len(neutron_ir_ops) == len( + expected_ops + ), "Neutron IR model doesn't have the expected number of ops." + + for op, expected_op in zip(neutron_ir_ops, expected_ops, strict=True): + assert isinstance( + op.builtin_options, expected_op + ), f"Expected {expected_op}, got {op}." + + @pytest.mark.parametrize( + "dim", + [ + 1, + [0, -3], + (-4, 1, 2), + [-3, 3], + [1, 2, 3], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__reducing_channels(self, mocker, request, dim): + # If the channels dimension is reduced (removed), the `amin` output will always be equal in channels first + # and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolAminModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amin: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + ReduceMin, + ], + ) + + @pytest.mark.parametrize( + "dim", + [ + (2, 3), + [1, -2, 3], + [-1, -2, 0], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__reducing_all_spatial_dims( + self, mocker, request, dim + ): + # If the spatial dimensions are reduced (removed), the `amin` output will always be equal in channels + # first and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolAminModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amin: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + ReduceMin, + ], + ) + + @pytest.mark.parametrize( + "dim", + [ + 0, + (2,), + [-1, 0], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__not_reducing_channels_or_all_spatial_dims( + self, mocker, request, dim + ): + # If the channels dimension is not reduced, a `Transpose` operator must be added to make the input channels + # first in Neutron IR. + + input_shape = (1, 7, 3, 3) + model = MaxPoolAminModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amin: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Transpose, # The necessary `Transpose` operator. + ReduceMin, + ], + ) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((2, 3, 4, 5, 6), 0, id="dim=0, 5D->4D"), + pytest.param((2, 3, 4, 5, 6), [-3], id="dim=[-3], 5D->4D"), + pytest.param((1, 2, 3, 4, 5, 6), (1, -1), id="dim=(1, -1), 6D->4D"), + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_output(self, mocker, request, input_shape, dim): + model = AminMaxPoolModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amin: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + ReduceMin, + Transpose, # The necessary `Transpose` operator. + MaxPool2D, + Transpose, + ], + ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py index f84471169ea..5f2238a37e6 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py @@ -305,7 +305,7 @@ def test__channels_first_input__reducing_channels(self, mocker, request, dim): def test__channels_first_input__reducing_all_spatial_dims( self, mocker, request, dim ): - # If tall he spatial dimensions are reduced (removed), the `mean` output will always be equal in channels + # If the spatial dimensions are reduced (removed), the `mean` output will always be equal in channels # first and channels last, so no `Transpose` ops are added. input_shape = (1, 7, 3, 3) model = MaxPoolMeanDimModule(dim, False) diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index da50d4dc0d9..f01ddcea75b 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -15,6 +15,7 @@ AdaptiveAvgPool2D = exir_ops.edge.aten._adaptive_avg_pool2d.default AddMm = exir_ops.edge.aten.addmm.default AddTensor = exir_ops.edge.aten.add.Tensor +Amin = exir_ops.edge.aten.amin.default AvgPool2D = exir_ops.edge.aten.avg_pool2d.default Bmm = exir_ops.edge.aten.bmm.default Cat = exir_ops.edge.aten.cat.default diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index fb67f47bf62..caa5aa02718 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -3,6 +3,7 @@ aten.abs.default,int8,static int8, aten._adaptive_avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" aten.addmm.default,int8,static int8,2D tensor only aten.add.Tensor,int8,static int8,"alpha = 1, input tensors of equal shape" +aten.amin.default,int8,static int8, aten.avg_pool1d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" aten.avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" aten.bmm.default,int8,static int8,"width and channels dim of both args %8 = 0, 3D tensors only" @@ -21,15 +22,15 @@ aten.max_pool1d.default,int8,static int8,"dilation=1, ceil_mode=False, channels% aten.max_pool2d.default,int8,static int8,"dilation=1, ceil_mode=False, channels%8=0, batch_size=1, stride_h=1 or 2" aten.max_pool2d_with_indices.default,int8,static int8,"dilation=1, ceil_mode=False, channels%8=0, batch_size=1, stride_h=1 or 2" aten.mean.dim,int8,static int8,"4D tensor only, dims = [-1,-2] or [-2,-1]" -aten.mul.Tensor, int8, static int8, "tensor-size % 8 = 0" +aten.mul.Tensor,int8, static int8, "tensor-size % 8 = 0" aten.mm.default,int8,static int8,"2D tensor only" aten.neg.default,int8,static int8, aten.permute_copy.default,int8, static int8, "Only specific transpositions supported, see backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py" -aten.prelu.default, int8, static int8, "rank = 4, channels % 8 = 0, flat input size / channels <= 4096" +aten.prelu.default,int8, static int8, "rank = 4, channels % 8 = 0, flat input size / channels <= 4096" aten.relu.default,int8,static int8, aten.sigmoid.default,int8,static int8, -aten.slice_copy.Tensor, int8, static int8 -aten._softmax.default, int8, static int8, "rank > 1, channels % 8 = 0, channels < 2048, flat input size / channels <= 4096, flat input size <= 524288" +aten.slice_copy.Tensor,int8, static int8 +aten._softmax.default,int8, static int8, "rank > 1, channels % 8 = 0, channels < 2048, flat input size / channels <= 4096, flat input size <= 524288" aten.split.default,N/A, N/A, "transforming split -> getitem to slice, see aten.slice_copy.Tensor" aten.split.Tensor,N/A, N/A, "transforming split -> getitem to slice, see aten.slice_copy.Tensor" aten.split_with_sizes.default,N/A, N/A, "transforming split -> getitem to slice, see aten.slice_copy.Tensor"