diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py index 4b9ff6fe85a..a2b21b73b35 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py @@ -1,31 +1,59 @@ -# 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 logging import executorch.backends.nxp.backend.ir.lib.tflite.Padding as tflPadding +import torch + +from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT from executorch.backends.nxp.backend.ir.converter.conversion import common from executorch.backends.nxp.backend.ir.converter.node_converter import ( CustomDelegationOptions, NodeConverter, ) -from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( average_pool_2d_options, ) -from torch import Size + +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec from torch.fx import Node from torch.nn import Parameter +KernelSize = tuple[int, int] +Stride = tuple[int, int] + class AdaptiveAvgPool2dConverter(NodeConverter): + @staticmethod + def _get_equivalent_avg_pool_parameters(node: Node) -> tuple[KernelSize, Stride]: + input_size = node.args[0].meta["val"].shape[2:] # Spatial dims from NCHW shape. + output_size = node.args[1] + stride = (input_size[0] // output_size[0], input_size[1] // output_size[1]) + kernel_size = ( + input_size[0] - (output_size[0] - 1) * stride[0], + input_size[1] - (output_size[1] - 1) * stride[1], + ) + + return kernel_size, stride + @staticmethod def _is_supported_in_IR( node: Node, parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + if ( + format_ := node.meta.get(NXP_NODE_FORMAT) + ) is None or not format_.is_channels_first(): + logging.warning( + "NXP backend: `adaptive_avg_pool_2d` doesn't have the required input format for delegation. " + "Please run `NodeFormatInference.identify_node_formats()` during lowering or report this issue." + ) + return False + input_size = node.args[0].meta["val"].shape output_size = node.args[1] @@ -39,30 +67,53 @@ def _is_supported_in_IR( return True - # noinspection PyMethodMayBeStatic - def _convert_adaptive_avg_pool_2d( - self, input_size: Size, output_size: list[int], t_op: tflite_model.Operator - ): - t_op.builtin_options = average_pool_2d_options.AveragePool2D() - stride = [input_size[-2] // output_size[-2], input_size[-1] // output_size[-1]] - common.assign_2d_strides(t_op.builtin_options, stride) - t_op.builtin_options.filter_h = ( - input_size[-2] - (output_size[-2] - 1) * stride[-2] - ) - t_op.builtin_options.filter_w = ( - input_size[-1] - (output_size[-1] - 1) * stride[-1] + @staticmethod + def _is_supported_on_target( + node: Node, + neutron_target_spec: NeutronTargetSpec, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + kernel_size, stride = ( + AdaptiveAvgPool2dConverter._get_equivalent_avg_pool_parameters(node) ) - t_op.builtin_options.padding = tflPadding.Padding.VALID - # AdaptiveAvgPool2d Node format: (Tensor self, SymInt[2] output_size) + if custom_delegation_options.use_new_flow_neutron_c: + # Requirements specified by the new Neutron flow documentation. + + if not NodeConverter.uses_quantization_type_for_io( + node, + supported_types=[torch.int8, torch.uint8], + input_indices=[0], + output_indices=[0], + ): + return False + + if any(k > 4096 for k in kernel_size): + return False + + if any(s > 4096 for s in stride): + return False + + return True + def convert(self, node: Node): - """Convert '_adaptive_avg_pool2d' operator to TFLite 'AveragePool2D'.""" + """Convert the '_adaptive_avg_pool2d' operator to NeutronIR 'AveragePool2D'. + The ExecuTorch schema is: + _adaptive_avg_pool2d( + Tensor self, + SymInt[2] output_size + ) -> Tensor + """ self.assert_convertible(node) - input_size = node.args[0].meta["val"].shape - output_size = node.args[1] - t_op = self._create_tflite_op_with_io_tensors(node) + t_op.builtin_options = average_pool_2d_options.AveragePool2D() + + kernel_size, stride = self._get_equivalent_avg_pool_parameters(node) + + common.assign_2d_strides(t_op.builtin_options, stride) + t_op.builtin_options.filter_h, t_op.builtin_options.filter_w = kernel_size + t_op.builtin_options.padding = tflPadding.Padding.VALID - self._convert_adaptive_avg_pool_2d(input_size, output_size, t_op) self.builder.append_operators([t_op]) diff --git a/backends/nxp/backend/node_format_inference.py b/backends/nxp/backend/node_format_inference.py index 80d9badbce0..65e34b7fbde 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -25,6 +25,8 @@ class NodeFormatInference: # The op in the dictionary is mapped to a dictionary, which holds indices to input nodes # that are always channels first. ops_with_channels_first_nodes = { + exir_ops.edge.aten._adaptive_avg_pool2d.default: {"inputs": [0]}, + torch.ops.aten.adaptive_avg_pool2d.default: {"inputs": [0]}, exir_ops.edge.aten.avg_pool2d.default: {"inputs": [0]}, exir_ops.edge.aten.convolution.default: {"inputs": [0, 1]}, exir_ops.edge.aten.max_pool2d_with_indices.default: {"inputs": [0]}, diff --git a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py index 85ada41a095..1e6c916efa0 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py @@ -1,4 +1,4 @@ -# 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. @@ -10,15 +10,29 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +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 ( convert_run_compare, + graph_contains_any_of_ops, ToChannelFirstPreprocess, ToChannelLastPreprocess, ) +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + AllCloseOutputComparator, +) from executorch.backends.nxp.tests.models import ( AdaptiveAvgPool2dConvMeanDimModule, AdaptiveAvgPool2dConvModule, + AdaptiveAvgPool2dModule, +) + +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare + +from executorch.backends.nxp.tests.ops_aliases import ( + AdaptiveAvgPool2D, + ExecutorchDelegateCall, ) from torch.export import ExportedProgram from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -151,3 +165,104 @@ def test_adaptive_avg_pool_2d_mean_dim_quant_conversion(mocker, use_qat): tflite_output_preprocess=ToChannelFirstPreprocess(), input_data=input_data, ) + + +class TestAdaptiveAvgPool2DNewNeutronFlow: + def test__basic_nsys_inference(self, mocker, use_qat): + input_shape = (2, 3, 4, 6) + output_size = (2, 3) + model = AdaptiveAvgPool2dModule(output_size) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={AdaptiveAvgPool2D: 1}, + expected_non_delegated_ops={}, + ) + + output_comparator = AllCloseOutputComparator( + 7.8e-3 + ) # Accept small error due to Neutron bug (AIR-14585). + + lower_run_compare( + model, + input_shape, + graph_verifier, + RandomDatasetCreator(low=-1, high=1), + output_comparator=output_comparator, + use_qat=use_qat, + use_new_flow_neutron_c=True, + ) + + @pytest.mark.xfail( + strict=True, + reason="Known Neutron bad compute issue. Will be fixed in Neutron SW 3.1.2.", + ) + def test__know_neutron_issue(self, mocker): + input_shape = (2, 3, 10, 15) + output_size = (5, 5) + model = AdaptiveAvgPool2dModule(output_size) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={AdaptiveAvgPool2D: 1}, + expected_non_delegated_ops={}, + ) + + # Use high tolerance so we notice when the issue is fixed. + output_comparator = AllCloseOutputComparator(7.8e-3) + + lower_run_compare( + model, + input_shape, + graph_verifier, + RandomDatasetCreator(low=-1, high=1), + output_comparator=output_comparator, + use_new_flow_neutron_c=True, + ) + + def test__kernel_size_and_stride_limit(self, mocker): + input_shape = (1, 3, 4, 4096) # input_size = (1, 4096) + output_size = ( + 2, + 1, + ) # If we reduced both dims to 1, ExecuTorch would replace the op with mean. + # stride = input_size // output_size = 4096 / 1 = 4096 + # kernel_size = input_size - (output_size - 1) * stride = 4096 - 0 * 4096 = 4096 + + model = AdaptiveAvgPool2dModule(output_size) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={AdaptiveAvgPool2D: 1}, + expected_non_delegated_ops={}, + ) + + output_comparator = AllCloseOutputComparator( + 7.9e-3 + ) # Accept small error due to Neutron bug (AIR-14585). + + lower_run_compare( + model, + input_shape, + graph_verifier, + RandomDatasetCreator(low=-1, high=1), + output_comparator=output_comparator, + use_new_flow_neutron_c=True, + ) + + def test__kernel_size_and_stride_limit_exceeded(self): + input_shape = (1, 3, 4, 4097) # input_size = (1, 4097) + output_size = ( + 2, + 1, + ) # If we reduced both dims to 1, ExecuTorch would replace the op with mean. + # stride = input_size // output_size = 4097 / 1 = 4097 + # kernel_size = input_size - (output_size - 1) * stride = 4097 - 0 * 4097 = 4097 + + model = AdaptiveAvgPool2dModule(output_size) + delegated_ep = to_quantized_edge_program( + model, input_shape, use_new_flow_neutron_c=True + ).exported_program() + + # Make sure the `adaptive_avg_pool2d` was NOT delegated. + assert not graph_contains_any_of_ops( + delegated_ep.graph, [ExecutorchDelegateCall] + ) + assert graph_contains_any_of_ops(delegated_ep.graph, [AdaptiveAvgPool2D]) diff --git a/backends/nxp/tests/model_output_comparator.py b/backends/nxp/tests/model_output_comparator.py index 4efa01011b8..f0dd7cd2d60 100644 --- a/backends/nxp/tests/model_output_comparator.py +++ b/backends/nxp/tests/model_output_comparator.py @@ -91,7 +91,13 @@ def compare_sample(self, sample_dir, cpu_output_tensors, npu_output_tensors): assert np.any( cpu_tensor ), "Output tensor contains only zeros. This is suspicious." - assert np.allclose(cpu_tensor, npu_tensor, atol=self.atol) + all_close = np.allclose(cpu_tensor, npu_tensor, atol=self.atol) + if not all_close: + max_diff = np.abs(cpu_tensor - npu_tensor).max() + print( + f"NPU output doesn't match reference. Maximum absolute difference: {max_diff}" + ) + assert all_close def _default_postprocess_fn(outputs: np.ndarray, _: str): diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index f190ca91e1f..3bf2e5bc2d9 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -12,6 +12,7 @@ from executorch.exir.dialects._ops import ops as exir_ops Abs = exir_ops.edge.aten.abs.default +AdaptiveAvgPool2D = exir_ops.edge.aten._adaptive_avg_pool2d.default AvgPool2D = exir_ops.edge.aten.avg_pool2d.default Bmm = exir_ops.edge.aten.bmm.default Convolution = exir_ops.edge.aten.convolution.default