From 4072ce5cfe4cf1a2a209856cc96aa4eb99d8307b Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Tue, 3 Feb 2026 09:13:30 +0100 Subject: [PATCH] NXP backend: Add support for `aten.upsample_nearest2d.vec`. --- backends/nxp/backend/edge_helper.py | 8 + .../nxp/backend/edge_program_converter.py | 3 +- .../ops_converters/__init__.py | 4 + .../upsample_nearest2d_converter.py | 107 +++++++++++ backends/nxp/backend/node_format_inference.py | 1 + backends/nxp/neutron_partitioner.py | 1 + backends/nxp/quantizer/neutron_quantizer.py | 2 + backends/nxp/quantizer/patterns.py | 9 + .../tests/test_convert_upsample_nearest2d.py | 181 ++++++++++++++++++ docs/source/backends/nxp/op-support.csv | 5 +- 10 files changed, 318 insertions(+), 3 deletions(-) create mode 100644 backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py create mode 100644 backends/nxp/tests/test_convert_upsample_nearest2d.py diff --git a/backends/nxp/backend/edge_helper.py b/backends/nxp/backend/edge_helper.py index 81b792b5298..9d97b952b24 100644 --- a/backends/nxp/backend/edge_helper.py +++ b/backends/nxp/backend/edge_helper.py @@ -354,3 +354,11 @@ def is_no_op_on_neutron(node: Node, parameters_mapping: dict[str, Parameter]) -> except Exception: # If execution fails, assume it's not a no-op. return False + + +def node_has_well_defined_shape(node: Node) -> bool: + if (val := node.meta.get("val")) is None: + # The node doesn't have a shape stored at all. + return False + + return all(isinstance(dim, int) and dim > 0 for dim in val.shape) diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index e6496e3eb16..0715500db27 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -42,12 +42,13 @@ exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405 exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter, # noqa F405 exir_ops.edge.aten.relu.default: ReLUConverter, # noqa F405 + exir_ops.edge.aten.sigmoid.default: SigmoidConverter, # noqa F405 exir_ops.edge.aten.slice_copy.Tensor: SliceTensorConverter, # noqa F405 exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 exir_ops.edge.aten.sub.Tensor: SubTensorConverter, # noqa F405 exir_ops.edge.aten.tanh.default: TanhConverter, # noqa F405 + exir_ops.edge.aten.upsample_nearest2d.vec: UpsampleNearest2DConverter, # noqa F405 exir_ops.edge.aten.view_copy.default: ViewCopyConverter, # noqa F405 - exir_ops.edge.aten.sigmoid.default: SigmoidConverter, # 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 3b8b9bf9b3f..19ca6c5346e 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 @@ -68,6 +68,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.tanh_converter import ( TanhConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.upsample_nearest2d_converter import ( + UpsampleNearest2DConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.view_copy_converter import ( ViewCopyConverter, ) @@ -97,5 +100,6 @@ "SoftmaxConverter", "SubTensorConverter", "TanhConverter", + "UpsampleNearest2DConverter", "ViewCopyConverter", ] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py new file mode 100644 index 00000000000..1ddc71425ef --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py @@ -0,0 +1,107 @@ +# 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 + +from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT +from executorch.backends.nxp.backend.edge_helper import node_has_well_defined_shape +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, + NodeConverter, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.resize_nearest_neighbor_options import ( + ResizeNearestNeighbor, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from torch.fx import Node +from torch.nn import Parameter + + +# noinspection SpellCheckingInspection +class UpsampleNearest2DConverter(NodeConverter): + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + + if node.meta.get(NXP_NODE_FORMAT, DataFormat.NONE) != DataFormat.CHANNELS_FIRST: + # This should never happen. + raise NotImplementedError( + "NXP backend: `aten.upsample_nearest2d.vec` didn't have correctly identified data" + " format. Please report this." + ) + + return True + + @staticmethod + def _is_supported_on_target( + node: Node, + neutron_target_spec: NeutronTargetSpec, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + # Neutron requires static shapes. + # neutron-converter/src/OperatorC/UpsamplePlugin.cpp?at=NEUTRON_SOFTWARE_2.2.3#74 + if not node_has_well_defined_shape(node): + return False + + if len(node.meta["val"].shape) != 4: + # Unexpected case. The input should always be 4D. + return False + + # The tensors here use the channels first format (NCHW). + _, in_c, in_h, in_w = node.all_input_nodes[0].meta["val"].shape + _, _, out_h, out_w = node.meta["val"].shape + + # Neutron supports only the doubling and quadrupleing of both height and width at the same time. + # neutron-library/src/utils/NeutronLibraryInterrogation.cpp?at=refs%2Ftags%2FNEUTRON_SOFTWARE_2.2.3#768 + # neutron-library/src/utils/NeutronLibraryInterrogation.cpp?at=refs%2Ftags%2FNEUTRON_SOFTWARE_2.2.3#778 + supported_scales = [2, 4] + if not any( + in_h * scale == out_h and in_w * scale == out_w + for scale in supported_scales + ): + return False + + # Neutron requires the input channels to be a multiple of `num_macs`. + # neutron-library/src/utils/NeutronLibraryInterrogation.cpp?at=refs%2Ftags%2FNEUTRON_SOFTWARE_2.2.3#767 + if in_c % neutron_target_spec.get_num_macs() != 0: + return False + + return True + + def convert(self, node: Node): + """Convert the `aten.upsample_nearest2d.vec` operator to Neutron IR `ResizeNearestNeighbor`. + The schema is: + aten::upsample_nearest2d.vec( + Tensor input, + SymInt[]? output_size, + float[]? scale_factors + ) -> Tensor + """ + self.assert_convertible(node) + + t_op = self._create_tflite_op_with_io_tensors(node) + x = t_op.tmp_inputs[0] + y = t_op.tmp_outputs[0] + + t_op.builtin_options = ResizeNearestNeighbor(False, False) + + # The `aten.upsample_nearest2d` can use either the `size` attribute or the `scale_factor` to define the output + # size. The Neutron IR `ResizeNearestNeighbor` only supports the `sizes` (output spatial dimensions). + # Both `size` and `scale_factor` can be easily supported by extracting the output spatial size from the output + # tensor's shape and using it as the `sizes`. + # The `self.assert_convertible(node)` call guarantees that the shape is 4D, channels last (NHWC), and static. + _, out_h, out_w, _ = y.shape + sizes = self.builder.create_tensor_for_data( + np.array([out_h, out_w], np.int32), "sizes" + ) + + t_op.tmp_inputs = [x, sizes] # Assign the NeutronIR inputs. + + 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 9154434f44b..a9e952746d1 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -27,6 +27,7 @@ class NodeFormatInference: exir_ops.edge.aten.convolution.default: {"inputs": [0, 1]}, exir_ops.edge.aten.max_pool2d_with_indices.default: {"inputs": [0]}, exir_ops.edge.aten.max_pool2d.default: {"inputs": [0]}, + exir_ops.edge.aten.upsample_nearest2d.vec: {"inputs": [0]}, } # A set of Edge Aten ops, which have the ability to change the format (for example - input nodes diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index ec0dd656e77..3fce757604e 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -220,6 +220,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]): exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 exir_ops.edge.aten.sub.Tensor: SubTensorConverter, # noqa F405 exir_ops.edge.aten.tanh.default: TanhConverter, # noqa F405 + exir_ops.edge.aten.upsample_nearest2d.vec: UpsampleNearest2DConverter, # noqa F405 exir_ops.edge.aten.view_copy.default: ViewCopyConverter, # noqa F405 } diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 62774fcb51d..ec91d31cdc5 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -47,6 +47,7 @@ TanhInPlacePattern, TanhPattern, TransposeIntPattern, + UpsampleNearest2DPattern, ViewPattern, ) from executorch.backends.nxp.quantizer.utils import ( @@ -273,6 +274,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) OpQuantizer(TanhPattern(is_qat=is_qat), static_qconfig), OpQuantizer(TanhInPlacePattern(is_qat=is_qat), static_qconfig), OpQuantizer(TransposeIntPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(UpsampleNearest2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(ViewPattern(is_qat=is_qat), static_qconfig), ] ) diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index f8ab4dd26e8..3bd3dbffeb3 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -944,6 +944,15 @@ def get_anchors( ) +class UpsampleNearest2DPattern(SharedSpecPattern): + """ + Quantizer for `aten.upsample_nearest2d.vec` operator. + """ + + def partition_types(self): + return [torch.ops.aten.upsample_nearest2d.vec] + + class ActivationsConcatClusterPattern(QuantizationPattern): """ Quantizer for activations concat cluster pattern. diff --git a/backends/nxp/tests/test_convert_upsample_nearest2d.py b/backends/nxp/tests/test_convert_upsample_nearest2d.py new file mode 100644 index 00000000000..4c1901d4f3c --- /dev/null +++ b/backends/nxp/tests/test_convert_upsample_nearest2d.py @@ -0,0 +1,181 @@ +# 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 +import pytest +import torch + +from executorch.backends.nxp.backend.edge_program_converter import ( + EdgeProgramToIRConverter, +) +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.exir.dialects._ops import ops as exir_ops + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(42) + np.random.seed(23) + + +# noinspection PyProtectedMember +ExecutorchDelegateCall = torch._higher_order_ops.executorch_call_delegate +UpsampleNearest2D = exir_ops.edge.aten.upsample_nearest2d.vec + + +class UpsampleNearestModule(torch.nn.Module): + + def __init__(self, size=None, scale=None): + super().__init__() + self.upsample = torch.nn.Upsample(size=size, scale_factor=scale, mode="nearest") + + def forward(self, x): + return self.upsample(x) + + +@pytest.mark.parametrize( + "input_shape, size", + [ + pytest.param((1, 8, 2, 3), (4, 6), id="2x upscale, 8 channels, tuple size"), + pytest.param((1, 8, 3, 3), 6, id="2x upscale, 8 channels, scalar size"), + pytest.param((1, 8, 2, 3), (8, 12), id="4x upscale, 8 channels, tuple size"), + pytest.param((1, 8, 3, 3), 12, id="4x upscale, 8 channels, scalar size"), + ], +) +def test_convert_upsample_nearest2d__size(mocker, input_shape, size): + model = UpsampleNearestModule(size=size) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + delegated_ep = to_quantized_edge_program( + model, input_shape, use_neutron_for_format_conversion=False + ).exported_program() + + # Make sure the `upsample` was delegated. + assert graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert not graph_contains_any_of_ops(delegated_ep.graph, [UpsampleNearest2D]) + + # Verify correct behavior of the converted NeutronIR model. + intermediate_ep = converter_spy.call_args.args[1] + neutron_ir_model, _ = converter_spy.spy_return + + input_data = ( + np.random.random(input_shape).astype(np.float32) * 256.0 - 128.0 + ).astype(np.int8) + + # Make sure the tested program contains the `upsample`. + assert graph_contains_any_of_ops(intermediate_ep.graph, [UpsampleNearest2D]) + + convert_run_compare( + intermediate_ep, + tfl_model=neutron_ir_model, + input_data=input_data, + tflite_input_preprocess=ToChannelLastPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), + ) + + +@pytest.mark.parametrize( + "input_shape, scale_factor", + [ + pytest.param((1, 8, 2, 3), 2, id="2x upscale, 8 channels, scalar scale"), + pytest.param((1, 8, 3, 3), 2.0, id="2x upscale, 8 channels, float scale"), + pytest.param((1, 8, 4, 5), (2, 2), id="2x upscale, 8 channels, tuple scale"), + pytest.param((1, 8, 2, 3), 4, id="4x upscale, 8 channels, scalar scale"), + pytest.param((1, 8, 2, 3), (4, 4), id="4x upscale, 8 channels, tuple scale"), + ], +) +def test_convert_upsample_nearest2d__scale_factor(mocker, input_shape, scale_factor): + model = UpsampleNearestModule(scale=scale_factor) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + delegated_ep = to_quantized_edge_program( + model, input_shape, use_neutron_for_format_conversion=False + ).exported_program() + + # Make sure the `upsample` was delegated. + assert graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert not graph_contains_any_of_ops(delegated_ep.graph, [UpsampleNearest2D]) + + # Verify correct behavior of the converted NeutronIR model. + intermediate_ep = converter_spy.call_args.args[1] + neutron_ir_model, _ = converter_spy.spy_return + + input_data = ( + np.random.random(input_shape).astype(np.float32) * 256.0 - 128.0 + ).astype(np.int8) + + # Make sure the tested program contains the `upsample`. + assert graph_contains_any_of_ops(intermediate_ep.graph, [UpsampleNearest2D]) + + convert_run_compare( + intermediate_ep, + tfl_model=neutron_ir_model, + input_data=input_data, + tflite_input_preprocess=ToChannelLastPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), + ) + + +def test_convert_upsample_nearest2d__no_delegation__unsupported_channels(): + size = 6 + input_shape = (1, 2, size // 2, size // 2) # 2 channels, not `num_macs`. + model = UpsampleNearestModule(size=size) + + delegated_ep = to_quantized_edge_program( + model, input_shape, use_neutron_for_format_conversion=False + ).exported_program() + + # Make sure the `upsample` was NOT delegated (channels != 8). + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [UpsampleNearest2D]) + + +@pytest.mark.parametrize( + "input_shape, scale_factor", + [ + pytest.param((1, 8, 4, 4), 3, id="3x upscale"), + pytest.param((1, 8, 4, 4), 1.5, id="1.5x upscale"), + pytest.param((1, 8, 4, 4), (2, 4), id="2x and 4x mixed upscale"), + pytest.param((1, 8, 10, 10), 1.99, id="1.99x upscale"), + ], +) +def test_convert_upsample_nearest2d__no_delegation__unsupported_scale( + input_shape, scale_factor +): + model = UpsampleNearestModule(scale=scale_factor) + + delegated_ep = to_quantized_edge_program( + model, input_shape, use_neutron_for_format_conversion=False + ).exported_program() + + # Make sure the `upsample` was NOT delegated (scale != 2). + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [UpsampleNearest2D]) + + +@pytest.mark.parametrize( + "input_shape, size", + [ + pytest.param((1, 8, 2, 3), (6, 9), id="3x upscale"), + pytest.param((1, 8, 2, 4), (3, 6), id="1.5x upscale"), + pytest.param((1, 8, 3, 4), 6, id="non-uniform upscale"), + ], +) +def test_convert_upsample_nearest2d__no_delegation__unsupported_size(input_shape, size): + model = UpsampleNearestModule(size=size) + + delegated_ep = to_quantized_edge_program( + model, input_shape, use_neutron_for_format_conversion=False + ).exported_program() + + # Make sure the `upsample` was NOT delegated (size != double of input). + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [UpsampleNearest2D]) diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index 581ec3ffb94..0a2de451b7a 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -15,7 +15,8 @@ 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.mm.default,int8,static int8,"2D tensor only" aten.relu.default,int8,static int8, -aten.tanh.default,int8,static int8, -aten.view_copy.default,int8,static int8, aten.sigmoid.default,int8,static int8, aten.slice_copy.Tensor, int8, static int8 +aten.tanh.default,int8,static int8, +aten.upsample_nearest2d.vec,int8,static int8,"channels % 8 = 0, H_scale = W_scale = 2 or 4" +aten.view_copy.default,int8,static int8,