Layout transform

Author: winskuo-quic

backends/qualcomm/_passes/layout_transform.py

@@ -79,6 +79,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.elu.default,
         exir_ops.edge.aten.eq.Tensor,
         exir_ops.edge.aten.exp.default,
+        exir_ops.edge.aten.flip.default,
         exir_ops.edge.aten.floor.default,
         exir_ops.edge.aten.floor_divide.default,
         exir_ops.edge.aten.full.default,
@@ -111,6 +112,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.round.default,
         exir_ops.edge.aten.sigmoid.default,
         exir_ops.edge.aten.sign.default,
+        exir_ops.edge.aten.slice_copy.Tensor,
         exir_ops.edge.aten.split_with_sizes.default,
         exir_ops.edge.aten.split_with_sizes_copy.default,
         exir_ops.edge.aten.sqrt.default,

backends/qualcomm/builders/op_flip.py

@@ -10,6 +10,8 @@
 import numpy as np
 import torch
 
+from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER
+
 from .node_visitor import NodeVisitor
 from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpStridedSlice, QNN_OP_PACKAGE_NAME_QTI_AISW
@@ -47,11 +49,14 @@ def define_node(
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
         )
-
         ranges = []
 
+        dims = node.args[1]
+        if QCOM_AXIS_ORDER in node.meta:
+            dims = [node.meta[QCOM_AXIS_ORDER].index(dim) for dim in dims]
+
         for dim, size in enumerate(output_tensor.shape):
-            if dim in node.args[1]:
+            if dim in dims:
                 ranges.extend([size - 1, -1, -1])
             else:
                 ranges.extend([0, size, 1])

backends/qualcomm/builders/op_slice_copy.py

@@ -6,9 +6,9 @@
 from typing import cast, Dict
 
 import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper
-
 import numpy as np
 import torch
+from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER
 
 from .node_visitor import NodeVisitor
 from .node_visitor_manager import register_node_visitor
@@ -47,8 +47,9 @@ def define_node(
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
         )
-
         dim = cast(int, node.args[1])
+        if QCOM_AXIS_ORDER in node.meta:
+            dim = node.meta[QCOM_AXIS_ORDER].index(dim)
         if dim < 0:
             dim = dim % len(input_tensor.shape)
 
@@ -62,7 +63,6 @@ def define_node(
                 end = end % input_tensor.shape[dim]
         else:
             end = input_tensor.shape[dim]
-
         input_tensor_rank = len(input_tensor.shape)
         ranges = []
         for i in range(input_tensor_rank):

backends/qualcomm/tests/models.py

@@ -646,6 +646,40 @@ def forward(self, x):
         return self.conv_transpose(self.conv(x))
 
 
+class Conv2dFlip(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            in_channels=16,
+            out_channels=16,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            bias=False,
+        )
+        self.dims = [1, 3]
+
+    def forward(self, x):
+        x = self.conv(x)
+        return torch.flip(x, self.dims)
+
+
+class Conv2dSliceCopy(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            in_channels=1,
+            out_channels=4,
+            kernel_size=(3, 3),
+            padding=1,
+            bias=True,
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x[:, 2:, :, :]
+
+
 class Conv2dSumReduceDim(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -823,19 +857,6 @@ def forward(self, x):
         return torch.flip(x, self.dims)
 
 
-class FlipDecomp(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.dims = [0, 2]
-
-    def forward(self, x):
-        for dim in self.dims:
-            idx = torch.arange(start=x.size(dim) - 1, end=-1, step=-1)
-            # Select using reverse index, equivalent to flipping.
-            x = torch.index_select(x, dim, idx)
-        return x
-
-
 class Floor(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1346,11 +1346,21 @@ def test_qnn_backend_conv2d_down_up_sample(self):
         sample_input = (torch.randn(1, 16, 224, 224),)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_conv2d_flip(self):
+        module = Conv2dFlip()  # noqa: F405
+        sample_input = (torch.randn(1, 16, 224, 224),)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_conv2d_max_pool2d(self):
         module = Conv2dMaxPool2d()  # noqa: F405
         sample_input = (torch.rand(1, 2, 14, 14),)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_conv2d_slice_copy(self):
+        module = Conv2dSliceCopy()  # noqa: F405
+        sample_input = (torch.randn([2, 1, 3, 3]),)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_conv2d_sum_reduce_dim(self):
         module = Conv2dSumReduceDim()  # noqa: F405
         sample_input = (torch.randn([1, 1, 3, 3]),)
@@ -2955,12 +2965,24 @@ def test_qnn_backend_conv2d_down_up_sample(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_conv2d_flip(self):
+        module = Conv2dFlip()  # noqa: F405
+        sample_input = (torch.randn(1, 16, 224, 224),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_conv2d_max_pool2d(self):
         module = Conv2dMaxPool2d()  # noqa: F405
         sample_input = (torch.rand(1, 2, 14, 14),)
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_conv2d_slice_copy(self):
+        module = Conv2dSliceCopy()  # noqa: F405
+        sample_input = (torch.randn([2, 1, 3, 3]),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_conv2d_sum_reduce_dim(self):
         module = Conv2dSumReduceDim()  # noqa: F405
         sample_input = (torch.randn([1, 1, 3, 3]),)