[WebGPU] Vectorize Split when every output segment is vec4-aligned - #32251
Open
Ananya Anand (4n4ny4) wants to merge 3 commits into
Open
[WebGPU] Vectorize Split when every output segment is vec4-aligned#32251Ananya Anand (4n4ny4) wants to merge 3 commits into
Ananya Anand (4n4ny4) wants to merge 3 commits into
Conversation
Vectorize eligible float32 splits by treating each axis slice as a contiguous segment and dispatching one vec4 copy per four elements. Map vectors to outputs with outer and within-segment offsets while preserving byte and dispatch counts. Keep the existing rank-aware scalar program as the fallback for non-float32 inputs or segments not divisible by four. Bind eligible tensors through flat 1-D storage views so complete contiguous segments can vectorize even when their trailing dimensions are not individually vec4-reducible. Frozen measurements show identical 22,118,400 bytes, throughput 84.375 -> 168.750 GB/s against NVML 652.8 GB/s (instruction-bound), paired GPU-resident p50 delta -0.175 ms, and natural wall delta -0.250 ms. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Contributor
Author
|
/run azp |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a vec4-optimized WebGPU Split path for aligned float32 segments.
Changes:
- Introduces a contiguous vec4 Split shader.
- Selects it for eligible float32 inputs.
- Adds a WebGPU Split test case.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
onnxruntime/core/providers/webgpu/tensor/split.h |
Declares the optimized program. |
onnxruntime/core/providers/webgpu/tensor/split.cc |
Implements dispatch and shader generation. |
js/web/test/data/ops/split.jsonc |
Adds aligned-split coverage. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| ] | ||
| }, | ||
| { | ||
| "name": "Split on Axis 0 - vec4 across trailing dimensions", |
Comment on lines
+78
to
+79
| std::vector<const ShaderVariableHelper*> outputs; | ||
| outputs.reserve(segment_vector_sizes_.size()); |
The vec4 test data only split on axis 0, where outer_index is always zero, so the per-outer output offset was never exercised. Add an axis 1 case over two outer slices with unequal aligned segments. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
The Web CI precheck runs prettier and fails on any resulting diff. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
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.
Description
SplitProgramcopies one element per invocation, and it works out where that element goes from scratch every time: flat offset to indices, read the split axis, loop over the cumulative split sizes to find the output, subtract the segment base, convert back to an offset. If the input is float32 and every output segment is a whole number of vec4s, none of that is necessary, because the split is just a contiguous partition of the buffer. This addsSplitContiguousVec4Programfor that case, picked when the input is float32 and everysplit_size * inner_sizedivides by 4. It treats the input and every output as vec4, so each invocation moves 16 bytes instead of 4 and the dispatch is a quarter the size, and it finds the destination with anifchain over segment bounds baked into the shader instead of a runtime loop. Anything that does not qualify, meaning non-float32 inputs or unaligned segments, still goes down the existing path untouched.segment_vector_sizesis the only thing the new shader bakes in and it is also the cache hint, so the pipeline key covers exactly what the WGSL depends on. The test is injs/web/test/data/ops/split.jsonc: a[2,4,3]float32 split on axis 0 givesinner_size12, so both segments are three whole vec4s and the case takes the new path.Motivation and Context
The index arithmetic in the generic path costs more than the copy it is wrapped around, and it gets worse as tensor rank and output count go up. Detection heads are where you notice it. YOLO26n splits repeatedly in its head, and those splits are float32 with channel counts that are already multiples of 4, so they qualify. Classifiers barely split at all, so I would not expect this to do much for them.