Vulkan: support index.Tensor with the index on any dimension

[msluszniak]

Summary

x[:, :, idx] reaches the edge dialect as index.Tensor(x, [None, None, idx]). The partitioner only accepted the index on dim 0, so every other gather fell back to the CPU.

That is not just slower than it needs to be. An unsupported node in the middle of a graph splits the delegate in two, and a model that gathers once per block ends up with one delegate call per block. A TTS vector estimator I was lowering came out as 29 delegate calls, 28 of them created by this single op. On a Galaxy S26 Ultra each of those boundaries costs about 3 ms in staging and submission, measured with a synthetic 4-block A/B.

Gathering along a dimension is exactly what index_select does, and the Vulkan implementation of it already handles every dimension. So dispatch there when the index does not sit on dim 0, and leave the index_tensor shader, which gathers along the leading dim, alone.

Two things were in the way of simply accepting the node

The dimension. It is recoverable at runtime: the serialized indices list keeps a null for every dimension that is not indexed, so the position of the tensor entry is the dimension. The list pointer has to be released before add_scalar(), which may reallocate the graph's value store (check_no_active_value_ptrs catches this).

The index tensor never got a representation. It lives inside a list whose other entries are None, and both is_tensor_arg_node() and is_non_constant_tensor_node() only recognise lists that are entirely tensors, so the whole argument was treated as a non-tensor and skipped by TagMemoryMetaPass. index_select reads channels-packed textures and asserts on anything else.

This adds tensor_nodes_in_arg(), which returns the tensor nodes of an argument and skips the Nones, and uses it everywhere the tagging pass walks a list argument. Lists that are entirely tensors, which is every other operator today, behave exactly as before.

Test plan

New test_vulkan_backend_index_tensor_non_leading_dim in backends/vulkan/test/test_vulkan_delegate.py, alongside the existing dim-0 test.

Verified on a Galaxy S26 Ultra (Adreno 840). A minimal x[:, :, idx] * 2 model aborts on check_index_select_args before this change and matches its CPU reference after (max abs diff 6e-4, fp16). The TTS model that motivated it goes from 29 delegate calls to 2 with unchanged outputs:

sub-model	delegate calls	cosine vs CPU
duration_predictor	11 -> 6	1.000000
text_encoder	15 -> 10	0.999414
vector_estimator	29 -> 2	0.999994
vocoder	14 -> 3	0.999977

End to end that is 11% faster on device, measured with the arms alternated forward and reverse inside each round to cancel thermal drift.

Depends on nothing, but note that lowering this model also needs #22403; without it FuseViewCopyTransform trips over the larger partition.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit 7436c03 with merge base 5428092 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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.