From 6bf8da0458421a0ab80608b322610ace51ee865d Mon Sep 17 00:00:00 2001 From: Vaclav Novak Date: Fri, 31 Jul 2026 18:22:40 +0200 Subject: [PATCH] feat: added support for `aten.pad` with mode `reflect` --- .../nxp/backend/edge_program_converter.py | 1 + .../ops_converters/__init__.py | 4 + .../ops_converters/pad_converter.py | 120 +++++++++++++++++ backends/nxp/neutron_partitioner.py | 89 +++++++------ backends/nxp/tests/executorch_pipeline.py | 6 +- .../test_constant_pad_nd_converter.py | 22 +-- .../node_converter/test_pad_converter.py | 126 ++++++++++++++++++ backends/nxp/tests/models.py | 49 +++---- backends/nxp/tests/ops_aliases.py | 1 + 9 files changed, 345 insertions(+), 73 deletions(-) create mode 100644 backends/nxp/backend/ir/converter/node_converters/ops_converters/pad_converter.py create mode 100644 backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index ad9dd2a49da..e400231146d 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -54,6 +54,7 @@ exir_ops.edge.aten.mm.default: MMConverter, # noqa F405 exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405 exir_ops.edge.aten.neg.default: NegConverter, # noqa F405 + exir_ops.edge.aten.pad.default: PadConverter, # noqa F405 exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter, # noqa F405 exir_ops.edge.aten.prelu.default: PReLUConverter, # noqa F405 exir_ops.edge.aten.relu.default: ReLUConverter, # 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 e21647f6e90..d32fdd08177 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 @@ -76,6 +76,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.neg_converter import ( NegConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.pad_converter import ( + PadConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.permute_copy_converter import ( PermuteCopyConverter, ) @@ -151,6 +154,7 @@ "MMConverter", "MulTensorConverter", "NegConverter", + "PadConverter", "PermuteCopyConverter", "PReLUConverter", "QDQPerChannelDequantizeConverter", diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/pad_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/pad_converter.py new file mode 100644 index 00000000000..75a0da3052c --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/pad_converter.py @@ -0,0 +1,120 @@ +# 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. + +from typing import Collection + +import numpy as np +import torch +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) + +from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( + apply_permutation_to, + create_channels_first_to_channels_last_permutation, +) +from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter +from executorch.backends.nxp.backend.ir.lib.tflite.MirrorPadMode import MirrorPadMode +from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( + mirror_pad_options, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from torch.fx import Node +from torch.nn import Parameter + + +class PadConverter(NodeConverter): + """Convert `aten.pad.default` to a TFLite padding operator. + + The ExecuTorch schema is: + aten::pad( + Tensor self, + SymInt[] pad, + str mode="constant", + float? value=None + ) -> Tensor + + Only mode "reflect" is handled in this converter. + Mode "constant" is decomposed to `aten.constant_pad_nd` and handled in its own converter. + Modes "replicate"/"circular" are not handled for now, such conversion to TFLite is not trivial. + """ + + @staticmethod + def _get_mode(node: Node) -> str: + return node.args[2] if len(node.args) > 2 else "constant" + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + mode = PadConverter._get_mode(node) + # `constant` mode is decomposed to `aten.constant_pad_nd`. + # Conversion of `replicate`/`circular` Torch padding to TFLite + # is more complicated and is skipped for now. + if mode != "reflect": + 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 _convert_paddings_to_tflite( + paddings: Collection[int], input_tensor: tflite_model.Tensor + ) -> list[int]: + # Group `padding` by two elements per list. + paddings_grouped = np.array(paddings).reshape(-1, 2) + + # In TFLite, `padding` order is reversed. + paddings_reversed = list(reversed(paddings_grouped)) + + # Add complementary zero pairs to `padding` to match input tensor rank. + zero_pair_compl = [[0, 0]] * (input_tensor.rank - len(paddings_reversed)) + padding_tfl = zero_pair_compl + paddings_reversed + + if input_tensor.tensor_format.is_channels_last(): + # Permute padding to match tensor format. + to_tflite_perm = create_channels_first_to_channels_last_permutation( + input_tensor.rank + ) + padding_tfl = apply_permutation_to(padding_tfl, to_tflite_perm) + + return padding_tfl + + def convert(self, node: Node): + """Convert `aten.pad.default` to a TFLite padding operator.""" + self.assert_convertible(node) + + t_op = self._create_tflite_op_with_io_tensors(node) + x = t_op.tmp_inputs[0] + + paddings = self._convert_paddings_to_tflite(node.args[1], x) + + paddings_tensor = self.builder.create_tensor_for_data( + np.asarray(paddings, "int32"), "paddings" + ) + + t_op.builtin_options = mirror_pad_options.MirrorPad(MirrorPadMode.REFLECT) + t_op.tmp_inputs = [x, paddings_tensor] + + self.builder.append_operators([t_op]) diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index e244cb22b57..59c590d5561 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -8,6 +8,7 @@ import logging import operator from dataclasses import dataclass +from functools import partial from typing import Callable, final, Mapping import torch @@ -227,6 +228,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]): exir_ops.edge.aten.mm.default: MMConverter, # noqa F405 exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405 exir_ops.edge.aten.neg.default: NegConverter, # noqa F405 + exir_ops.edge.aten.pad.default: PadConverter, # noqa F405 exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter, # noqa F405 exir_ops.edge.aten.prelu.default: PReLUConverter, # noqa F405 exir_ops.edge.aten.relu.default: ReLUConverter, # noqa F405 @@ -360,6 +362,11 @@ def __init__( self.preserve_ops = preserve_ops or [] self.check_op_support = check_op_support + # dict: exir_op -> converter + self.edge_op_to_converter = {} + for exir_op, converter in supported_ops.items(): + self.edge_op_to_converter[exir_op._op] = converter + @staticmethod def _partition_contains_compute_nodes( partition: Partition, parameters_mapping: dict[str, Parameter] @@ -501,57 +508,63 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: tagged_exported_program=exported_program, partition_tags=partition_tags ) + def _preserve_ops_filter( + self, + parameters_mapping: dict[str, Parameter], + node: torch.fx.Node, + ): + """Filter for `preserve_ops`. This filter makes `preserve_ops` a candidate list of ops that can be or cannot be preserved, + based on conditions below. + """ + # Make `preserve_ops` a candidate list, otherwise the conditions would be checked for all graph ops. + if node.target not in self.preserve_ops: + return False + + # If the converter for the node does not exist, the operator should be decomposed. + # When it does not have a converter, it usually means it is natively decomposed into simpler ops that have a converter. + node_converter = self.edge_op_to_converter.get(node.target) + if node_converter is None: + logger.w( + f"Node of target {str(node.target)} might be decomposed to other edge operators " + + "because no appropriate NodeConverter exists." + ) + return False + + # op must be decomposed if it is not delegable + delegable = node_converter.is_supported( + node, + self.neutron_target_spec, + parameters_mapping, + self.custom_delegation_options, + ) + + # if the node is delegable and the base check_op callable determines the node should not be decomposed, + # then do not decompose + base_check_op = ( + True if self.check_op_support is None else self.check_op_support(node) + ) + return delegable and base_check_op + def ops_to_not_decompose( self, ep: ExportedProgram, ) -> tuple[list[torch._ops.OpOverload], Callable[[torch.fx.Node], bool] | None]: """ - Method to determine which operators SHOULD NOT be decomposed to edge dialect. + Method to determine which operators SHOULD NOT be decomposed to simpler edge ops. - The method returns: - a list of operators NOT to decompose - optional callable - filter - If the operator is in the list and the evaluation of the callable is True (or the callable is not specified), - the operator should not be decomposed (i.e. it truly belongs to the list). Otherwise, it should be decomposed. + Returns: + List[torch._ops.OpOverload]: a list of ops that should not be decomposed. + Optional[Callable[[torch.fx.Node], bool]]]: an optional filter, called for each node in the + graph, that lets a node be decomposed even though its op is in the list above. A node is kept + (not decomposed) only if the filter returns True for it; if it returns False, the node is decomposed. """ # The parameters_mapping will only contain `FakeTensor` objects since the partitioning is done with a "Fake model". # If a node is selected to not be decomposed and it requires the real static parameters in its 'is_supported_*' methods, # the state dict from the quantized model will need to be provided here. parameters_mapping = EdgeProgramToIRConverter.map_inputs_to_parameters(ep) - aten_op_to_converter = {} - for exir_op, converter in supported_ops.items(): - aten_op_to_converter[exir_op._op] = converter # determine node format so it can be checked if the nodes are delegable NodeFormatInference(ep, only_for_op_support_check=True).identify_node_formats() - def check_op_support_extended(node: torch.fx.Node): - # if the operator should not be preserved, then it should be decomposed - if node.target not in self.preserve_ops: - return False - - # if the converter for the node does not exist, the operator should be decomposed - node_converter = aten_op_to_converter.get(node.target) - if node_converter is None: - logger.w( - f"Node of target {str(node.target)} might be decomposed to other edge operators " - + "because no appropriate NodeConverter exists." - ) - return False - - delegable = node_converter.is_supported( - node, - self.neutron_target_spec, - parameters_mapping, - self.custom_delegation_options, - ) - - # if the node is delegable and the base check_op callable determines the node should not be decomposed, - # then do not decompose - base_check_op = ( - True if self.check_op_support is None else self.check_op_support(node) - ) - return delegable and base_check_op - - return self.preserve_ops, check_op_support_extended + return self.preserve_ops, partial(self._preserve_ops_filter, parameters_mapping) diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py index a186535e53d..964b8de159c 100644 --- a/backends/nxp/tests/executorch_pipeline.py +++ b/backends/nxp/tests/executorch_pipeline.py @@ -216,7 +216,11 @@ def to_quantized_edge_program( ) # List of operators to not decompose during the lowering. - preserve_ops = [torch.ops.aten.prelu.default, torch.ops.aten.hardswish.default] + preserve_ops = [ + torch.ops.aten.prelu.default, + torch.ops.aten.pad.default, + torch.ops.aten.hardswish.default, + ] compile_spec = generate_neutron_compile_spec( target, intermediates_dir=intermediates_dir, diff --git a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py index 32bbf93fae4..b4a64447aa6 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py @@ -13,10 +13,7 @@ ModelBuilder, ) from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( - ConstantPadNDConvModule, - ConstantPadNDModule, -) +from executorch.backends.nxp.tests.models import PadConvModule, PadModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.ops_aliases import ConstantPadND, Convolution from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -78,8 +75,7 @@ def assert_delegated_and_output_shape_equals( def test__basic_nsys_inference( self, mocker, request, input_shape, paddings, use_qat ): - # These test cases are also supported by the old flow. - model = ConstantPadNDModule(paddings) + model = PadModule(paddings, "constant") self.assert_delegated(model, input_shape, mocker, request, use_qat) def test__channels_padding(self, mocker, request): @@ -87,7 +83,7 @@ def test__channels_padding(self, mocker, request): # These paddings will be applied to the last dimension, which is the channels as the input is formatless. paddings = (1, 1) expected_output_shape = (2, 4, 8) # Padded channels. - model = ConstantPadNDModule(paddings) + model = PadModule(paddings, "constant") self.assert_delegated_and_output_shape_equals( model, input_shape, expected_output_shape, mocker, request @@ -97,7 +93,7 @@ def test__batch_padding(self, mocker, request): input_shape = (2, 4, 6) paddings = (0, 0, 0, 0, 1, 1) # Padding applied to the batch dimension. expected_output_shape = (4, 4, 6) # Padded batch. - model = ConstantPadNDModule(paddings) + model = PadModule(paddings, "constant") self.assert_delegated_and_output_shape_equals( model, input_shape, expected_output_shape, mocker, request @@ -107,7 +103,7 @@ def test__batch_padding(self, mocker, request): def test__specific_constant(self, mocker, request, constant): input_shape = (2, 4, 6) paddings = (1, 1) - model = ConstantPadNDModule(paddings, constant) + model = PadModule(paddings, "constant", constant) self.assert_delegated(model, input_shape, mocker, request) @pytest.mark.parametrize( @@ -119,7 +115,13 @@ def test__specific_constant(self, mocker, request, constant): ], ) def test__channels_first(self, mocker, request, input_shape, paddings): - model = ConstantPadNDConvModule(paddings) + # compute channels size after padding + if len(paddings) // 2 == len(input_shape) - 1: + conv_in_channels = input_shape[1] + paddings[-1] + paddings[-2] + else: + conv_in_channels = input_shape[1] + + model = PadConvModule(conv_in_channels, paddings, "constant") graph_verifier = DetailedGraphVerifier( mocker, expected_delegated_ops={ConstantPadND: 1, Convolution: 1}, diff --git a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py new file mode 100644 index 00000000000..266260f9e1f --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py @@ -0,0 +1,126 @@ +# 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.ir.converter.builder.model_builder import ( + ModelBuilder, +) +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.models import PadConvModule, PadModule +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import Convolution, Pad + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(23) + np.random.seed(23) + + +class TestPadConverter: + """The PyTorch padding is added to the individual dimensions from the back (slightly confusing), see: + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html#torch.nn.functional.pad + + Current `pad_converter` currently converts only padding with mode `reflect`. Mode `constant` padding + is decomposed into `constant_pad_nd` and has its own converter. + + Padding with `reflect` mode has the following constraints: + - len(padding) = 2 ... only 2D/3D input + - len(padding) = 4 ... only 3D/4D input + - len(padding) = 6 ... only 4D/5D input + - len(padding) = 8 and higher not implemented in Torch + + Thus padding the first dim directly (without permute) is not possible. + Mainly, this constraint disables padding of the outer-most dimension (ie. batch). + """ + + # noinspection PyMethodMayBeStatic + def assert_delegated( + self, model, input_shape, mocker, request, exp_deleg_nodes=None, use_qat=False + ): + if exp_deleg_nodes is None: + exp_deleg_nodes = {Pad: 1} + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops=exp_deleg_nodes, + expected_non_delegated_ops={}, + ) + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + use_qat=use_qat, + ) + + def assert_delegated_and_output_shape_equals( + self, model, input_shape, expected_output_shape, mocker, request + ): + model_builder_spy = mocker.spy(ModelBuilder, "finish") + + self.assert_delegated(model, input_shape, mocker, request) + + neutron_ir_subgraph = model_builder_spy.call_args[0][0].get_sub_graph() + assert neutron_ir_subgraph.outputs.tmp_outputs[0].shape.vector == list( + expected_output_shape + ) + + @pytest.mark.parametrize( + "input_shape, paddings", + [ + pytest.param((3, 5), tuple(range(2)), id="2D, padding one dim"), + pytest.param((3, 3, 5), tuple(range(2)), id="3D, padding one dim"), + pytest.param((3, 7, 5), tuple(range(4)), id="3D, padding two dims"), + pytest.param((3, 3, 7, 5), tuple(range(4)), id="4D, padding two dims"), + pytest.param((3, 9, 7, 5), tuple(range(6)), id="4D, padding three dims"), + pytest.param((3, 3, 9, 7, 5), tuple(range(6)), id="5D, padding three dims"), + ], + ) + def test__basic__reflect(self, mocker, request, input_shape, paddings): + model = PadModule(paddings=paddings, mode="reflect") + + self.assert_delegated(model, input_shape, mocker, request) + + def test__channels_padding__reflect(self, mocker, request): + input_shape = (2, 4, 6) + # These paddings will be applied to the last dimension, which is the channels as the input is formatless. + paddings = (1, 1) + expected_output_shape = (2, 4, 8) # Padded channels. + model = PadModule(paddings, "reflect") + + self.assert_delegated_and_output_shape_equals( + model, input_shape, expected_output_shape, mocker, request + ) + + @pytest.mark.parametrize( + "input_shape, paddings", + [ + pytest.param((1, 10, 8, 6), tuple(range(4)), id="4D, padding H, W"), + pytest.param((1, 10, 8, 6), tuple(range(6)), id="4D, padding C, H, W"), + ], + ) + def test__channels_first__reflect(self, mocker, request, input_shape, paddings): + # compute channels size after padding + if len(paddings) // 2 == len(input_shape) - 1: + conv_in_channels = input_shape[1] + paddings[-1] + paddings[-2] + else: + conv_in_channels = input_shape[1] + model = PadConvModule(conv_in_channels, paddings, "reflect") + + self.assert_delegated( + model, input_shape, mocker, request, {Pad: 1, Convolution: 1} + ) + + def test__qat__reflect(self, mocker, request): + input_shape = (2, 4, 6) + paddings = (1, 1) + model = PadModule(paddings, "reflect") + + self.assert_delegated(model, input_shape, mocker, request, use_qat=True) diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py index bcf485acfe6..9f22519bc52 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/models.py @@ -174,17 +174,6 @@ def forward(self, x): return self.maxpool(x) -class Conv2dConstantPadNDModule(torch.nn.Module): - def __init__(self, paddings: Collection[int], constant: float | int | None = None): - super().__init__() - self.pad = ConstantPadNDModule(paddings, constant) - self.conv = Conv2dModule() - - def forward(self, x): - x = self.conv(x) - return self.pad(x) - - class SoftmaxModule(torch.nn.Module): def __init__(self, dim: int): super().__init__() @@ -372,26 +361,38 @@ def forward(self, x): return x -class ConstantPadNDModule(torch.nn.Module): - def __init__(self, paddings: Collection[int], constant: float | int | None = None): +class PadModule(torch.nn.Module): + def __init__( + self, + paddings: Collection[int], + mode: str = "constant", + value: float | None = None, + ): super().__init__() - self.paddings = paddings - self.constant = constant + self.paddings = tuple(paddings) + self.mode = mode + self.value = value def forward(self, x): - if self.constant is None: - return torch.nn.functional.pad(x, tuple(self.paddings), "constant") + if self.mode == "constant": + return torch.nn.functional.pad(x, self.paddings, self.mode, self.value) + elif self.mode == "reflect": + return torch.nn.functional.pad(x, self.paddings, self.mode) else: - return torch.nn.functional.pad( - x, tuple(self.paddings), "constant", self.constant - ) + raise ValueError("Unsupported pad mode.") -class ConstantPadNDConvModule(torch.nn.Module): - def __init__(self, paddings: Collection[int], constant: float | int | None = None): +class PadConvModule(torch.nn.Module): + def __init__( + self, + in_channels: int, + paddings: Collection[int], + mode: str = "constant", + value: float | None = None, + ): super().__init__() - self.pad = ConstantPadNDModule(paddings, constant) - self.conv = Conv2dModule() + self.pad = PadModule(paddings, mode, value) + self.conv = Conv2dModule(in_channels=in_channels) def forward(self, x): x = self.pad(x) diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index d71c5c9ab44..5ccc1d67de0 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -44,6 +44,7 @@ Minimum = exir_ops.edge.aten.minimum.default MulTensor = exir_ops.edge.aten.mul.Tensor Neg = exir_ops.edge.aten.neg.default +Pad = exir_ops.edge.aten.pad.default PermuteCopy = exir_ops.edge.aten.permute_copy.default Prelu = exir_ops.edge.aten.prelu.default QuantizePerChannel = exir_ops.edge.quantized_decomposed.quantize_per_channel.default