NXP backend: Enable constant_pad_nd with new Neutron flow.

Author: MartinPavella

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

@@ -7,6 +7,7 @@
 from typing import Collection
 
 import numpy as np
+import torch
 
 from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT
 
@@ -41,19 +42,33 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        paddings = node.args[1]
-        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
+        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
 
-        return True
+            return True
+
+        else:
+            paddings = node.args[1]
+            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
 
     @staticmethod
     def _is_supported_in_IR(
@@ -110,7 +125,14 @@ def _convert_paddings_to_tflite(
         return paddings
 
     def convert(self, node: Node):
-        """Convert the `aten.constant_pad_nd` operator to TFLite `PadV2`."""
+        """Convert the `aten.constant_pad_nd` operator to NeutronIR `PadV2`.
+        The ExecuTorch schema is:
+            constant_pad_nd(
+                Tensor self,
+                SymInt[] pad,
+                Scalar value=0
+            ) -> Tensor
+        """
         self.assert_convertible(node)
 
         t_op = self._create_tflite_op_with_io_tensors(node)

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

@@ -4,9 +4,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import numpy as np
+
+# noinspection PyUnusedImports
 import pytest
 import torch
+
 from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig
+from executorch.backends.nxp.backend.ir.converter.builder.model_builder import (
+    ModelBuilder,
+)
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.constant_pad_nd_converter import (
     ConstantPadNDConverter,
 )
@@ -17,16 +23,18 @@
 from executorch.backends.nxp.tests.executors import (
     convert_run_compare,
     graph_contains_any_of_ops,
+    OverrideTargetSupportCheck,
     ToNCHWPreprocess,
     ToNHWCPreprocess,
 )
+from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier
 from executorch.backends.nxp.tests.models import (
     ConstantPadNDConvModule,
     ConstantPadNDModule,
 )
+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
-from executorch.backends.nxp.tests.executors import OverrideTargetSupportCheck
-from executorch.exir.dialects._ops import ops as exir_ops
 
 
 @pytest.fixture(autouse=True)
@@ -158,9 +166,8 @@ def test_constant_pad_nd__unsupported_paddings(input_shape, paddings, use_qat):
         model, input_shape, use_qat=use_qat
     ).exported_program()
 
-    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)
+    assert graph_contains_any_of_ops(exec_program.graph, [ConstantPadND])
 
 
 @pytest.mark.xfail(reason="EIEX=855")
@@ -173,9 +180,7 @@ def test_constant_pad_nd__delegation__formatless__supported_padding(use_qat):
     ).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]
-    )
+    assert not graph_contains_any_of_ops(exec_program.graph, [ConstantPadND])
 
 
 def test_constant_pad_nd__delegation__formatless__unsupported_padding(use_qat):
@@ -187,9 +192,7 @@ def test_constant_pad_nd__delegation__formatless__unsupported_padding(use_qat):
     ).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]
-    )
+    assert graph_contains_any_of_ops(exec_program.graph, [ConstantPadND])
 
 
 @pytest.mark.xfail(reason="Regression in Neutron SW 3.0.1 (AIR-14264)", strict=True)
@@ -202,9 +205,7 @@ def test_constant_pad_nd__delegation__channels_first__supported_padding(use_qat)
     ).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]
-    )
+    assert not graph_contains_any_of_ops(exec_program.graph, [ConstantPadND])
 
 
 def test_constant_pad_nd__delegation__channels_first__unsupported_padding(use_qat):
@@ -216,6 +217,122 @@ def test_constant_pad_nd__delegation__channels_first__unsupported_padding(use_qa
     ).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]
+    assert graph_contains_any_of_ops(exec_program.graph, [ConstantPadND])
+
+
+class TestConstantPadNDNewNeutronFlow:
+    """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
+    """
+
+    # noinspection PyMethodMayBeStatic
+    def assert_delegated(self, model, input_shape, mocker, use_qat=False):
+        graph_verifier = DetailedGraphVerifier(
+            mocker,
+            expected_delegated_ops={ConstantPadND: 1},
+            expected_non_delegated_ops={},
+        )
+
+        lower_run_compare(
+            model,
+            input_shape,
+            graph_verifier,
+            use_qat=use_qat,
+            use_new_flow_neutron_c=True,
+        )
+
+    def assert_delegated_and_output_shape_equals(
+        self, model, input_shape, expected_output_shape, mocker
+    ):
+        model_builder_spy = mocker.spy(ModelBuilder, "finish")
+
+        self.assert_delegated(model, input_shape, mocker)
+
+        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((2,), tuple(range(2)), id="1D, padding H"),
+            pytest.param((2, 4), tuple(range(2)), id="2D, padding H"),
+            pytest.param((2, 4), tuple(range(4)), id="2D, padding N, H"),
+            pytest.param((2, 4, 6), tuple(range(2)), id="3D, padding H"),
+            pytest.param((2, 4, 6), tuple(range(4)), id="3D, padding C, H"),
+            pytest.param((2, 4, 6, 8), tuple(range(2)), id="4D, padding W"),
+            pytest.param((2, 4, 6, 8), tuple(range(4)), id="4D, padding H, W"),
+            pytest.param((1, 2, 3, 4, 5), tuple(range(2)), id="5D, padding D"),
+            pytest.param((1, 2, 3, 4, 5), tuple(range(4)), id="5D, padding W, D"),
+        ],
+    )
+    def test__basic_nsys_inference(self, mocker, input_shape, paddings, use_qat):
+        # These test cases are also supported by the old flow.
+        model = ConstantPadNDModule(paddings)
+        self.assert_delegated(model, input_shape, mocker, use_qat)
+
+    def test__channels_padding(self, mocker):
+        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)
+
+        self.assert_delegated_and_output_shape_equals(
+            model, input_shape, expected_output_shape, mocker
+        )
+
+    def test__batch_padding(self, mocker):
+        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)
+
+        self.assert_delegated_and_output_shape_equals(
+            model, input_shape, expected_output_shape, mocker
+        )
+
+    @pytest.mark.parametrize("constant", [0.0, -13.37])
+    def test__specific_constant(self, mocker, constant):
+        input_shape = (2, 4, 6)
+        paddings = (1, 1)
+        model = ConstantPadNDModule(paddings, constant)
+        self.assert_delegated(model, input_shape, mocker)
+
+    @pytest.mark.parametrize(
+        "input_shape, paddings",
+        [
+            pytest.param((1, 4, 6, 8), tuple(range(2)), id="4D, padding W"),
+            pytest.param((1, 4, 6, 8), tuple(range(4)), id="4D, padding H, W"),
+        ],
     )
+    def test__channels_first(self, mocker, input_shape, paddings):
+        model = ConstantPadNDConvModule(paddings)
+        graph_verifier = DetailedGraphVerifier(
+            mocker,
+            expected_delegated_ops={ConstantPadND: 1, Convolution: 1},
+            expected_non_delegated_ops={},
+        )
+
+        lower_run_compare(
+            model, input_shape, graph_verifier, use_new_flow_neutron_c=True
+        )
+
+    @pytest.mark.xfail(
+        strict=True,
+        raises=RuntimeError,
+        reason="Known issue in Neutron: https://jira.sw.nxp.com/browse/AIR-14624",  # @lint-ignore
+    )
+    def test__bugged_channels_first_case(self, mocker):
+        input_shape, paddings = (1, 2, 6, 8), (0, 1, 2, 3, 1, 1)
+        model = ConstantPadNDConvModule(paddings)
+        graph_verifier = DetailedGraphVerifier(
+            mocker,
+            expected_delegated_ops={ConstantPadND: 1, Convolution: 1},
+            expected_non_delegated_ops={},
+        )
+
+        lower_run_compare(
+            model, input_shape, graph_verifier, use_new_flow_neutron_c=True
+        )

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

backends/nxp/tests/ops_aliases.py

@@ -14,6 +14,7 @@
 Abs = exir_ops.edge.aten.abs.default
 AvgPool2D = exir_ops.edge.aten.avg_pool2d.default
 Bmm = exir_ops.edge.aten.bmm.default
+ConstantPadND = exir_ops.edge.aten.constant_pad_nd.default
 Convolution = exir_ops.edge.aten.convolution.default
 DequantizePerChannel = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default
 DequantizePerTensor = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default