Fix Concat Memory

Summary:
There is a bug with some of the tensor manipulation functions when we operate them in NHWC. We serialize the permuted tensor shape. But other arguments which reference dimensions within the shape are not properly permuted as well.

For example, if we are concating two tensors along dim 1. If we add a convolution before this node our memory format tagging will tag this node to operate in NHWC. As a result, we will actually be concatenating against dim 3 because the tensor was permuted.

Reviewed By: digantdesai

Differential Revision: D47413372

fbshipit-source-id: 3a1cd9b0ca678d2d673fbb09ca2a35888d61f986

Author: mcr229

backends/xnnpack/operators/op_cat.py

@@ -13,6 +13,7 @@
     XNode,
 )
 from executorch.backends.xnnpack.utils.quant_utils import QuantParams
+from executorch.backends.xnnpack.utils.utils import PERM_NHWC_TO_NCHW
 from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_INVALID_VALUE_ID
 
 
@@ -49,6 +50,9 @@ def define_node(
         if len(node.args) > 1:
             axis = cast(int, node.args[1])
 
+        if "XNN_NHWC_NODE" in node.meta:
+            axis = PERM_NHWC_TO_NCHW[axis]
+
         if num_tensors_to_cat == 2:
             xnode = XNNConcatenate2(
                 axis=axis,

backends/xnnpack/test/test_xnnpack.py

@@ -926,6 +926,47 @@ def forward(self, x, y):
             (torch.randn(1, 2, 3), torch.randn(1, 2, 5)),
         )
 
+    def test_xnnpack_backend_concatenate_nhwc(self):
+        class Concat(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(
+                    in_channels=1,
+                    out_channels=3,
+                    kernel_size=(3, 3),
+                    padding=1,
+                    bias=False,
+                )
+
+            def forward(self, x, y):
+                x = self.conv(x)
+                return torch.concatenate((y, x, y, x), 1)
+
+        self.lower_and_test_with_partitioner(
+            Concat(), (torch.randn(1, 1, 3, 3), torch.randn(1, 1, 3, 3))
+        )
+
+    def test_xnnpack_backend_concatenate_nhwc2(self):
+        class Concat(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(
+                    in_channels=1,
+                    out_channels=3,
+                    kernel_size=(3, 3),
+                    padding=1,
+                    bias=False,
+                )
+
+            def forward(self, x, y):
+                x = self.conv(x)
+                y = self.conv(y)
+                return torch.concatenate((y, x, y, x), 3)
+
+        self.lower_and_test_with_partitioner(
+            Concat(), (torch.randn(1, 1, 3, 3), torch.randn(1, 1, 3, 3))
+        )
+
     def test_xnnpack_backend_slice_copy(self):
         class Slice(torch.nn.Module):
             def forward(self, x):