Qualcomm AI Engine Direct - LPAI Support Partition

Author: winskuo-quic

backends/qualcomm/_passes/__init__.py

@@ -52,6 +52,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
@@ -113,6 +114,7 @@
     InsertRequantize,
     LayoutTransform,
     LiftConstantScalarOperands,
+    LpaiPartitionFallbackSupport,
     RecomposePadMaxPool2d,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,

backends/qualcomm/_passes/backends/lpai/__init__.py

@@ -4,10 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .fold_qdq import LpaiFoldQDQ
 from .qnn_lpai_pass_manager import QnnLpaiPassManager
 
 __all__ = [
-    LpaiFoldQDQ,
     QnnLpaiPassManager,
 ]

backends/qualcomm/_passes/backends/lpai/fold_qdq.py

@@ -5,12 +5,19 @@
 # LICENSE file in the root directory of this source tree.
 
 from executorch.backends.qualcomm._passes import (
+    ConvertMhaToSha,
     DecomposeHardsigmoid,
     DecomposeReciprocal,
     FoldQDQ,
+    FuseConsecutiveCast,
+    FuseConsecutiveTranspose,
+    InsertRequantize,
+    LayoutTransform,
+    LpaiPartitionFallbackSupport,
     RemoveRedundancy,
+    ResolveDebugHandle,
+    TagQuantIO,
 )
-from executorch.backends.qualcomm._passes.backends.lpai.fold_qdq import LpaiFoldQDQ
 from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager
 
 
@@ -24,37 +31,38 @@ class QnnLpaiPassManager(QnnPassManager):
     @classmethod
     def get_default_pass_activations(cls):
         pass_activations = super().get_default_pass_activations()
-        pass_activations = [
-            (LpaiFoldQDQ if p is FoldQDQ else p, act) for p, act in pass_activations
-        ]
         # Hardsigmoid and Reciprocal no longer appear at to_edge stage as it is decomposed in the export/annotation pipeline.
         # The current change is intended to proactively prepare for the upcoming deprecation of the export pipeline.
         pass_activations.extend(
             [
                 (DecomposeHardsigmoid, True),
                 (DecomposeReciprocal, True),
+                (LpaiPartitionFallbackSupport, True),
             ]
         )
         return pass_activations
 
     @classmethod
     def get_passes_dependency_for_capture_program(cls):
         deps = super().get_passes_dependency_for_capture_program()
-        # Replace FoldQDQ with LpaiFoldQDQ in the dependency table
-        if FoldQDQ in deps:
-            deps[LpaiFoldQDQ] = deps.pop(FoldQDQ)
-        for key in deps:
-            deps[key] = [LpaiFoldQDQ if v is FoldQDQ else v for v in deps[key]]
         # Hardsigmoid and Reciprocal no longer appear at to_edge stage as it is decomposed in the export/annotation pipeline.
         # The current change is intended to proactively prepare for the upcoming deprecation of the export pipeline.
         deps.update(
             {
                 DecomposeHardsigmoid: [RemoveRedundancy],
                 DecomposeReciprocal: [RemoveRedundancy],
+                LpaiPartitionFallbackSupport: [TagQuantIO],
+                ResolveDebugHandle: [LpaiPartitionFallbackSupport],
             }
         )
         return deps
 
+    def _validate_edge_passes(self) -> None:
+        super()._validate_edge_passes()
+        assert isinstance(
+            self.passes[-2], LpaiPartitionFallbackSupport
+        ), "Please ensure LpaiPartitionFallbackSupport is the last edge pass before ResolveDebugHandle."
+
     @classmethod
     def get_annotation_passes(cls):
         passes = [DecomposeHardsigmoid, DecomposeReciprocal]
@@ -74,5 +82,14 @@ def get_export_passes(
 
     @classmethod
     def get_preprocess_passes(cls, use_mha2sha=False):
-        passes = super().get_preprocess_passes(use_mha2sha)
-        return [LpaiFoldQDQ if p is FoldQDQ else p for p in passes]
+        passes = [
+            FoldQDQ,
+            ConvertMhaToSha,
+            InsertRequantize,
+            LayoutTransform,
+            FuseConsecutiveCast,
+            FuseConsecutiveTranspose,
+        ]
+        if not use_mha2sha:
+            passes.remove(ConvertMhaToSha)
+        return passes

backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py

@@ -6,10 +6,7 @@
 import torch
 from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops
 from executorch.backends.qualcomm.builders.utils import is_parameter
-from executorch.backends.qualcomm.utils.constants import (
-    QCOM_BYPASS_NODE,
-    QCOM_FALLBACK_NODE,
-)
+from executorch.backends.qualcomm.utils.constants import QCOM_BYPASS_NODE
 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
@@ -42,26 +39,20 @@ def _fold_dq(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
             if n.target not in dq_ops:
                 continue
 
-            if not self.force_fold and (
-                QCOM_BYPASS_NODE in n.meta or QCOM_FALLBACK_NODE in n.meta
-            ):
-                continue
-
-            for user_n in user_list:
-                user_n.replace_input_with(n, n.args[0])
-            graph_module.graph.erase_node(n)
+            # skip parameters & buffers
+            if not self.force_fold and is_parameter(n.args[0], self.edge_program):
+                self._annotate_bypass(n)
+            else:
+                for user_n in user_list:
+                    user_n.replace_input_with(n, n.args[0])
+                graph_module.graph.erase_node(n)
 
     def _fold_q(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
         # remove q
         for n in graph_module.graph.nodes:
             if n.target not in q_ops:
                 continue
 
-            if not self.force_fold and (
-                QCOM_BYPASS_NODE in n.meta or QCOM_FALLBACK_NODE in n.meta
-            ):
-                continue
-
             to_be_removed = [n]
             source_n = n.args[0]
 
@@ -85,16 +76,7 @@ def _fold_q(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
             for n in to_be_removed:
                 graph_module.graph.erase_node(n)
 
-    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
-
     def call(self, graph_module: torch.fx.GraphModule):
-        if not self.force_fold:
-            self._preserve_qdq(graph_module)
         self._fold_dq(graph_module)
         self._fold_q(graph_module)
         dead_code_elimination_pass(graph_module)