NXP backend: Improve partitioning and conversion of Pad op

Author: skywall

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

@@ -9,7 +9,6 @@
 import numpy as np
 
 from executorch.backends.nxp.backend.edge_helper import input_rank
-from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
 from executorch.backends.nxp.backend.ir.converter.conversion.translator import (
     apply_permutation_to,
     create_channels_first_to_channels_last_permutation,
@@ -24,6 +23,7 @@
 )
 from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
+    pad_options,
     pad_v2_options,
 )
 from torch.fx import Node
@@ -50,6 +50,10 @@ 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
@@ -101,6 +105,15 @@ def convert(self, node: Node):
             np.asarray(paddings, "int32"), "paddings"
         )
 
+        if constant == 0.0:
+            # We're padding with zeros, we can use traditional Pad op
+            t_op.tmp_inputs = [x, paddings_tensor]
+            t_op.tmp_outputs = [y]
+            t_op.builtin_options = pad_options.Pad()
+
+            self.builder.append_operators([t_op])
+            return
+
         if x.quantization is None:
             constant_tensor = self.builder.create_tensor_for_data(
                 np.array([constant], tf_lite_type_to_numpy(x.type)), "constant"
@@ -124,6 +137,4 @@ def convert(self, node: Node):
         t_op.tmp_outputs = [y]
         t_op.builtin_options = pad_v2_options.PadV2()
 
-        ops_to_add = OpsList(middle_op=t_op)
-
-        self.builder.append_operators(ops_to_add.flatten())
+        self.builder.append_operators([t_op])

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

@@ -7,7 +7,10 @@
 import pytest
 import torch
 
-from executorch.backends.nxp.tests.executorch_pipeline import to_edge_program
+from executorch.backends.nxp.tests.executorch_pipeline import (
+    to_edge_program,
+    to_quantized_edge_program,
+)
 from executorch.backends.nxp.tests.executors import (
     convert_run_compare,
     ToNCHWPreprocess,
@@ -60,16 +63,10 @@ def test_constant_pad_nd_conversion__default_constant():
         pytest.param([2, 4], list(range(4)), id="2D, padding N, H"),
         pytest.param([2, 4, 6], list(range(2)), id="3D, padding H"),
         pytest.param([2, 4, 6], list(range(4)), id="3D, padding C, H"),
-        pytest.param([2, 4, 6], list(range(6)), id="3D, padding N, C, H"),
         pytest.param([2, 4, 6, 8], list(range(2)), id="4D, padding W"),
         pytest.param([2, 4, 6, 8], list(range(4)), id="4D, padding H, W"),
-        pytest.param([2, 4, 6, 8], list(range(6)), id="4D, padding C, H, W"),
-        pytest.param([2, 4, 6, 8], list(range(8)), id="4D, padding N, C, H, W"),
         pytest.param([1, 2, 3, 4, 5], list(range(2)), id="5D, padding D"),
         pytest.param([1, 2, 3, 4, 5], list(range(4)), id="5D, padding W, D"),
-        pytest.param([1, 2, 3, 4, 5], list(range(6)), id="5D, padding H, W, D"),
-        pytest.param([1, 2, 3, 4, 5], list(range(8)), id="5D, padding C, H, W, D"),
-        pytest.param([1, 2, 3, 4, 5], list(range(10)), id="5D, padding N, C, H, W, D"),
     ],
 )
 def test_constant_pad_nd_conversion__format_less(input_shape, paddings):
@@ -87,13 +84,12 @@ def test_constant_pad_nd_conversion__format_less(input_shape, paddings):
     [
         pytest.param([1, 4, 6, 8], list(range(2)), id="4D, padding W"),
         pytest.param([1, 4, 6, 8], list(range(4)), id="4D, padding H, W"),
-        pytest.param([1, 1, 6, 8], [1, 2, 3, 4, 2, 1], id="4D, padding C, H, W"),
-        # pytest.param([1, 1, 6, 8], [1, 2, 3, 4, 2, 1, 5, 6], id='4D, padding N, C, H, W'),  # Batch size must stay 0.
     ],
 )
 def test_constant_pad_nd_conversion__channels_first(input_shape, paddings):
+    model = ConstantPadNDConvModule(paddings)
     edge_program = to_edge_program(
-        ConstantPadNDConvModule(paddings), input_shape
+        model, input_shape
     ).exported_program()  # Extra `Conv` after the padding.
 
     input_data = np.random.random(input_shape).astype(np.float32)
@@ -104,3 +100,24 @@ def test_constant_pad_nd_conversion__channels_first(input_shape, paddings):
         tflite_input_preprocess=ToNHWCPreprocess(),
         tflite_output_preprocess=ToNCHWPreprocess(),
     )
+
+
+@pytest.mark.parametrize(
+    "input_shape, paddings",
+    [
+        pytest.param([2, 4, 6], list(range(6)), id="3D, padding N, C, H"),
+        pytest.param([2, 4, 6, 8], list(range(6)), id="4D, padding C, H, W"),
+        pytest.param([2, 4, 6, 8], list(range(8)), id="4D, padding N, C, H, W"),
+        pytest.param([1, 2, 3, 4, 5], list(range(6)), id="5D, padding H, W, D"),
+        pytest.param([1, 2, 3, 4, 5], list(range(8)), id="5D, padding C, H, W, D"),
+        pytest.param([1, 2, 3, 4, 5], list(range(10)), id="5D, padding N, C, H, W, D"),
+        pytest.param([1, 1, 6, 8], [1, 2, 3, 4, 2, 1], id="4D, padding C, H, W"),
+    ],
+)
+def test_constant_pad_nd__unsupported_paddings(input_shape, paddings):
+    model = ConstantPadNDModule(paddings)
+    exec_program = to_quantized_edge_program(model, input_shape).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)