Qualcomm AI Engine Direct - Support Constant Propogation

[winskuo-quic]

Summary

• Support Constant Folding through edge pass.

Test plan

pytest backends/qualcomm/tests/rework/passes/test.py -k "test_constant_folding"

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 e460e1e with merge base dba8a83 ():

NEW FAILURE - The following job has failed:

• pull / test-qnn-delegate-linux / linux-job (gh)
  backends/qualcomm/tests/test_qnn_delegate.py::TestQNNQuantizedOperator::test_qnn_backend_slice_copy

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.

[psiddh]

@claude review this PR

[claude]

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

---

Review: Qualcomm AI Engine Direct - Support Constant Propagation

• Gather context / read diff
• Review constant_folding.py
• Review pass manager, utils, and tests
• Post review

The overall approach makes sense: seed from constant placeholders, propagate in topological order, and only materialize at the boundary. I found two correctness issues and a few design points to look at before merging. I only read the code; I couldn't run the tests here.

---

🔴 Correctness

1. Impure ops (RNG) get folded into a frozen value — constant_folding.py:152

_propagate only checks that node.op == "call_function" and that all args are const. Ops like aten.rand, randn, randint, bernoulli and empty have only primitive args, so _is_const returns True. The pass then runs them once at compile time and bakes the result in as a buffer, so the model stops producing fresh random values at runtime.

The upstream pass this code was copied from (exir/passes/constant_prop_pass.py:143-153) guards against this explicitly:

if (
    node.op != "call_function"
    or node.target is memory.alloc
    or node.target in all_skip_targets
    # Ops with side effects (RNG draws, mutation) have to run at runtime.
    or node.is_impure()
):
    continue

Please add at least node.is_impure(), and a test case like x + torch.rand(4, 4).

Fix this →

2. Non-tensor results crash when materialized — constant_folding.py:177, :191

node_to_tensor can hold values that aren't tensors: Python scalars from _local_scalar_dense or item, a SymInt from sym_size, or tuples from multi-output ops. If one of those has a user that can't be folded, _materialize_as_buffer calls graph_module.register_buffer(name, <int/float>), which raises TypeError. Even if it didn't raise, it would swap a scalar arg for a tensor get_attr. Skip materialization when not isinstance(tensor, torch.Tensor).

Fix this →

---

🟡 Design / robustness

3. run_decompositions({}) in __init__ is expensive, and its snapshot is taken early — constant_folding.py:61

• get_to_edge_transform_passes builds every pass before any of them runs (qnn_pass_manager.py, self.add_pass(p(**kwargs)) in the loop). So this snapshot of the signature and constants is taken before FoldQDQ, I64toI32, LayoutTransform and the rest have run. It still works today only because the lookups match on placeholder names and get_parameter re-casts to meta["val"].dtype. That's fragile.
• It re-traces the whole program on every QNN lowering, even when nothing can be folded. For LLM-sized graphs that adds real time and memory.
• The comment says decomposition is needed so the signature records mutable buffers. But an edge ExportedProgram should already carry buffers_to_mutate. Can you give the case where edge_program.graph_signature was missing it? If there is one, please document it. Otherwise use edge_program directly, as AnnotateGetAttr, FoldQDQ and LayoutTransform do. At minimum, move this work lazily into call().

4. No limit on constant size, and full isn't skipped

Upstream skips aten.full by default (_DEFAULT_SKIP_TARGETS_NO_QUANT). Here, full, arange, zeros, and expand/repeat of small constants are all folded. For example, a causal mask built with full([1, 1, S, S]) followed by triu becomes an S×S constant in the .pte. With S=4096 in fp32 that's 64 MB, where the ops would otherwise compute it on-device. Consider a skip set (starting with upstream's full) and/or a size limit, such as not folding when the output numel is much larger than the input numel.

5. Quantized constant chains through a non-QDQ op aren't tested

Rule 3 only protects the case where the boundary node is a dq. Take weight -> dq -> permute/reshape -> linear after FoldQDQ: the boundary is the permute. It gets materialized as a float get_attr, the weight's dq node is removed by DCE, and the new node gets copy_meta(permute.meta). The result is only correct if permute already has QCOM_QUANT_ATTRS. If the quantizer didn't annotate that op, the linear gets a float weight in a quantized graph. AnnotateGetAttr can't recover the attrs because there's no dq user left. Please add a quantized subtest along these lines, e.g. a weight transposed in forward.

6. Pass ordering relative to AnnotateGetAttr — qnn_pass_manager.py

Both ConstantFolding and AnnotateGetAttr depend only on LayoutTransform, so which one runs first depends on how the solver breaks the tie. ConstantFolding creates new get_attr nodes. If any of them can end up with a dq user, AnnotateGetAttr has to run after it. Consider adding ConstantFolding to AnnotateGetAttr's dependency list so the order is fixed.

7. Seeds only come from placeholders — _get_const_placeholders

get_attr nodes created by earlier passes (CanonicalizeConv, ConvertLinearToConv2d, DecomposeAsStrided) are never used as seeds, even though is_parameter already returns True for get_attr. This isn't a bug, just a missed chance to fold more. Please mention it in the docstring if it's intentional.

---

🟢 Nits

• constant_folding.py:29,44,159 — the logic is copied from exir/passes/constant_prop_pass.py, not const_prop_pass.py. The docstring's claim also mixes the two files up: constant_prop_pass operates on an edge ExportedProgram. The real limitation is that QNN edge passes run on a GraphModule without updating the signature. Please reword it.
• You could import _PRIMITIVE_TYPES, is_const and get_data from exir/passes/constant_prop_pass.py. Upstream get_data needs exported_program only for its signature, so a tiny wrapper would do. That's better than keeping a third copy.
• _EDGE_QDQ_TARGETS repeats what q_ops | dq_ops in builders/node_visitor.py already gives you, which is also what the tests use.
• constant_folding.py:177 — upstream only calls .contiguous() when a stride is 0. Here every result is cloned, which doubles peak memory for large folded constants. Only the final materialized boundary tensors need to be contiguous.
• The docstrings have stray trailing whitespace and "Boundry" / "afer" typos. _propagate's docstring has trailing spaces on several lines, which lintrunner will flag.
• PR title and file header: "Propogation" → "Propagation".
• pattern.py:768 — the comment "neither it nor the following add is_fp" looks like a find/replace mistake ("is folded").
• utils.py: the tensor → node rename is fine, but it's unrelated to this change.
• The PR needs a release notes: label (see the bot comment).

Summary

The two items I'd fix before merging are #1 (impure/RNG ops get folded) and #2 (a non-tensor result crashes materialization). #3 (the run_decompositions cost and early snapshot) and #4 (constant-size blowup) are worth fixing before this is on by default for LLM workloads. The tests are good for the float cases. What's missing is an impure-op case and a quantized chain that goes through a non-dq op.