NXP backend: added support for aten.pad with mode reflect

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

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",

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])

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)

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,

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,16 +75,15 @@ 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):
         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 = 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},