Repository navigation
Qualcomm AI Engine Direct - Decompose kept-whole ops that HTP leaves on CPU - #24429
shoumikhin wants to merge 6 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24429
Note: Links to docs will display an error until the docs builds have been completed. ❌ 12 Pending, 2 Unclassified FailuresAs of commit 78cdf65 with merge base 893d815 ( UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 34m 51s —— View job This PR makes the HTP Detailed reviewPR Review: #24429SummaryOn HTP,
All four changed files were reviewed. Backward Compatibility
API Design
Testing
Code Quality
RecommendationRequest Changes What blocks approval:
|
…on CPU On HTP, a partition that reads nothing at runtime stays on the CPU, because HTP refuses a graph without runtime inputs. That is decided after decomposition. So a constant-only op that QNN keeps whole, such as linear, could land in such a partition undecomposed, and lowering failed with "was not decomposed or delegated". The HTP filter of ops_to_not_decompose now runs the same drop rule before decomposition and decomposes only the kept-whole ops it would leave on the CPU. The real partitions are not known yet, so it partitions the graph as if QNN took every op except the ones whose data types or shapes rule it out (the checks filter_fn already makes, now shared as must_stay_on_cpu). Ops in partitions that stay delegated keep their whole form, and with it their quantization encodings. The caller's skip_node_op_set and skip_node_id_set are left out of that estimate on purpose. Decomposing adds nodes, and the new nodes take names that the later nodes had before. If a skip could change what is decomposed, adding one skip could rename the node that a skip_node_id_set entry means, and the wrong node would be skipped. Ops that QNN only rejects while partitioning are not known yet either. So a kept-whole op that only a skip or such an op cuts off still fails as before. exir calls ops_to_not_decompose several times before decomposition and uses only one of the filters, so the estimate runs on the filter's first use. On a 10,240-node graph the partition search alone takes about 10 s. On the lowered program the filter stays plain filter_fn, so exir still reports kept-whole ops left on the CPU. The constant-only partition test now stubs the QNN manager and the compile, so it runs without the QNN SDK. It checks that a constant-only linear cut off by bfloat16 casts is decomposed, that a constant-only linear in a delegated partition stays whole, that the filter is plain filter_fn on the lowered program, and that a node ID skip still names the same add when an op skip and an ID-skipped feeder are present. Same change as #24429 on main, including its Partitioner.ops_to_not_decompose docstring. It replaces the earlier version carried here, which decomposed every op that reads only constants.
8e4fe06 to
5dab5e9
Compare
|
Thanks for the review. The latest push changes this:
|
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 26m 55s —— View job I recommend approval: I found no correctness problems in this revision, and two minor requests remain: run the new scenario against the real QNN SDK in CI, and keep the HTP-only drop rule in one place. Detailed reviewPR Review: #24429SummaryOn HTP, I found no correctness problems beyond the limits the description already lists. Two minor problems remain:
All four changed files were reviewed. Code Quality
Testing
RecommendationApprove Nothing I found blocks approval. The real-SDK test is the more important of the two requests, because it is the only way CI can catch the estimate disagreeing with the real HTP checks. The |
… Samsung passes (#24316) ## What is wrong today torch.export does not record some of the node names it assigns: the names of inputs, and the names of ops it renamed to make room for an input. With an input called `mul`, the model's own `input * 2` becomes `mul_1`, and nothing marks `mul_1` as taken. A later pass that adds a node can pick one of those names again, and the graph breaks: ``` RuntimeError: Node redefined name t_1! ``` QNN models with autocast regions or attention hit this in the QNN export passes. The Samsung passes hit it too: with an input named `cat`, `torch.cat` followed by `torch.roll` fails in `DecomposeRoll`. ## What this change does `reserve_node_names` in `backends/transforms/utils.py` records the name of every node in the top-level graph, so nodes added later get names that are not taken. The QNN and Samsung pass managers call it at the start of `_transform`. torch.fx has no public way to reserve a name, so the helper adds the names to the graph's name set directly. Copying the graph would register the names through public API, but it also renames inputs named after Python builtins (`input` becomes `input_1`) while the graph signature keeps the old names. The real fix belongs in torch.export (pytorch/pytorch#202927), and a TODO says to drop the helper once that lands. ## Tests `test_export_pipeline_keeps_node_names_unique` runs the QNN export passes on an autocast model and on an attention model whose inputs are named `input` and `mul`. It checks that the graph lints, keeps the input names the graph signature uses, and matches eager, and that the collision really happens without the helper. If a torch update removes the collision, the test skips with a pointer to the torch issue instead of failing. On Linux x86_64 with torch 2.14 and a 2.16 nightly: the test fails with `Node redefined name t_1!` and `Node redefined name mul_1!` when the call is removed, and it fails on the `input` case when the helper copies the graph instead. The Samsung case above fails without the helper and passes with it (torch only). `test_passes.py` passes, except two tests that need torchtune or the quantized CPU kernels. Earlier versions of this PR also carried an HTP change that decomposes constant-only ops left on the CPU, and an RMSNorm dtype fix. Both now have their own PRs: RMSNorm dtype fix: #24428 HTP constant-only decomposition: #24429 cc @cbilgin @psiddh
…on CPU On HTP, a partition that reads nothing at runtime stays on the CPU, because HTP refuses a graph without runtime inputs. That is decided after decomposition. So a constant-only op that QNN keeps whole, such as linear, could land in such a partition undecomposed, and lowering failed with "was not decomposed or delegated". The HTP filter of ops_to_not_decompose now runs the same drop rule before the kept-whole ops are decomposed, and decomposes only the ones it would leave on the CPU. The real partitions are not known yet, so it partitions the graph as if QNN took every op except the ones whose data types or shapes rule it out (the checks filter_fn already makes, now shared as must_stay_on_cpu). Ops in partitions that stay delegated keep their whole form, and with it their quantization encodings. The caller's skip_node_op_set and skip_node_id_set are left out of that estimate on purpose. Decomposing adds nodes, and the new nodes take names that the later nodes had before. If a skip could change what is decomposed, adding one skip could rename the node that a skip_node_id_set entry means, and the wrong node would be skipped. Ops that QNN only rejects while partitioning are not known yet either. So a kept-whole op that only a skip or such an op cuts off still fails as before. And when such an op reads only constants and feeds a kept-whole op, the estimate decomposes the kept-whole op although it ends up delegated, so it is delegated in pieces. exir calls ops_to_not_decompose several times on the ATen program and uses only one of the filters, so the estimate runs on the filter's first use. On a 10,240-node graph the partition search alone takes about 10 s. On the lowered program the filter stays plain filter_fn, so exir still reports kept-whole ops left on the CPU. A new test stubs the QNN SDK setup, the QNN manager and the compile, so it runs without the SDK. It checks that constant-only linears cut off by bfloat16 casts or by a masked select are decomposed, that a constant-only linear in a delegated partition stays whole, that the filter is plain filter_fn on the lowered program, and that a node ID skip still names the same add when an op skip and an ID-skipped feeder are present. The existing constant-only partition test still runs against the real SDK.
to_edge_transform_and_lower calls ops_to_not_decompose during decomposition, on the ATen program, to choose which nodes to keep whole, and after lowering, on the EDGE program, to report kept-whole nodes that were neither decomposed nor delegated. The docstring only described the first use. QnnPartitioner tells the two apart by ep.dialect, and its HTP filter only works on the nodes of the program it was given, so the docstring now names both stages, their dialects, and that rule. It also says that the list checked after lowering includes the lists of earlier partitioners for the same method. Docstring only.
To decide which kept-whole ops to decompose, the HTP filter of ops_to_not_decompose estimated which ops HTP leaves on the CPU by running the capability-based partitioner. That partitioner's time grows with the square of the graph size: it was about 4 times slower for each doubling, from 2,000 to 32,000 nodes. On a large vision-language-action model, with about 43,000 nodes, the export sat in its cycle search for more than 2 hours. Only a partition whose ops read no runtime value is dropped, so the estimate now looks only at those ops. Only ops whose data types or shapes rule QNN out count as CPU ops, and every getitem counts as a QNN op, since QNN's op check accepts it. One pass over the graph in order and one in reverse find, for every op QNN could take: - whether it reads a runtime value, directly or through other QNN ops; - whether it reaches a QNN op through a CPU op, without passing an op that reads a runtime value. The partitioner may split such an op off to avoid a cycle. A walk from the ops that read a runtime value then finds the ops linked to them by QNN ops alone. An op that reads no runtime value is left on the CPU when nothing links it to one, or when it has such a CPU detour. Each pass visits every node and edge once. Skips and ops that QNN rejects while partitioning are still not known here, as before. Apart from those, the estimate errs toward decomposing. The partitioner also merges partitions that share no data, so a constant-only op that nothing links to a runtime value, or one with a CPU detour, can still end up delegated. Then it is delegated in pieces. The constant-only linear that stays whole in the HTP test now also feeds an add that reaches another QNN op through CPU casts, so the test checks that such a path does not decompose it. New cases in the HTP test, where a constant-only linear must be decomposed: one feeds a delegated add both directly and through CPU casts, one feeds a delegated add and a CPU split whose getitems QNN takes, and one sits in a model whose input reaches only a CPU op. A new test times the estimate on a 6,000-node chain: well under 1 s now, about 10 s with the partitioner on the same machine.
reserve_node_names was only tested through the Qualcomm pass manager, and only in a test that needs the Qualcomm backend. The Samsung pass manager calls it too. A torch-only test now exports a model whose input takes the name "mul", so export renames the multiply to "mul_1". After reserve_node_names, a new node asked to be named "mul_1" gets another name, and the graph lints.
When the pass manager failed for a reason other than the name collision, the test reported an assertion about the message, and the real error showed up only as the context of that assertion. The test now re-raises any error that is not the collision, with its own traceback.
The test skips when a torch update removes the name collision, since the cases then no longer test reserve_node_names. It skipped as soon as any case stopped colliding, which also hid the cases that still did. It now skips only when none of them collide.
5dab5e9 to
78cdf65
Compare
|
Pushed 4 new commits, rebased on main:
|
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 56m 24s —— View job Requesting changes: the PR's new HTP estimate of which kept-whole ops end up on CPU counts export's runtime-assert nodes as QNN ops, so a masked select makes it decompose linears that main delegates whole, and its cost grows quadratically with the number of constants. Detailed reviewPR Review: #24429SummaryThe PR adds an HTP-only estimate to Code Quality
Testing
API Design
Backward Compatibility
Performance
RecommendationRequest Changes Two problems block approval:
Both are small fixes inside
|
What is wrong today
On HTP, a partition that reads nothing at runtime stays on the CPU, because HTP refuses a graph without runtime inputs. That decision happens after decomposition. So a constant-only op that QNN keeps whole, such as
linear, can land in such a partition undecomposed, and lowering fails:What this change does
Before the kept-whole ops are decomposed, the HTP filter of
ops_to_not_decomposeestimates which ops HTP will leave on the CPU, and decomposes only the kept-whole ops among them. An op counts as a CPU op only when its data types or shapes rule QNN out (the checksfilter_fnalready makes, now shared asmust_stay_on_cpu), and every getitem counts as a QNN op, since QNN's op check accepts it. Ops that the estimate expects to be delegated stay whole and keep their quantization encodings.Only a partition whose ops read no runtime value is dropped, so the estimate looks only at those ops. It leaves such an op on the CPU when no chain of QNN ops links it to an op that reads a runtime value, or when it reaches a QNN op through a CPU op without passing an op that reads a runtime value, because the partitioner may then split it off to avoid a cycle. The estimate is a few passes over the graph, each visiting every node and edge once. Running the partitioner here instead would be far too slow on large models, since its time grows with the square of the graph size (about 4 times slower for each doubling, measured from 2,000 to 32,000 nodes).
The caller's
skip_node_op_setandskip_node_id_setare left out of this estimate on purpose. Decomposing adds nodes that take the names of later nodes, so a skip that changed what is decomposed could rename the node a skip entry means.The
Partitioner.ops_to_not_decomposedocstring now describes both stages.What it covers
Covered: kept-whole ops cut off from the model's inputs by data types (such as bfloat16) or data-dependent shapes.
Not covered, and failing as before:
rms_normorpixel_unshuffle;scatter.src,scatter_addandscatter_reduce(no default decomposition);The estimate can also err the other way. The partitioner also merges partitions that share no data, so a constant-only kept-whole op that nothing links to a runtime value, or one with a CPU detour, can still end up delegated. Then it is delegated in pieces.
Tests
test_htp_decomposes_kept_whole_ops_left_on_cpustubs the QNN SDK and the compile. It checks that:filter_fn, so exir still reports kept-whole ops left on the CPU;test_htp_kept_whole_estimate_is_fast_on_large_graphstimes the estimate on a 6,000-node chain. It takes well under 1 s, where running the partitioner on the same machine takes about 10 s.Each check fails when the rule it covers is removed, and the tests pass with torch 2.14 and 2.16 nightly on Linux x86_64.
test_partitioner_keeps_constant_only_subgraph_on_cpustill runs against the real SDK. With the QNN SDK (SM8650), Swin-v2-T lowers as on main in 8a8w and fp16: one partition, nothing on the CPU. On a 43,000-node vision-language-action model the estimate takes under 4 s.Follow-ups from the review of #24316
test_export_pipeline_keeps_node_names_uniquenow skips only when no case collides any more, and it re-raises any error that is not the name collision, so the original traceback is kept.test_reserve_node_nameschecks that a node added afterreserve_node_namesdoes not take a name that export already used. It tests the helper directly, so it needs no backend.cc @cbilgin @psiddh