Qualcomm AI Engine Direct - LPAI Support Partition

Author: winskuo-quic

backends/qualcomm/_passes/__init__.py

@@ -49,6 +49,7 @@
 from .insert_reshape_for_reduce_ops import InsertReshapeForReduceOps
 from .layout_transform import LayoutTransform
 from .lift_constant_scalar_operands import LiftConstantScalarOperands
+from .lpai_partition_fallback_support import LpaiPartitionFallbackSupport
 from .recompose_pad_maxpool2d import RecomposePadMaxPool2d
 from .recompose_pixel_unshuffle import RecomposePixelUnshuffle
 from .recompose_rms_norm import RecomposeRmsNorm
@@ -107,6 +108,7 @@
     InsertRequantize,
     LayoutTransform,
     LiftConstantScalarOperands,
+    LpaiPartitionFallbackSupport,
     RecomposePadMaxPool2d,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,

backends/qualcomm/_passes/fold_qdq.py

@@ -100,52 +100,49 @@ def _fold_q(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
 
     def _preserve_qdq(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
         for n in graph_module.graph.nodes:
-            # skip parameters & buffers
-            if n.target in dq_ops and is_parameter(n.args[0], self.edge_program):
-                self._annotate_bypass(n)
-                continue
-
-            # TODO: In LPAI backend v6, there is an accuracy drop for the quantize and dequantize operations.
-            # To address this, keep the quantize/dequantize operations at the model's input and output.
-            # For example, input -> q_1 (Fallback) -> dq_1 (Bypass) -> graph -> q_2 (Bypass) -> dq_2 (Fallback) -> output
-            # Here, q_1 and dq_2 will fallback to CPU, while q_2 and dq_1 will be bypassed in qnn_partition and folded in qnn_preprocess.
-            if self.backend_type == QnnExecuTorchBackendType.kLpaiBackend:
-                if (
-                    is_graph_input(n, self.edge_program)
-                    # For tagged quantized I/O, we should not fallback quantize operation.
-                    and QCOM_QUANTIZED_IO not in n.meta
-                ):
-                    user_list = list(n.users.keys())
-                    if len(user_list) > 0:
-                        q_node = user_list[0]
-                        q_node.meta[QCOM_FALLBACK_NODE] = True
-                        # Annotate the q_node since it will serve as the input for the first node during operator validation
-                        q_node.meta[QCOM_QUANT_ATTRS] = get_quant_attrs(
-                            self.edge_program, q_node
-                        )
-                        q_node.meta[QCOM_QUANTIZED_IO] = q_node.args[-1]
-                        dq_node = list(q_node.users.keys())[0]
-                        # Bypass dequantize op for graph validation by torch
-                        dq_node.meta[QCOM_BYPASS_NODE] = True
-                        # Make sure that the quantize operator isn't inserted for input in insert_io_qdq.py
-                        n.meta[QCOM_QUANTIZED_IO] = q_node.args[-1]
-                elif (
-                    is_graph_output(n)
-                    and n.target in dq_ops
-                    # For tagged quantized I/O, we should not fallback dequantize operation.
-                    and QCOM_QUANTIZED_IO not in n.args[0].args[0].meta
-                ):
-                    n.meta[QCOM_FALLBACK_NODE] = True
-                    q_node = n.args[0]
-                    # Bypass quantize op for graph validation by torch
-                    q_node.meta[QCOM_BYPASS_NODE] = True
-                    op_node = q_node.args[0]
-                    # Make sure that the dequantize operator isn't inserted for output in insert_io_qdq.py
-                    op_node.meta[QCOM_QUANTIZED_IO] = q_node.args[-1]
+            if (
+                is_graph_input(n, self.edge_program)
+                # For tagged quantized I/O, we should not fall back quantize operation.
+                and QCOM_QUANTIZED_IO not in n.meta
+            ):
+                user_list = list(n.users.keys())
+                if len(user_list) > 0:
+                    q_node = user_list[0]
+                    q_node.meta[QCOM_FALLBACK_NODE] = True
+                    # Annotate the q_node since it will serve as the input for the first node during operator validation
+                    q_node.meta[QCOM_QUANT_ATTRS] = get_quant_attrs(
+                        self.edge_program, q_node
+                    )
+                    dq_node = list(q_node.users.keys())[0]
+                    # Bypass dequantize op for graph validation by torch
+                    dq_node.meta[QCOM_BYPASS_NODE] = True
+            elif (
+                is_graph_output(n)
+                and n.target in dq_ops
+                # For tagged quantized I/O, we should not fall back dequantize operation.
+                and QCOM_QUANTIZED_IO not in n.args[0].args[0].meta
+            ):
+                n.meta[QCOM_FALLBACK_NODE] = True
+                q_node = n.args[0]
+                # Bypass quantize op for graph validation by torch
+                q_node.meta[QCOM_BYPASS_NODE] = True
 
     def call(self, graph_module: torch.fx.GraphModule):
         if not self.force_fold:
-            self._preserve_qdq(graph_module)
+            for node in graph_module.graph.nodes:
+                # skip parameters & buffers
+                if node.target in dq_ops and is_parameter(
+                    node.args[0], self.edge_program
+                ):
+                    self._annotate_bypass(node)
+                    continue
+
+            if self.backend_type == QnnExecuTorchBackendType.kLpaiBackend:
+                # In LPAI backend v6, there is an accuracy drop for the quantize and dequantize operations.
+                # To address this, keep the quantize/dequantize operations at the model's input and output.
+                # For example, input -> q_1 (Fall back) -> dq_1 (Bypass) -> graph -> q_2 (Bypass) -> dq_2 (Fall back) -> output
+                # Here, q_1 and dq_2 will fall back to CPU, while q_2 and dq_1 will be bypassed in qnn_partition and folded in qnn_preprocess.
+                self._preserve_qdq(graph_module)
         self._fold_dq(graph_module)
         self._fold_q(graph_module)
         dead_code_elimination_pass(graph_module)

backends/qualcomm/_passes/insert_io_qdq.py

@@ -3,11 +3,14 @@
 #
 # 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 torch
+from executorch.backends.qualcomm._passes.utils import (
+    insert_dequant_node,
+    insert_quant_node,
+)
 
-from executorch.backends.qualcomm.builders.node_visitor import q_ops
+from executorch.backends.qualcomm.builders.node_visitor import q_dq_map, q_ops
 
 from executorch.backends.qualcomm.builders.utils import (
     is_mutable_buffer_input,
@@ -18,7 +21,6 @@
     QCOM_QUANT_ATTRS,
     QCOM_QUANTIZED_IO,
 )
-from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.passes import dead_code_elimination_pass
 
@@ -31,98 +33,10 @@ class InsertIOQDQ(ExportPass):
     right before outputs according to stored quantization encodings.
     """
 
-    q_dq_map = {
-        # per tensor (quantize -> dequantize)
-        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
-        exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
-        # per tensor (dequantize -> dequantize, for pre-quantized params)
-        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
-        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor,
-        # per channel (quantize -> dequantize)
-        exir_ops.edge.quantized_decomposed.quantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
-        # per channel (dequantize -> dequantize, for pre-quantized params)
-        exir_ops.edge.quantized_decomposed.dequantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
-    }
-
     def __init__(self, edge_program: torch.export.ExportedProgram):
         super(InsertIOQDQ, self).__init__()
         self.edge_program = edge_program
 
-    def _ceate_args(self, target: torch.fx.node.Target, quant_attrs: Dict):
-        ret = []
-
-        arg_schemas = list(target._schema.arguments)[1:]
-        for arg_schema in arg_schemas:
-            name = arg_schema.name
-            # TODO: Due to the new parameter "out_dtype" in the dequantize node,
-            # it could not be found in the quant_attrs of other nodes,
-            # and it will cause a key error. For now, the output type
-            # of our dequantize node is only float. (by default in pytorch)
-            if name == "out_dtype":
-                continue
-            value = quant_attrs[name]
-            if isinstance(arg_schema.type, torch.Tensor) and (
-                isinstance(value, int) or isinstance(value, float)
-            ):
-                value = torch.tensor(value)
-            ret.append(value)
-        return ret
-
-    def _create_node(
-        self,
-        graph_module: torch.fx.GraphModule,
-        node: torch.fx.node,
-        target: torch.fx.node.Target,
-        quant_attrs: Dict = None,
-    ) -> torch.fx.node:
-        # check if there has a specified quant_attrs
-        # if not, use the existent info. from current node
-        if quant_attrs is None:
-            quant_attrs = node.meta.get(QCOM_QUANT_ATTRS)
-
-        inserted_node = graph_module.graph.create_node(
-            "call_function",
-            target,
-            (node, *self._ceate_args(target, quant_attrs)),
-        )
-        meta_val = node.meta["val"]
-        if target in q_ops:
-            inserted_node.meta[QCOM_QUANT_ATTRS] = node.meta.pop(QCOM_QUANT_ATTRS)
-            meta_val = meta_val.to(quant_attrs["dtype"])
-
-        inserted_node.meta["val"] = meta_val
-        return inserted_node
-
-    def _insert_quant_node(
-        self,
-        graph_module: torch.fx.GraphModule,
-        node: torch.fx.node,
-        target: torch.fx.node.Target,
-        quant_attrs: Dict = None,
-    ) -> torch.fx.Node:
-        with graph_module.graph.inserting_after(node):
-            users = list(node.users.keys())
-            inserted_node = self._create_node(graph_module, node, target, quant_attrs)
-            for user in users:
-                # If we found mix quantization pattern and reuse the existing q_node, we skip adding a new q node.
-                if user.target not in q_ops:
-                    user.replace_input_with(node, inserted_node)
-
-        return inserted_node
-
-    def _insert_dequant_node(
-        self,
-        graph_module: torch.fx.GraphModule,
-        node: torch.fx.node,
-        target: torch.fx.node.Target,
-    ) -> None:
-        with graph_module.graph.inserting_after(node):
-            users = list(node.users.keys())
-            inserted_node = self._create_node(graph_module, node, target)
-            for user in users:
-                if user.op == "output":
-                    user.replace_input_with(node, inserted_node)
-
     def _insert(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
         # Snapshot nodes: inserting Q/DQ nodes mutates the graph's linked list,
         # so iterating the live list can revisit newly inserted nodes.
@@ -140,20 +54,30 @@ def _insert(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
                     or is_mutable_buffer_input(n, self.edge_program)
                 )
             ):
-                self._insert_quant_node(
-                    graph_module, n, n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]
-                )
+                # Insert a single Q node and connect single Q node to all users
+                users = list(n.users.keys())
+                if users:
+                    input_q_node = insert_quant_node(
+                        graph_module=graph_module,
+                        input_node=n,
+                        output_node=users[0],
+                        target=n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING],
+                    )
+                    for user in users[1:]:
+                        if user.target not in q_ops:
+                            user.replace_input_with(n, input_q_node)
 
             # insert dq before output or fold mix_quantization q if applicable
-            users = list(n.users.keys())
-            if n.meta.get(QCOM_QUANT_ATTRS) and any(
-                user.op == "output" for user in users
-            ):
-                self._insert_dequant_node(
-                    graph_module,
-                    n,
-                    self.q_dq_map[n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]],
-                )
+            users = [user for user in n.users if user.op == "output"]
+            if n.meta.get(QCOM_QUANT_ATTRS) and len(users) != 0:
+                # We should always only have 1 output node. Expect len(users) == 1
+                for user in users:
+                    _ = insert_dequant_node(
+                        graph_module=graph_module,
+                        input_node=n,
+                        output_node=user,
+                        target=q_dq_map[n.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]],
+                    )
 
     def call(self, graph_module: torch.fx.GraphModule):
         self._insert(graph_module)