Allow context-binary lowering to use edge dialect ops (#20518)

[JakeStevens]

Summary:

Unblocks the QNN context-binary path from lowering through to_edge with _use_edge_ops=True (the default). Previously it was pinned to EdgeCompileConfig(_use_edge_ops=False) purely to keep the qaisw context-loader custom op's original name, because loader detection was name-based.

We can simply unwrap the op overload and maintain the same check.

Differential Revision: D109598309

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 db9263c with merge base 7013c8d ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

@JakeStevens has exported this pull request. If you are a Meta employee, you can view the originating Diff in D109598309.

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

[JakeStevens]

I have updated the PR to the minimal changes:

1. remove the _use_edge_ops=False
2. add the unwrap_op_overload call to have the existing checks work

[psiddh]

@claude review this pr

[claude]

Claude finished @psiddh's task in 3m 13s —— View job

---

Code Review: PR #20518 — Allow context-binary lowering to use edge dialect ops

• Read PR diff
• Analyze qnn_preprocess.py unwrap change
• Verify partitioner / is_node_supported path still works
• Review new tests
• Post review

---

Overall this is a clean, well-scoped change. Dropping the _use_edge_ops=False pin lets from_context_binary lower through the default to_edge path, and the name-based loader detection is preserved by unwrapping the overload first. The logic is correct and the new tests cover the important paths. A few notes below.

Correctness — looks good ✅

The key concern with _use_edge_ops=True is that the qaisw custom op becomes an EdgeOpOverload rather than a raw OpOverload. I traced both detection sites and both remain correct:

• qnn_preprocess.py:94-99 — unwrap_op_overload(node.target) returns the underlying _op, whose __name__ is ctx_loader.default (not the edge-prefixed qaisw.ctx_loader.default), so eval("torch.ops.qaisw.ctx_loader.default") resolves to the same OpOverload object and the op == context_loader_target assertion holds.
• partition/qnn_partitioner.py:98 — the untouched OpContextLoader.namespace == node.target.namespace check still works because EdgeOpOverload.__getattr__ delegates .namespace to the wrapped _op, which reports "qaisw". Good that this didn't need a change.

Minor observations

1. Pre-existing bare except: (qnn_preprocess.py:102) — not introduced here, but worth flagging since the diff touches this block. It swallows everything, including a KeyError from node.meta[OpContextLoader.meta_ctx_bin] (line 101) if the binary was never stamped, and re-raises the generic "not supported in Qnn Delegate" message — which would be misleading for that failure mode. Also, unwrap_op_overload now raises TypeError for a target that isn't an op-overload wrapper, which similarly collapses into err_msg. Behavior is acceptable for this PR; just noting the error surface got slightly broader.

2. Test boilerplate duplication — all three new tests repeat the same ~15-line setup (define Library, torch.library.impl, the Model, torch.export.export, then filter for context-loader nodes). Per the repo's "minimal helpers unless they significantly improve readability" guidance, this is a case where a small shared helper like _build_context_loader_edge_program(op_name, check_ir_validity) returning (edge_program, loader_nodes) would genuinely cut ~45 lines and make each test's intent (which config/path it exercises) the only thing that varies. Fix this →

3. Redundant assertion (test_context_loader_edge_op_is_delegated) — after context_loader_nodes[0].meta[...] = ctx_bin, the immediately following assertEqual(ctx_bin, context_loader_nodes[0].meta[...]) just re-reads the value that was set on the previous line; it doesn't exercise any product code. Can be dropped.

Test coverage — good

• test_context_loader_op_lowers_with_ir_validation is the direct regression test for this change: it uses the default to_edge (edge ops + IR validity on) and confirms the loader node survives. 👍
• test_build_op_wrappers_returns_context_binary correctly exercises the unwrap path in _build_op_wrappers.
• The MagicMock-driven is_node_supported call is a reasonable way to avoid needing a live QNN manager, and the inline comment explaining it is helpful.

Note re: shewu-quic's earlier question

The op_name is None concern was on a prior revision; the current "unwrap + reuse existing check" approach no longer has that construct, so that thread is resolved.

Nothing here is blocking — items 1–3 are optional polish.
· branch export-D109598309

[harshs-qti]

Summary:

Unblocks the QNN context-binary path from lowering through to_edge with _use_edge_ops=True (the default). Previously it was pinned to EdgeCompileConfig(_use_edge_ops=False) purely to keep the qaisw context-loader custom op's original name, because loader detection was name-based.

We can simply unwrap the op overload and maintain the same check.

Differential Revision: D109598309

Can you add more context on why this is needed or what usecase is enabled with this change

[JakeStevens]

Summary:
Unblocks the QNN context-binary path from lowering through to_edge with _use_edge_ops=True (the default). Previously it was pinned to EdgeCompileConfig(_use_edge_ops=False) purely to keep the qaisw context-loader custom op's original name, because loader detection was name-based.
We can simply unwrap the op overload and maintain the same check.
Differential Revision: D109598309

Can you add more context on why this is needed or what usecase is enabled with this change

_use_edge_ops is being deprecated, this instance is the last remaining (in-tree) usage of _use_edge_ops=False. Once this is removed, we can finally deprecate. See the TODO here:

https://github.com/pytorch/executorch/blob/main/exir/capture/_config.py

[JakeStevens]

Addressed claude review

[JakeStevens]

@harshs-qti @shewu-quic are there any specific concerns you have or additional tests I can run?