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

• 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

[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.

❌ 1 New Failure

As of commit 735df03 with merge base e500fee ():

NEW FAILURE - The following job has failed:

• Test ARM Backend / test-arm / test-backend-linux (arm_tosa_fp, operators) / linux-job (gh)
  RuntimeError: Command docker exec -t b3735cf95ebd41fbc35281e6951dbb0be347cbc98120d5b59b54048341ce21a9 /exec failed with exit code 2

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]

This is the shared export code that every backend goes through, not just specific to QNN. wouldn't it silently change the other backend's input signatures? Instead of here can we restrict it into QNN specific flow ?

[zhaoxul-qti]

This is the most critical part of the entire migration series.

One of the main challenges is that transform_passes currently operates on a GraphModule rather than an ExportedProgram. As a result, any pass that needs to modify constants at this stage (for example, converting Conv1d to Conv2d, which requires updating the weight tensor and propagating the change back to the ExportedProgram's input signature) cannot be handled correctly without lifting the constant tensors. This is the reason lift_constant_tensor_pass is required here.

In the original workflow, this logic was confined to the QNN-specific path within to_edge_transform_and_lower_to_qnn, which is exactly the flow that this PR is attempting to deprecate and migrate away from.

Given that this is the shared export code that every backend goes through, I was wondering if these two lines of change be considered acceptable once the full CI suite passes across all backends.

[zhaoxul-qti]

Hi @psiddh,
I'd appreciate your perspective on this part when you have a chance.

[psiddh]

Is it possible that get_attr has no users ? If so it'll IndexError and abort the export ? Also since it only checks the first user (is this guaranteed ), a constant whose dq isn't first would silently miss other attrs? Could we guard on node.users and scan for any dq user instead of assuming [0] ?

[zhaoxul-qti]

Thanks for pointing this out.

I don't expect the get_attr node to have no users in this case, but I agree that adding a defensive check would make it more robust. I'll also add the check and validate for any dq user in node.users.

[psiddh]

should it call super().init()

[zhaoxul-qti]

I don't think this is required, as the current implementation works correctly without it.

[psiddh]

Also this pass exists because other passes drop quant metadata on get_attr nodes ? If so, would it be better to have other passes preserve the it instead ? What do you think ?

[zhaoxul-qti]

Yes, this pass exists because other passes can inadvertently drop quant metadata on get_attr nodes.

Although we could patch each affected pass individually, that approach does not scale well. Any future pass that recreates or transforms the graph through

graph_module = super().call(graph_module).graph_module

would need to remember to preserve the quant attributes for get_attr nodes as well. Missing that step could easily introduce regressions.

For that reason, I added a generic AnnotateGetAttr pass to address this problem once for all. I believe this provides a more uniform and maintainable solution, while also reducing the risk of future passes accidentally breaking the get_attr's quantization metadata propagation.

[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!