Fix Slice 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 arguemtns which reference dimensions within the shape are not properly permuted as well.

For example, if we are slicing 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 are actually slicing dim 3 because the tensor was permuted.

Reviewed By: digantdesai

Differential Revision: D47413374

fbshipit-source-id: d6bd661a30d084a48756a9f06fc5d8d16e12aebd

Author: mcr229

backends/xnnpack/operators/op_slice_copy.py

@@ -10,7 +10,12 @@
     XNNStaticSlice,
     XNode,
 )
-from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node
+from executorch.backends.xnnpack.utils.utils import (
+    check_or_raise,
+    get_input_node,
+    PERM_NCHW_TO_NHWC,
+    PERM_NHWC_TO_NCHW,
+)
 
 
 @register_node_visitor
@@ -52,8 +57,13 @@ def define_node(
         )
         output_tensor_val = node.meta["val"]
         output_shape = list(output_tensor_val.shape)
-
         dim_of_slice = cast(int, node.args[1])
+
+        if "XNN_NHWC_NODE" in node.meta:
+            input_shape = [input_shape[i] for i in PERM_NCHW_TO_NHWC]
+            output_shape = [output_shape[i] for i in PERM_NCHW_TO_NHWC]
+            dim_of_slice = PERM_NHWC_TO_NCHW[dim_of_slice]
+
         slice_begin_index = cast(int, node.args[2])
         if slice_begin_index < 0:
             slice_begin_index = input_shape[dim_of_slice] + slice_begin_index

backends/xnnpack/test/test_xnnpack.py

@@ -955,3 +955,21 @@ def forward(self, x):
             self.lower_module_and_test_output(
                 Slice(), (torch.randn(5, 5, 5),), use_partitioner=True
             )
+
+    def test_xnnpack_backend_slice_copy_memory_format(self):
+        class ConvSlice(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 = self.conv(x)
+                return y[:, :, 2:3, -2:]
+
+        self.lower_and_test_with_partitioner(ConvSlice(), (torch.randn(1, 1, 3, 3),))