Qualcomm AI Engine Direct - oss model enablement (EfficientSAM)

 - e2e script for https://github.com/yformer/EfficientSAM
 - Fastvit breakage fix
 - Passes order correction
 - Add support for cum_sum

Author: DannyYuyang-quic

backends/qualcomm/_passes/layout_transform.py

@@ -54,6 +54,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.ceil.default,
         exir_ops.edge.aten.clamp.default,
         exir_ops.edge.aten.constant_pad_nd.default,
+        exir_ops.edge.aten.cumsum.default,
         exir_ops.edge.aten.div.Tensor,
         exir_ops.edge.aten.elu.default,
         exir_ops.edge.aten.eq.Tensor,

backends/qualcomm/_passes/recompose_pixel_unshuffle.py

@@ -45,13 +45,11 @@ def call(self, graph_module: torch.fx.GraphModule):
                         continue
 
                     view_node = premute_node.args[0]
-                    if any(
-                        [
-                            view_node.op != "call_function",
-                            view_node.target != self.view_target,
-                            len(view_node.args[1]) != 6,
-                            len(premute_node.args[1]) != 6,
-                        ]
+                    if (
+                        view_node.op != "call_function"
+                        or view_node.target != self.view_target
+                        or len(view_node.args[1]) != 6
+                        or len(premute_node.args[1]) != 6
                     ):
                         continue
 

backends/qualcomm/_passes/utils.py

@@ -107,7 +107,7 @@ def get_passes_dependency_for_capture_program():
         ConvertToLinear: [RecomposePixelUnshuffle],
         DecomposeAny: [RemoveRedundancy],
         DecomposeLinalgVectorNorm: [RemoveRedundancy],
-        ExpandBroadcastTensorShape: [RemoveRedundancy],
+        ExpandBroadcastTensorShape: [ConstantI64toI32, TensorI64toI32],
         FoldQDQ: [AnnotateQuantAttrs, AnnotateDecomposed],
         LayoutTransform: [
             AnnotateQuantAttrs,

backends/qualcomm/builders/__init__.py

@@ -20,6 +20,7 @@
     op_clamp,
     op_conv2d,
     op_cos,
+    op_cum_sum,
     op_depth_to_space,
     op_dequantize,
     op_div,
@@ -106,6 +107,7 @@
     op_clamp,
     op_conv2d,
     op_cos,
+    op_cum_sum,
     op_depth_to_space,
     op_dequantize,
     op_div,

backends/qualcomm/builders/op_cos.py

@@ -3,7 +3,6 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-
 from typing import Dict
 
 import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper

backends/qualcomm/builders/op_cum_sum.py

@@ -0,0 +1,84 @@
+# 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.
+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
+
+from .node_visitor import NodeVisitor, register_node_visitor
+from .qnn_constants import OpCumulativeSum, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class CumulativeSum(NodeVisitor):
+    target = ["aten.cumsum.default"]
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def get_param(self, node, input_tensor):
+        dim = node.args[1]
+
+        if dim < 0:
+            dim = dim % len(input_tensor.shape)
+        if QCOM_AXIS_ORDER in node.meta:
+            dim = node.meta[QCOM_AXIS_ORDER].index(dim)
+
+        return cast(np.uint32, dim)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper],
+    ) -> PyQnnWrapper.PyQnnOpWrapper:
+        input_node = node.args[0]
+        input_tensor = self.get_tensor(input_node, node)
+        input_tensor_wrapper = self.define_tensor(
+            input_node,
+            node,
+            input_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+
+        dim = self.get_param(node, input_tensor)
+
+        output_tensor = self.get_tensor(node, node)
+        output_tensor_wrapper = self.define_tensor(
+            node,
+            node,
+            output_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+
+        cumsum_op = PyQnnWrapper.PyQnnOpWrapper(
+            node.name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpCumulativeSum.op_name,
+        )
+        cumsum_op.AddInputTensors([input_tensor_wrapper])
+        cumsum_op.AddOutputTensors([output_tensor_wrapper])
+        cumsum_op.AddScalarParam(
+            OpCumulativeSum.param_axis,
+            PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32,
+            {QCOM_DATA: dim},
+        )
+        cumsum_op.AddScalarParam(
+            OpCumulativeSum.param_exclusive,
+            PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_BOOL_8,
+            {QCOM_DATA: False},
+        )
+        cumsum_op.AddScalarParam(
+            OpCumulativeSum.param_reverse,
+            PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_BOOL_8,
+            {QCOM_DATA: False},
+        )
+
+        return cumsum_op

backends/qualcomm/builders/op_sin.py

@@ -3,7 +3,6 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-
 from typing import Dict
 
 import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper

backends/qualcomm/builders/qnn_constants.py

@@ -50,6 +50,14 @@ class OpConvert:
     op_name: str = "Convert"
 
 
+@dataclass(init=False, frozen=True)
+class OpCumulativeSum:
+    op_name = "CumulativeSum"
+    param_axis = "axis"
+    param_exclusive = "exclusive"
+    param_reverse = "reverse"
+
+
 @dataclass(init=False, frozen=True)
 class OpDepthToSpace:
     op_name: str = "DepthToSpace"

backends/qualcomm/quantizer/annotators.py

@@ -971,6 +971,11 @@ def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None
     )
 
 
+@register_annotator([torch.ops.aten.cumsum.default])
+def annotate_cumsum(node: Node, quantization_config: QuantizationConfig) -> None:
+    annotate_single_in_single_out(node, quantization_config)
+
+
 @register_annotator([torch.ops.aten.linear.default])
 def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None:
     act_node = node.args[0]

backends/qualcomm/tests/models.py

@@ -560,6 +560,14 @@ def forward(self, x):
         return torch.cos(x)
 
 
+class CumSum(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.cumsum(dim=0)
+
+
 class Div(torch.nn.Module):
     def __init__(self):
         super().__init__()