NXP backend: Improve constant_pad_nd delegation by using inferred node formats.

Author: MartinPavella

backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py

@@ -27,6 +27,8 @@
     pad_options,
     pad_v2_options,
 )
+
+from executorch.backends.nxp.backend.node_format_inference import NXP_NODE_FORMAT
 from torch.fx import Node
 from torch.nn import Parameter
 
@@ -41,11 +43,17 @@ def _is_supported_on_target(
     ) -> bool:
         match target:
             case Target.RT700:
-                # TODO: Consider different tensor formats (dim-order)
                 paddings = node.args[1]
-                if len(paddings) > 4 and paddings[4:6] != [0, 0]:
-                    # Attempt to Pad channels dimension, which is not supported on Neutron.
-                    return False
+                if node.meta[NXP_NODE_FORMAT].is_channels_first():
+                    # Dim `1` will end up being the channels. It is padded by paddings[4:6].
+                    if len(paddings) > 4 and paddings[4:6] != [0, 0]:
+                        # Attempt to Pad channels dimension -> currently not supported
+                        return False
+                else:
+                    # Dim `-1` will end up being the channels. It is padded by paddings[:2].
+                    if len(paddings) > 0 and paddings[:2] != [0, 0]:
+                        # Attempt to Pad channels dimension -> currently not supported
+                        return False
 
                 return True
 
@@ -71,10 +79,6 @@ def _is_supported_in_IR(
         if not NodeConverter._has_shared_q_params_if_quantized(node):
             return False
 
-        if len(paddings) > 4 and paddings[4:6] != [0, 0]:
-            # Attempt to Pad channels dimension -> currently not supported
-            return False
-
         return True
 
     # noinspection PyMethodMayBeStatic

backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py

@@ -1,4 +1,4 @@
-# Copyright 2024 NXP
+# Copyright 2024-2025 NXP
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -13,13 +13,15 @@
 )
 from executorch.backends.nxp.tests.executors import (
     convert_run_compare,
+    graph_contains_any_of_ops,
     ToNCHWPreprocess,
     ToNHWCPreprocess,
 )
 from executorch.backends.nxp.tests.models import (
     ConstantPadNDConvModule,
     ConstantPadNDModule,
 )
+from executorch.exir.dialects._ops import ops as exir_ops
 
 
 @pytest.fixture(autouse=True)
@@ -121,3 +123,51 @@ def test_constant_pad_nd__unsupported_paddings(input_shape, paddings):
     nodes = list(exec_program.graph.nodes)
     # There is at least one non-delegated Pad node
     assert any(node.name == "aten_constant_pad_nd_default" for node in nodes)
+
+
+def test_constant_pad_nd__delegation__formatless__supported_padding():
+    input_shape = (2, 4, 6, 8)  # Formatless -> the last dim (8) will be padded.
+    paddings = [0, 0, 1, 2, 3, 4]  # The last dim is padded using the first 2 paddings.
+    model = ConstantPadNDModule(paddings)
+    exec_program = to_quantized_edge_program(model, input_shape).exported_program()
+
+    # Make sure the `pad` was delegated.
+    assert not graph_contains_any_of_ops(
+        exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default]
+    )
+
+
+def test_constant_pad_nd__delegation__formatless__unsupported_padding():
+    input_shape = (2, 4, 6, 8)  # Formatless -> the last dim (8) will be padded.
+    paddings = [0, 1]  # The last dim is padded using the first 2 paddings.
+    model = ConstantPadNDModule(paddings)
+    exec_program = to_quantized_edge_program(model, input_shape).exported_program()
+
+    # Make sure the `pad` was NOT delegated.
+    assert graph_contains_any_of_ops(
+        exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default]
+    )
+
+
+def test_constant_pad_nd__delegation__channels_first__supported_padding():
+    input_shape = (2, 4, 6, 8)  # Channels first -> the second dim (4) will be padded.
+    paddings = [1, 2, 3, 4, 0, 0]  # The second dim is padded using the paddings[4:6].
+    model = ConstantPadNDConvModule(paddings)
+    exec_program = to_quantized_edge_program(model, input_shape).exported_program()
+
+    # Make sure the `pad` was delegated.
+    assert not graph_contains_any_of_ops(
+        exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default]
+    )
+
+
+def test_constant_pad_nd__delegation__channels_first__unsupported_padding():
+    input_shape = (2, 3, 6, 8)  # Channels first -> the second dim (3) will be padded.
+    paddings = [0, 0, 0, 0, 1, 0]  # The second dim is padded using the paddings[4:6].
+    model = ConstantPadNDConvModule(paddings)
+    exec_program = to_quantized_edge_program(model, input_shape).exported_program()
+
+    # Make sure the `pad` was NOT delegated.
+    assert graph_contains_any_of_ops(
+        exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default]
+    )