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

@@ -63,16 +63,10 @@ def test_constant_pad_nd_conversion__default_constant():
         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), list(range(6)), id="3D, padding N, 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((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), tuple(range(2)), id="5D, padding D"),
         pytest.param((1, 2, 3, 4, 5), tuple(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), tuple(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):
@@ -88,13 +82,16 @@ def test_constant_pad_nd_conversion__format_less(input_shape, paddings):
 @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"),
+        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)