Qualcomm AI Engine Direct - GA MobileVit V2 (#11279)

### Summary
- GA Model: MobileVit V2 (**PLEASE use QNN2.29** when running this model
or else the accuracy is bad.)
- Support Fold/Unfold Operator (Aten is Col2Im/Im2Col) with passes

### Test plan
#### Accuracy
top1: ~50%
top5: ~85%

#### Speed
SM8750: 2.34ms/inf

### Script
python examples/qualcomm/oss_scripts/mobilevit_v2.py -b build-android -H
$HOST -s $DEVICE --dataset ../imagenet-mini/val/



### Test plan
E2E UT

Author: winskuo-quic

backends/qualcomm/_passes/__init__.py

@@ -13,6 +13,7 @@
 from .convert_square_to_pow import ConvertSquareToPow
 from .decompose_any import DecomposeAny
 from .decompose_cdist import DecomposeCDist
+from .decompose_col_im import DecomposeColIm
 from .decompose_einsum import DecomposeEinsum
 from .decompose_expm1 import DecomposeExpM1
 from .decompose_linalg_vector_norm import DecomposeLinalgVectorNorm
@@ -49,6 +50,7 @@
     ConvertSquareToPow,
     DecomposeAny,
     DecomposeCDist,
+    DecomposeColIm,
     DecomposeEinsum,
     DecomposeExpM1,
     DecomposeLinalgVectorNorm,

backends/qualcomm/_passes/decompose_col_im.py

@@ -0,0 +1,121 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+from .utils import copy_meta
+
+
+class DecomposeColIm(ExportPass):
+    """
+    Decompose im2col(unfold) to pixel_unshuffle + view_copy
+    Decompose col2im(fold) to view_copy + pixel_shuffle
+    """
+
+    def __init__(self):
+        super(DecomposeColIm, self).__init__()
+        self.im2col_op = exir_ops.edge.aten.im2col.default
+        self.col2im_op = exir_ops.edge.aten.col2im.default
+        self.pixel_unshuffle_op = exir_ops.edge.aten.pixel_unshuffle.default
+        self.pixel_shuffle_op = exir_ops.edge.aten.pixel_shuffle.default
+        self.view_copy_op = exir_ops.edge.aten.view_copy.default
+
+    def _decompose_im2col(self, graph_module: torch.fx.GraphModule):
+        for node in graph_module.graph.nodes:
+            if node.target == self.im2col_op:
+                input_node = node.args[0]
+                kernel_size = node.args[1]
+                stride = node.args[4]
+                batch_size = node.meta["val"].shape[0]
+                assert (
+                    stride == kernel_size
+                ), "im2col can only be converted when stride == kernel_size"
+                assert (
+                    input_node.meta["val"].dim() == 4
+                ), "im2col can only be converted when input dims == 4"
+                assert (
+                    kernel_size[0] == kernel_size[1]
+                ), "im2col can only be converted when kernel height == width"
+                users = list(node.users.keys())
+                with graph_module.graph.inserting_after(input_node):
+                    pixel_unshuffle_node = graph_module.graph.create_node(
+                        "call_function",
+                        self.pixel_unshuffle_op,
+                        (input_node, kernel_size[0]),
+                    )
+                    pixel_unshuffle_node.meta = copy_meta(node.meta)
+                    orig_height = input_node.meta["val"].shape[2]
+                    orig_width = input_node.meta["val"].shape[3]
+
+                    pixel_unshuffle_node.meta["val"] = pixel_unshuffle_node.meta[
+                        "val"
+                    ].reshape(
+                        batch_size,
+                        -1,
+                        orig_height // kernel_size[0],
+                        orig_width // kernel_size[1],
+                    )
+
+                    with graph_module.graph.inserting_after(pixel_unshuffle_node):
+                        view_copy_node = graph_module.graph.create_node(
+                            "call_function",
+                            self.view_copy_op,
+                            (pixel_unshuffle_node, tuple(node.meta["val"].shape)),
+                        )
+                        view_copy_node.meta = copy_meta(node.meta)
+                        for user in users:
+                            user.replace_input_with(node, view_copy_node)
+
+    def _decompose_col2im(self, graph_module: torch.fx.GraphModule):
+        for node in graph_module.graph.nodes:
+            if node.target == self.col2im_op:
+                input_node = node.args[0]
+                output_size = node.args[1]
+                kernel_size = node.args[2]
+                stride = node.args[5]
+                batch_size = node.meta["val"].shape[0]
+                assert (
+                    stride == kernel_size
+                ), "col2im can only be converted when stride == kernel_size"
+                assert (
+                    node.meta["val"].dim() == 4
+                ), "col2im can only be converted when output dims == 4"
+                assert (
+                    kernel_size[0] == kernel_size[1]
+                ), "col2im can only be converted when kernel height == width"
+                users = list(node.users.keys())
+                with graph_module.graph.inserting_after(input_node):
+                    view_tensor = input_node.meta["val"].reshape(
+                        batch_size,
+                        -1,
+                        output_size[0] // kernel_size[0],
+                        output_size[1] // kernel_size[1],
+                    )
+                    view_copy_node = graph_module.graph.create_node(
+                        "call_function",
+                        self.view_copy_op,
+                        (input_node, tuple(view_tensor.shape)),
+                    )
+                    view_copy_node.meta = copy_meta(node.meta)
+                    view_copy_node.meta["val"] = view_tensor
+
+                    with graph_module.graph.inserting_after(view_copy_node):
+                        pixel_shuffle_node = graph_module.graph.create_node(
+                            "call_function",
+                            self.pixel_shuffle_op,
+                            (view_copy_node, kernel_size[0]),
+                        )
+                        pixel_shuffle_node.meta = copy_meta(node.meta)
+
+                        for user in users:
+                            user.replace_input_with(node, pixel_shuffle_node)
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        self._decompose_im2col(graph_module)
+        self._decompose_col2im(graph_module)
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/qualcomm/_passes/layout_transform.py

@@ -23,8 +23,21 @@
 class LayoutTransform(ExportPass):
     """
     QNN delegate requires channel last layout format, this pass aims to
-    help generate the correct transformation by inserting fewest ammount of
+    help generate the correct transformation by inserting fewest amount of
     'permute' operators in the graph.
+    Please notice that permute op is inserted during qnn_preprocess.
+
+    Operations are divided into 3 categories: sensitive_layout, agnostic_layout, and pytorch_layout.
+    sensitive_layout: These ops must be lowered to QNN in NHWC format. A permute(NCHW->NHWC) op will be inserted in front of the sensitive_layout op.
+    agnostic_layout: These ops are agnostic to layout format, which means it can be passed to QNN in either NCHW or NHWC format.
+    pytorch_layout: These ops must be lowered to QNN in NCHW format. A permute(NHWC->NCHW) op will be inserted in front of the pytorch_layout op.
+
+    For optimization purposes, permute is only inserted when it is necessary to switch between sensitive_layout and pytorch_layout.
+    For example, a model consists of three kinds of operations: conv(sensitive_layout), relu(agnostic_layout), and unsqueeze(pytorch_layout)
+    If a graph originally looks like : in -> conv -> relu -> conv -> relu -> unsqueeze -> out
+    After layout_transform pass: in -> permute(NCHW->NHWC) -> conv -> relu -> conv -> relu -> permute(NHWC->NCHW) -> unsqueeze -> out
+    The reason for inserting the 1st permute is because conv is layout sensitive. Since relu is agnostic to layout, it doesn't matter what format is used.
+    This format works fine until unsqueeze is encountered, which is a pytorch_format operation, so a 2nd permute is necessary to convert it back to pytorch format.
     """
 
     layout_sensitive_ops = {
@@ -76,7 +89,6 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.logical_not.default,
         exir_ops.edge.aten.lt.Scalar,
         exir_ops.edge.aten.lt.Tensor,
-        exir_ops.edge.aten._log_softmax.default,
         exir_ops.edge.aten.maximum.default,
         exir_ops.edge.aten.mean.dim,
         exir_ops.edge.aten.minimum.default,
@@ -88,7 +100,6 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.prelu.default,
         exir_ops.edge.aten.repeat.default,
         exir_ops.edge.aten.relu.default,
-        exir_ops.edge.aten._softmax.default,  # TODO: Need to find a new solution to do "axis_order" to transform axis.
         exir_ops.edge.aten.sigmoid.default,
         exir_ops.edge.aten.split_with_sizes.default,
         exir_ops.edge.aten.split_with_sizes_copy.default,
@@ -282,11 +293,29 @@ def check_arg(arg):
                     else:
                         check_arg(args)
 
+    def conditional_sensitive_check(self, node):
+        # For softmax and log_softmax, we must ensure axis == -1 since thats the only axis supported by QNN.
+        # Softmax and log_softmax is treated as pytorch_layout in default, and will be treated as sensitive_layout when axis is not given as last dim.
+        target_nodes = [
+            exir_ops.edge.aten._softmax.default,
+            exir_ops.edge.aten._log_softmax.default,
+        ]
+        if node.target in target_nodes:
+            dim = node.args[1]
+            if dim < 0:
+                dim = dim % node.meta["val"].dim()
+            if dim != node.meta["val"].dim() - 1:
+                return True
+        return False
+
     def call(self, graph_module: torch.fx.GraphModule):
         graph = graph_module.graph
         sensitive_nodes = [
-            node for node in graph.nodes if self.is_layout_sensitive(node)
+            node
+            for node in graph.nodes
+            if self.is_layout_sensitive(node) or self.conditional_sensitive_check(node)
         ]
+
         # perform first run traversal for identifying nodes subjected to layout changes
         if self.insert_permute:
             self.insert_permute, self.transformed_tag = False, QCOM_LAYOUT_CHANGE

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -18,6 +18,7 @@
     ConvertSquareToPow,
     DecomposeAny,
     DecomposeCDist,
+    DecomposeColIm,
     DecomposeEinsum,
     DecomposeExpM1,
     DecomposeLinalgVectorNorm,
@@ -82,6 +83,7 @@ def get_capture_program_passes():
         (ConvertBmmToMatmul, True),
         (ConvertConv1dToConv2d, True),
         (DecomposeAny, True),
+        (DecomposeColIm, True),
         (ExpandBroadcastTensorShape, False),
         (FixedLinearKeepDim, True),
         (FoldQDQ, True),

backends/qualcomm/_passes/utils.py

@@ -67,6 +67,7 @@ def get_passes_dependency_for_capture_program():
         ConvertBmmToMatmul,
         ConvertConv1dToConv2d,
         DecomposeAny,
+        DecomposeColIm,
         DecomposeLinalgVectorNorm,
         ExpandBroadcastTensorShape,
         FixedLinearKeepDim,
@@ -91,6 +92,7 @@ def get_passes_dependency_for_capture_program():
         AnnotateUnbind: [RemoveRedundancy],
         ConvertBmmToMatmul: [RecomposePixelUnshuffle],
         DecomposeAny: [RemoveRedundancy],
+        DecomposeColIm: [FoldQDQ],
         DecomposeLinalgVectorNorm: [RemoveRedundancy],
         ExpandBroadcastTensorShape: [FoldQDQ],
         FixedLinearKeepDim: [FoldQDQ],

backends/qualcomm/builders/op_index_put.py

@@ -41,7 +41,7 @@ def define_node(
         indices_qnn = torch.cat(indices_unpacked).unsqueeze(0)
         indice_node = [n for n in indicies_node if isinstance(n, torch.fx.Node)]
         # TODO consider to write a pass to combine to one input tensor for indices
-        assert len(indice_node) == 1, "Not support mutilple indices tensor"
+        assert len(indice_node) == 1, "Not support multiple indices tensor"
 
         indices_tensor_wrapper = self.define_tensor(
             indice_node[0],

backends/qualcomm/builders/op_log_softmax.py

@@ -3,10 +3,10 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+import warnings
 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, QCOM_DATA
@@ -58,6 +58,10 @@ def define_node(
 
         # logsoftmax only supports last dimension for now, which is channel in QNN
         if dim != input_tensor.dim() - 1:
+            warnings.warn(
+                "[QNN Delegate Op Builder]: LogSoftmax only supports channel axis.",
+                stacklevel=1,
+            )
             return None
 
         log_softmax_op = PyQnnWrapper.PyQnnOpWrapper(

backends/qualcomm/builders/op_softmax.py

@@ -3,10 +3,10 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+import warnings
 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, QCOM_DATA
@@ -53,9 +53,12 @@ def define_node(
             dim = dim % len(input_tensor.shape)
         if QCOM_AXIS_ORDER in node.meta:
             dim = node.meta[QCOM_AXIS_ORDER].index(dim)
-
         # softmax only supports last dimension for now, which is channel in QNN
         if dim != input_tensor.dim() - 1:
+            warnings.warn(
+                "[QNN Delegate Op Builder]: Softmax only supports channel axis.",
+                stacklevel=1,
+            )
             return None
 
         softmax_op = PyQnnWrapper.PyQnnOpWrapper(

backends/qualcomm/partition/utils.py

@@ -42,9 +42,11 @@ def filter_fn(node: torch.fx.Node) -> bool:
 def get_skip_decomp_table() -> List[torch._ops.OperatorBase]:
     do_not_decompose = [
         torch.ops.aten.adaptive_avg_pool2d.default,
+        torch.ops.aten.col2im.default,
         torch.ops.aten.elu.default,
         torch.ops.aten.hardsigmoid.default,
         torch.ops.aten.hardswish.default,
+        torch.ops.aten.im2col.default,
         torch.ops.aten.instance_norm.default,
         torch.ops.aten.leaky_relu.default,
         torch.ops.aten.linear.default,

backends/qualcomm/quantizer/annotators.py

@@ -399,6 +399,11 @@ def annotate_cos(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)
 
 
+@register_annotator([torch.ops.aten.col2im.default, torch.ops.aten.im2col.default])
+def annotate_col_im(node: Node, quantization_config: QuantizationConfig) -> None:
+    annotate_single_in_single_out(node, quantization_config)
+
+
 @register_annotator([torch.ops.aten.sin.default])
 def annotate_sin(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)
@@ -508,7 +513,13 @@ def annotate_prelu(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_binary(node, quantization_config)
 
 
-@register_annotator([torch.ops.aten.view.default, torch.ops.aten._unsafe_view.default])
+@register_annotator(
+    [
+        torch.ops.aten.view_copy.default,
+        torch.ops.aten.view.default,
+        torch.ops.aten._unsafe_view.default,
+    ]
+)
 def annotate_view(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_in_out_obs_sharing_op(node, quantization_config)
     if not _is_annotated([node]):