Qualcomm AI Engine Direct - Fold views of parameters before annotation

backends/qualcomm/_passes/__init__.py

@@ -52,6 +52,7 @@
 from .decompose_wrap_with_autocast import DecomposeWrapWithAutocast
 from .expand_broadcast_tensor_shape import ExpandBroadcastTensorShape
 from .fixed_linear_keep_dim import FixedLinearKeepDim
+from .fold_constant_views import FoldConstantViews
 from .fold_qdq import FoldQDQ
 from .fuse_batch_norm_with_conv import FuseBatchNormWithConv
 from .fuse_consecutive_cast import FuseConsecutiveCast
@@ -125,6 +126,7 @@
     DecomposeWrapWithAutocast,
     ExpandBroadcastTensorShape,
     FixedLinearKeepDim,
+    FoldConstantViews,
     FoldQDQ,
     FuseBatchNormWithConv,
     FuseConsecutiveCast,

backends/qualcomm/_passes/fold_constant_views.py

@@ -0,0 +1,71 @@
+# 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 operator
+
+import torch
+from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.passes import dead_code_elimination_pass
+
+
+class FoldConstantViews(ExportPass):
+    """
+    Replace views and splits of a parameter with constants before annotation.
+
+    nn.MultiheadAttention splits in_proj_weight at runtime when query and value
+    differ, so each projection's weight is the output of a split rather than a
+    parameter. A per-channel weight encoding on such a tensor is never folded
+    into a static weight, and lowering fails. Folding gives each projection its
+    own parameter.
+    """
+
+    view_ops = {
+        torch.ops.aten.split.Tensor,
+        torch.ops.aten.split_with_sizes.default,
+        torch.ops.aten.chunk.default,
+        torch.ops.aten.unbind.int,
+        torch.ops.aten.slice.Tensor,
+        torch.ops.aten.select.int,
+    }
+
+    def __init__(self):
+        super(FoldConstantViews, self).__init__()
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        graph = graph_module.graph
+        for node in list(graph.nodes):
+            if (
+                node.target not in self.view_ops
+                or not isinstance(node.args[0], torch.fx.Node)
+                or node.args[0].op != "get_attr"
+                or any(isinstance(a, torch.fx.Node) for a in node.args[1:])
+            ):
+                continue
+            source = node.args[0]
+            *path, attr = source.target.split(".")
+            tensor = getattr(graph_module.get_submodule(".".join(path)), attr)
+            outputs = node.target(tensor, *node.args[1:], **node.kwargs)
+            if isinstance(outputs, torch.Tensor):
+                replaced = [(node, outputs)]
+            else:
+                replaced = [
+                    (user, outputs[user.args[1]])
+                    for user in node.users
+                    if user.target == operator.getitem
+                ]
+            for old, value in replaced:
+                name = f"{source.target.replace('.', '_')}_{old.name}"
+                constant = value.detach().clone().contiguous()
+                if isinstance(tensor, torch.nn.Parameter):
+                    constant = torch.nn.Parameter(constant, requires_grad=False)
+                setattr(graph_module, name, constant)
+                with graph.inserting_before(old):
+                    get_attr = graph.get_attr(name)
+                get_attr.meta = old.meta.copy()
+                old.replace_all_uses_with(get_attr)
+
+        dead_code_elimination_pass(graph_module)
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -55,6 +55,7 @@
     DecomposeWrapWithAutocast,
     ExpandBroadcastTensorShape,
     FixedLinearKeepDim,
+    FoldConstantViews,
     FoldQDQ,
     FuseBatchNormWithConv,
     FuseConsecutiveCast,
@@ -173,6 +174,7 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False):
         """Return annotation pipeline pass classes. Override in subclasses to add backend-specific passes."""
         passes = [
             RemoveRedundancy,
+            FoldConstantViews,
             # Retraced autocast regions may emit ops that the passes below decompose.
             DecomposeWrapWithAutocast,
             RecomposePixelUnshuffle,

backends/qualcomm/tests/test_passes.py

@@ -14,6 +14,7 @@
     DecomposeColIm,
     DecomposeWrapWithAutocast,
     ExpandBroadcastTensorShape,
+    FoldConstantViews,
     FoldQDQ,
     InsertIOQDQ,
     InsertReshapeForReduceOps,
@@ -920,6 +921,32 @@ def test_make_qnn_per_block_config_rejects_asymmetric(self):
         with self.assertRaisesRegex(ValueError, "symmetric"):
             visitor.make_qnn_per_block_config(parameter_node, quant_attrs)
 
+    def test_fold_constant_views_mha_in_proj(self):
+        # With query != value, nn.MultiheadAttention splits in_proj_weight at
+        # runtime; each projection must get its own constant weight.
+        class CrossAttention(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.attn = torch.nn.MultiheadAttention(16, 2, batch_first=True)
+
+            def forward(self, x):
+                q = x.sigmoid()
+                return self.attn(q, q, x, need_weights=False)[0]
+
+        mod = CrossAttention().eval()
+        x = torch.randn(1, 5, 16)
+        gm = torch.export.export(mod, (x,)).module()
+        FoldConstantViews()(gm)
+
+        weights = [
+            n.args[1]
+            for n in gm.graph.nodes
+            if n.target == torch.ops.aten.linear.default
+        ]
+        self.assertTrue(weights)
+        self.assertTrue(all(w.op == "get_attr" for w in weights))
+        self.assertTrue(torch.allclose(gm(x), mod(x)))
+
     def test_insert_reshape_for_argmax(self):
         class ArgmaxModule(torch.nn.Module):
             def forward(self, x):