Qualcomm AI Engine Direct - Pass migration - part 1

- Add `lift_constant_tensor_pass` after `EdgeProgramManager` completes the given passes.
- Refactor the execution phase of Qualcomm‑specific passes and move to `to_edge_transform_and_lower`.
- Introduce the `AnnotateGetAttr` pass to ensure quantization attributes are preserved for the `get_attr` node.

Author: zhaoxul-qti

backends/qualcomm/_passes/__init__.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from .annotate_avg_pool1d import AnnotateAvgPool1D
+from .annotate_get_attr import AnnotateGetAttr
 from .annotate_quant_attrs import AnnotateQuantAttrs
 from .annotate_stack import AnnotateStack
 from .annotate_unbind import AnnotateUnbind
@@ -70,6 +71,7 @@
 
 __all__ = [
     AnnotateAvgPool1D,
+    AnnotateGetAttr,
     AnnotateQuantAttrs,
     AnnotateStack,
     AnnotateUnbind,

backends/qualcomm/_passes/annotate_get_attr.py

@@ -0,0 +1,52 @@
+# 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.backends.qualcomm.builders.node_visitor import dq_ops
+from executorch.backends.qualcomm.utils.constants import QCOM_QUANT_ATTRS
+from executorch.exir.pass_base import ExportPass, PassResult
+
+from .utils import get_quant_attrs
+
+
+class AnnotateGetAttr(ExportPass):
+    """
+    Annotates quantization attributes for `get_attr` nodes.
+
+    In passes such as `CanonicalizeConv`, `ConvertLinearToConv2d`,
+    `LiftConstantScalarOperands`, and others, `get_attr` nodes and the
+    corresponding quantization attributes are inserted into the GraphModule
+    to store modified constant values.
+
+    However, the quantization attributes associated with `get_attr` nodes will
+    be discarded in certain passes, such as `I64toI32` and `LayoutTransform`.
+    This happens due to the following line:
+    `graph_module = super().call(graph_module).graph_module`
+
+    which reconstructs the GraphModule and drops any existing quantization
+    attributes stored in the metadata of `get_attr` nodes.
+
+    To guarantee correctness, this pass repopulates the quantization attributes
+    for `get_attr` nodes and ensures to be scheduled after the `I64toI32` and
+    `LayoutTransform` passes.
+    """
+
+    def __init__(self, edge_program: torch.export.ExportedProgram):
+        self.edge_program = edge_program
+
+    def _annotate_get_attr(
+        self, graph_module: torch.fx.GraphModule
+    ) -> torch.fx.GraphModule:
+        for node in graph_module.graph.nodes:
+            if node.op == "get_attr" and list(node.users)[0].target in dq_ops:
+                dq_op = list(node.users)[0]
+                quant_attrs = get_quant_attrs(self.edge_program, dq_op)
+                node.meta[QCOM_QUANT_ATTRS] = quant_attrs
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        self._annotate_get_attr(graph_module)
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/qualcomm/_passes/annotate_quant_attrs.py

@@ -55,12 +55,6 @@ def _annotate_source_nodes(
             source_n = quant_node.args[0]
         source_n.meta[QCOM_QUANT_ATTRS] = quant_attrs
 
-    def _expand(self, tensor, dim, axis) -> torch.Tensor:
-        tensor = tensor[(...,) + (None,) * (dim - 1)]
-        order = torch.arange(dim).tolist()
-        order[axis], order[0] = order[0], order[axis]
-        return tensor.permute(order)
-
     # Find the the last dq nodes between regular op nodes
     # Return dq2 in example below when q1 is given as node parameter:
     # ... -> n1 -> q1 -> dq1 -> q2 -> dq2 -> n2 -> ...

backends/qualcomm/_passes/backends/gpu/qnn_gpu_pass_manager.py

@@ -38,12 +38,9 @@ def get_annotation_passes(cls):
         return []
 
     @classmethod
-    def get_export_passes(
-        cls,
-        convert_linear_to_conv2d: bool = False,
-    ):
+    def get_export_passes(cls):
         # DecomposeReciprocal should be placed in the export pipeline, as it depends on
         # LiftConstantScalarOperands to lift the scalar operand.
         passes = [DecomposeReciprocal]
-        passes.extend(super().get_export_passes(convert_linear_to_conv2d))
+        passes.extend(super().get_export_passes())
         return passes

backends/qualcomm/_passes/backends/htp/qnn_htp_pass_manager.py

@@ -38,12 +38,9 @@ def get_annotation_passes(cls):
         return passes
 
     @classmethod
-    def get_export_passes(
-        cls,
-        convert_linear_to_conv2d: bool = False,
-    ):
+    def get_export_passes(cls):
         # DecomposeReciprocal should be placed in the export pipeline, as it depends on
         # LiftConstantScalarOperands to lift the scalar operand.
         passes = [DecomposeReciprocal]
-        passes.extend(super().get_export_passes(convert_linear_to_conv2d))
+        passes.extend(super().get_export_passes())
         return passes

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

@@ -70,14 +70,11 @@ def get_annotation_passes(cls):
         return passes
 
     @classmethod
-    def get_export_passes(
-        cls,
-        convert_linear_to_conv2d: bool = False,
-    ):
+    def get_export_passes(cls):
         # Both DecomposeHardSigmoid and DecomposeReciprocal should be placed in the export
         # pipeline, as they rely on LiftConstantScalarOperands to lift the scalar operand.
         passes = [DecomposeHardsigmoid, DecomposeReciprocal]
-        passes.extend(super().get_export_passes(convert_linear_to_conv2d))
+        passes.extend(super().get_export_passes())
         return passes
 
     @classmethod