Qualcomm AI Engine Direct - Pass migration - part 1

[zhaoxul-qti]

Summary

This is part 1 of a series whose goal is to deprecate to_edge_transform_and_lower_to_qnn and migrate Qualcomm-specific passes from the export pipeline into to_edge_transform_and_lower.

The core change is in exir/program/_program.py: after the EdgeProgramManager finishes running transform_passes, we now invoke lift_constant_tensor_pass on each method's ExportedProgram. Several QNN passes (CanonicalizeConv, ConvertLinearToConv2d, etc.) rewrite constant values and insert new get_attr nodes to store them. Without lifting, these mutated constants are not reflected back into the ExportedProgram; lift_constant_tensor_pass promotes the get_attr nodes into placeholders so the modified constants are correctly propagated.

Changes

• exir/program/_program.py: run lift_constant_tensor_pass over every method after edge_manager.transform(transform_passes) completes.
• Move CanonicalizeConv and ConvertLinearToConv2d out of the export pipeline (get_export_passes) and into the to-edge transform pipeline.
• Add new pass AnnotateGetAttr to repopulate QCOM_QUANT_ATTRS for get_attr nodes. Passes such as I64toI32 and LayoutTransform reconstruct the GraphModule (graph_module = super().call(graph_module).graph_module), which drops the quant attributes previously stored on get_attr node metadata.

Test plan

• CanonicalizeConv
  Test cases:
  TEST_MATRIX = {TestQNNFloatingPointOperator.test_qnn_backend_conv1d TestQNNFloatingPointOperator.test_qnn_conv1d_batch_norm TestQNNFloatingPointOperator.test_qnn_backend_conv2d TestQNNFloatingPointOperator.test_qnn_backend_conv3d_sequential TestQNNFloatingPointOperator.test_qnn_backend_conv_transpose1d TestQNNFloatingPointOperator.test_qnn_backend_conv_transpose2d TestQNNFloatingPointOperator.test_qnn_backend_conv_transpose3d TestQNNFloatingPointModel.test_qnn_backend_conv1d_relu_log_softmax TestQNNQuantizedOperator.test_qnn_backend_conv1d TestQNNQuantizedOperator.test_qnn_conv1d_batch_norm TestQNNQuantizedOperator.test_qnn_backend_conv2d TestQNNQuantizedOperator.test_qnn_backend_conv3d_sequential TestQNNQuantizedOperator.test_qnn_backend_conv_transpose1d TestQNNQuantizedOperator.test_qnn_backend_conv_transpose2d TestQNNQuantizedOperator.test_qnn_backend_conv_transpose3d TestQNNQuantizedModel.test_qnn_backend_conv1d_relu_log_softmax}
  Command:
  python backends/qualcomm/tests/test_qnn_delegate.py ${TEST_MATRIX} --build_folder build-android/ --host ${HOST_NAME} --device ${DEVICE_ID} --soc_model ${SOC_ID} --seed 1126 --backend htp
python -m pytest "backends/qualcomm/tests/rework/passes/test.py::test_canonicalize_conv[htp]" -v
python -m pytest "backends/qualcomm/tests/rework/passes/test.py::test_canonicalize_conv[lpai]" -v

• ConvertLinearToConv2d
  Test cases:
  TEST_MATRIX = {TestQNNFloatingPointOperator.test_qnn_backend_linear TestQNNFloatingPointOperator.test_qnn_backend_linear_to_conv2d TestQNNFloatingPointOperator.test_qnn_backend_linear_shared_weights TestQNNFloatingPointOperator.test_qnn_backend_linear_to_conv2d_shared_weights TestQNNQuantizedOperator.test_qnn_backend_linear TestQNNQuantizedOperator.test_qnn_backend_linear_to_conv2d TestQNNQuantizedOperator.test_qnn_backend_linear_shared_weights TestQNNQuantizedOperator.test_qnn_backend_linear_to_conv2d_shared_weights TestQNNQuantizedOperator.test_qnn_backend_linear_block TestQNNQuantizedOperator.test_qnn_backend_linear_to_conv2d_block TestQNNQuantizedOperator.test_qnn_backend_linear_qat}
  Command:
  python backends/qualcomm/tests/test_qnn_delegate.py ${TEST_MATRIX} --build_folder build-android/ --host ${HOST_NAME} --device ${DEVICE_ID} --soc_model ${SOC_ID} --seed 1126 --backend htp
python -m pytest "backends/qualcomm/tests/rework/passes/test.py::test_convert_linear_to_conv2d[htp]" -v
python -m pytest "backends/qualcomm/tests/rework/passes/test.py::test_convert_linear_to_conv2d[lpai]" -v
python -m pytest "backends/qualcomm/tests/rework/htp/op/v68/test.py::test_linear_shared_weights" -v

• AnnotateGetAttr
  Command:
  python -m pytest "backends/qualcomm/tests/rework/passes/test.py::test_annotate_get_attr" -v

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20738

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 6af438b with merge base a1eee71 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[zhaoxul-qti]

Hi @psiddh,
This PR is part one of the series to deprecate the to_edge_transform_and_lower_to_qnn.
Would you mind take a look on this PR?
Thanks for the review on this!

[psiddh]

@claude review

[claude]

Claude finished @psiddh's task in 2m 43s —— View job

---

Claude Code Review — Pass migration part 1

I focused on the four open review threads from @psiddh plus a couple of adjacent findings. Overall the migration direction is sound; the concerns raised are legitimate and I agree with most of them.

---

1. lift_constant_tensor_pass on the shared to_edge_transform_and_lower path (exir/program/_program.py:1345-1348)

if transform_passes is not None:
    edge_manager = edge_manager.transform(transform_passes)
    for method in edge_manager.methods:
        lift_constant_tensor_pass(edge_manager.exported_program(method))

@psiddh's concern is valid. This is the shared entry point for every backend, and lift_constant_tensor_pass is not a no-op signature-wise: it promotes get_attr nodes to placeholders and mutates graph_signature (appends InputSpec/buffers, see _program.py:263-288).

In practice this is usually a no-op for other backends because a post-to_edge graph rarely has get_attr nodes — constants are already lifted during export. But "usually" is the problem: any backend whose transform_passes introduce get_attr nodes (intentionally kept as attrs) would now silently have its input signature rewritten. That's a behavioral change hidden in shared code to serve a QNN-specific need.

I'd recommend scoping this to the QNN flow. Since the lifting only matters after QNN passes that mutate constants, the cleaner options are:

• Perform the lift inside a QNN pass (or a small terminal QNN pass) rather than in _program.py, or
• Gate it behind a condition so it only runs when a transform pass actually inserted get_attr nodes.

Note there's already precedent for QNN doing its own lifting: qnn_pass_manager.py:444,462 call lift_constant_tensor_pass in transform_for_export_pipeline / transform_for_to_edge_pipeline. Consolidating there keeps the shared path untouched.

2. AnnotateGetAttr._annotate_get_attr assumes users[0] is a dq (annotate_get_attr.py:44-47)

if node.op == "get_attr" and list(node.users)[0].target in dq_ops:
    dq_op = list(node.users)[0]

Two real issues here, both flagged by @psiddh:

• IndexError on a userless get_attr. A dangling get_attr (e.g. left behind before DCE) makes list(node.users)[0] throw and aborts export. Guard on node.users being non-empty first.
• Only the first user is inspected. If a get_attr feeds multiple users and the dq isn't index 0 (ordering is not guaranteed), the quant attrs are silently dropped for that constant.

Suggest scanning for any dq user instead:

if node.op != "get_attr":
    continue
dq_op = next((u for u in node.users if u.target in dq_ops), None)
if dq_op is None:
    continue
node.meta[QCOM_QUANT_ATTRS] = get_quant_attrs(self.edge_program, dq_op)

The same list(...users)[0] pattern appears in convert_linear_to_conv2d.py:90,94 and is worth the same defensive treatment.

Fix this →

3. Missing super().__init__() (annotate_get_attr.py:37)

def __init__(self, edge_program: torch.export.ExportedProgram):
    self.edge_program = edge_program

Correct catch. Every other pass in this PR (CanonicalizeConv:50, I64toI32:52) calls super().__init__(). It happens to work today only because call() bypasses the interpreter (super().call(...)) and returns PassResult directly — but that's fragile and inconsistent. Add super().__init__().

Fix this →

4. Re-annotate vs. preserve metadata (design)

@psiddh asks whether it's better to make I64toI32/LayoutTransform preserve the metadata rather than reconstruct it. My take:

The root cause is that graph_module = super().call(graph_module).graph_module (i64_to_i32.py:183) runs the ExportPass interpreter, which rebuilds the graph and does not carry get_attr node metadata across. That's a framework-level limitation, not something local to these two passes — so "fix it in the passes" would mean each pass individually copying meta back, and any future pass that reconstructs the GraphModule would silently reintroduce the bug.

A single re-annotation pass scheduled last (as done here, via the dependency table qnn_pass_manager.py:273-278) is a defensible centralization. The tradeoff to call out in the docstring/PR: it re-derives attrs from the surviving dq node rather than preserving the originals, so it's only correct as long as the get_attr → dq structure is intact at that point (which is why the users[0] robustness in #2 matters). If a constant's dq has already been folded away by then, its attrs won't be recoverable. Worth a comment noting that assumption.

---

Additional observations

• canonicalize_conv.py:146-152 — the weight resolution assumes weight is either a placeholder or a node with .args[0] (a dq). If a weight arrives as a bare get_attr (no dq), the else branch does node.args[ConvParamIdx.WEIGHT].args[0] on a get_attr, which has no args. Likely fine for the tested FP/QDQ shapes, but a get_attr weight would break it.
• node_visitor.py:93,100 — adding torchao.quantize_affine/dequantize_affine to q_ops/dq_ops is a broad change that affects every consumer of those sets (FoldQDQ, AnnotateQuantAttrs, etc.), not just this PR's passes. Since get_quant_attrs reads args positionally by schema, please confirm the affine op's schema args are handled (block quant path in convert_linear_to_conv2d.py:108-112 suggests they differ from per-tensor/per-channel).

Nice cleanups in the refactor — collapsing the conv variants onto aten.convolution.default with the ConvParamIdx enum reads much better than the old positional indexing.
· branch dev1/zhaoxul/pass_migration_pt1

[psiddh]

@zhaoxul-qti can you take a look at the comments above ?

[shewu-quic]

@zhaoxul-qti can you take a look at the comments above ?

Sorry, @zhaoxul-qti is currently OoO and should be back this week. He'll take a look at the comments and respond as soon as possible.

Thanks for the review!

[zhaoxul-qti]

Hi @psiddh,
I have addressed all the comments above and pushed a new commit on it.
I think this PR is ready to merge once there is no CI failure.
Thanks for reviewing it.