[ET-VK] Wire the missing resize functions for embedding and index_select - #22410
Open
msluszniak wants to merge 1 commit into
Open
[ET-VK] Wire the missing resize functions for embedding and index_select#22410msluszniak wants to merge 1 commit into
msluszniak wants to merge 1 commit into
Conversation
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.
🔗 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.
|
This PR needs a
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.