NXP backend: Fix delegation check for aten.cat.default.

Author: MartinPavella

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

@@ -10,6 +10,7 @@
 )
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
 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 (
@@ -85,9 +86,29 @@ def _is_supported_on_target(
 
         dim = CatConverter._get_normalized_dim(node)
 
-        # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1491
-        if dim == 0:
-            return False
+        # There is a bug in the NeutronConverter, where if none of the dimensions before the one referenced by
+        #  `dim` are `!= 1`, the `Concat` is not delegated.
+        # This only happens when the inputs to the `Concat` are model inputs, and not outputs of other
+        #  operators. However, such a requirement cannot be currently enforced during partitioning, so we must
+        #  take the conservative approach and decline delegation.
+        input_shapes = [list(n.meta["val"].shape) for n in node.args[0]]
+        if node.meta[NXP_NODE_FORMAT].is_channels_first():
+            # Transform the shapes to channels last.
+            to_nhwc_perm = create_channels_first_to_channels_last_permutation(
+                len(node.meta["val"].shape), True
+            )
+            input_shapes = [
+                apply_permutation_to(shape, to_nhwc_perm)
+                for shape in input_shapes
+            ]
+
+            # Transform the `dim` to refer to a channels last dimension.
+            dim = to_nhwc_perm.index(dim)
+
+        for input_shape in input_shapes:
+            if not any(d != 1 for d in input_shape[:dim]):
+                # Do not delegate if there are no "non-1" dimensions in the shape before the `dim` dimension.
+                return False
 
         # Neutron requires the channels to be a multiple of `num_macs`. The channels could either be the second or the
         #  last dimension, depending on the formats of the node.

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

@@ -73,7 +73,7 @@ def forward(self, *inputs: torch.Tensor):
     ],
 )
 def test_cat__same_shapes(dim, num_inputs, rank, mocker):
-    input_shape = tuple([2, 8, 8, 8, 8][-rank:])
+    input_shape = tuple([8, 8, 8, 8][:rank])
 
     converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
 
@@ -134,11 +134,20 @@ def test_cat__channels_first__same_shapes(dim, num_inputs, mocker):
     )
 
 
-@pytest.mark.parametrize("dim", [0, -4])
-@pytest.mark.parametrize("num_inputs", [2])
-def test_cat__unsupported_dim__imxrt700(dim, num_inputs):
-    input_shape = (2, 8, 6, 8)
-
+@pytest.mark.parametrize(
+    "dim, input_shape",
+    [
+        pytest.param(0, (1, 8, 8, 8), id="axis = 0"),
+        pytest.param(0, (8, 8, 8, 8), id="axis = 0, no `1s` in the shape."),
+        pytest.param(-4, (1, 8, 8, 8), id="axis = -4"),
+        pytest.param(1, (1, 1, 8, 8), id="axis = 1"),
+        pytest.param(-3, (1, 1, 8, 8), id="axis = -3"),
+        pytest.param(2, (1, 1, 1, 8), id="axis = 2"),
+        pytest.param(-2, (1, 1, 1, 8), id="axis = -2"),
+    ],
+)
+def test_cat__unsupported__imxrt700(dim, input_shape):
+    num_inputs = 2
     quantized_program = to_quantized_edge_program(
         CatModule(dim), [input_shape] * num_inputs, target="imxrt700"
     ).exported_program()