[ET-VK] Wire the missing resize functions for embedding and index_select - #22410
[ET-VK] Wire the missing resize functions for embedding and index_select#22410msluszniak wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22410
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
SS-JIA
left a comment
SS-JIA
left a comment
There was a problem hiding this comment.
AI-assisted review. Inline findings below.
| {}, | ||
| // Resizing Logic | ||
| nullptr)); | ||
| resize_index_select_channel_node)); |
There was a problem hiding this comment.
[High] resize_index_select_channel_node is not declared or defined, so this branch does not compile. Please add the callback and pass every value it needs through resize args.
AI reviewer note: this comment was generated by an AI reviewer.
There was a problem hiding this comment.
Right, and it was the whole CI failure: every red job was use of undeclared identifier on this line and line 110. resize_index_select_channel_node is now defined, and takes the index count from the index tensor in the resize args it already has.
| {}, | ||
| // Resizing Logic | ||
| nullptr)); | ||
| resize_fn)); |
There was a problem hiding this comment.
[High] resize_fn is also undeclared. Beyond the compile failure, the batch-axis path freezes a channel-derived stride in params at graph construction. A resize callback must recompute/update that stride when channels change, or explicitly reject such resizes; otherwise addressing is silently wrong. Add a below-bound dynamic-channel test.
AI reviewer note: this comment was generated by an AI reviewer.
There was a problem hiding this comment.
Both fixed. resize_fn is now resize_index_select_node, with the built dim passed through the resize args since the callback cannot recover it.
On the frozen stride: a resize callback gets only the graph, args and resize args, so it has no handle on the node params buffer and cannot rewrite it. Rather than rejecting the resize, the shader now derives the stride from sizes_ubo(in), which the runtime already keeps current across a resize, and the params buffer keeps only gpu_dim. Batch-dim index_select with a changing channel count therefore works instead of being refused.
Added the below-bound test you asked for: built at 8 channels, run at 1, 3, 4, 5 and 8, so the frozen stride (2 at build, 1 at C<=4) would read the wrong texel. Also covering the width dim, the channel dim, and embedding with a varying index length. I have no Vulkan device here, so I have verified all four partition to Vulkan under dynamic shapes and left execution to CI.
add_embedding_legacy_node() and both index_select node builders pass nullptr as their resizing logic, even though the matching resize functions are already defined right next to them. Under dynamic shapes their outputs therefore keep the extents they were built with instead of tracking the real input sizes, and downstream ops read them at the wrong size. Pass the resize functions that already exist, and mark index_select as supporting resize now that it actually does. Found while chasing a wrong-output bug in a TTS model on Adreno 840: the embedding output kept its upper-bound extents whenever a graph change moved the partition boundary so that a CPU round-trip no longer happened to re-establish the correct shape.
922bd83 to
b62a817
Compare
SS-JIA
left a comment
SS-JIA
left a comment
There was a problem hiding this comment.
Potentially valid comemnts from AI Review. RC'ing for now while CI runs.
The previous commit wired resize_index_select_channel_node and resize_fn without either existing, so the Vulkan backend did not build. Both are written here: the output takes the input's shape with the selected dim replaced by the number of entries in the index tensor. The dim the batch and width/height node was built for is passed through the resize args, since the callback cannot recover it otherwise. Selecting along the batch dim walks the z axis in units of channel texels, and that stride was computed from the channel count at graph construction and frozen in a params buffer. A resize callback has no way to rewrite a params buffer, so the shader derives the stride from the input sizes UBO instead, which the runtime already keeps current across a resize. The params buffer keeps only the dim. Adds dynamic-shape coverage for all four paths: the batch dim with the channel count varying below the built size, the width dim, the channel dim, and embedding with a varying index length.
3 participants
Summary
add_embedding_legacy_node()and bothindex_selectnode builders passnullptras their resizing logic:even though
resize_embedding_node,resize_index_select_channel_nodeand the localresize_fnare already defined immediately above them and do the right thing. Under dynamic shapes the outputs keep the extents they were built with (the upper bound) instead of tracking the real input sizes, so consumers read them at the wrong size.register_index_select()also does not setsupports_resize, so this patch sets it now that the op honours resize.Fix
Pass the resize functions that already exist. Three one-line changes plus the registry flag.
How it surfaced
A TTS model on Adreno 840 produced wrong outputs after an unrelated change moved a partition boundary. The embedding output had been keeping its upper-bound extents all along; previously a CPU round-trip happened to re-establish the correct shape downstream, so the defect was masked. Once the gather stayed on the GPU the stale extents reached the consumer.
Verified on a Galaxy S26 Ultra: with these wired up, the affected sub-models return correct shapes and match their CPU references (cosine >= 0.999 at sequence lengths well below the dynamic bound) where before they were wrong at every length except the bound itself.
Note this class of bug is hard to see with
executor_runner, which can only run at the dynamic upper bound -- exactly the one shape where a missing resize is harmless.