[WebGPU EP] Support uint8 for Expand - #29854
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
@hariharans29, PTAL, thanks! |
There was a problem hiding this comment.
Pull request overview
This PR extends the WebGPU Execution Provider’s Expand kernel to support uint8 tensors by reusing the existing packed-1-byte-type path (previously bool-only) and generalizing the shader logic to handle both types. It also adds unit tests intended to exercise the three packed-byte shader branches (whole-word copy, splat, and per-element pack).
Changes:
- Generalize the packed-byte
Expandshader path to support bothbool(vec4<bool>) anduint8(byte shift/mask + reassembly intou32). - Add
uint8to WebGPUExpandkernel type constraints. - Add
uint8Expandtest cases covering the packed-byte broadcast branches.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/providers/cpu/tensor/expand_test.cc | Adds uint8 Expand tests to exercise packed-byte broadcast variants. |
| onnxruntime/core/providers/webgpu/tensor/expand.cc | Extends the packed-byte path from bool to also handle uint8, and updates kernel type constraints accordingly. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Low
Review: PR #29854 — [WebGPU EP] Support uint8 for Expand (head
|
| Test | Input shape | Output shape | Input-last-dim | Output-last-dim | Shader branch |
|---|---|---|---|---|---|
Expand_3x3_uint8 |
{1} |
{3, 3} |
1 | 3 | Per-element assembly |
Expand_3x1_uint8 |
{3} |
{3, 3} |
3 | 3 | Per-element assembly |
Expand_1x3_uint8 |
{3, 1} |
{3, 3} |
1 | 3 | Per-element assembly |
Expand_1x4_uint8 |
{3, 1} |
{3, 4} |
1 | 4 | Splat |
Expand_4x1_uint8 |
{1, 4} |
{4, 4} |
4 | 4 | Whole-word copy |
All three shader branches (input_last_dim_divisible_by_4_ = whole-word; output_last_dim_divisible_by_4_ = splat; else = per-element pack) are exercised.
Byte-position robustness: Expand_4x1_uint8 uses distinct source bytes {10, 20, 30, 40} that fully populate a packed word — this catches any byte-lane ordering bug in the whole-word path even though that path doesn't invoke the arithmetic helpers, because it exercises the u8→u32 upload convention itself. Expand_1x4_uint8 uses distinct splat values (11, 22, 33) so a wrong splat idiom (e.g., e | (e << 4) instead of e * 0x01010101) would produce a corrupted 4-lane replica. Expand_3x1_uint8 (per-element assembly) constructs packed groups from up to 3 distinct source values (row {11, 22, 33} broadcast into 9-element flattened output), which catches pack_elements byte-shift bugs.
Non-blocking gap: No test constructs a 4-lane packed group from 4 distinct source bytes. Expand_3x3_uint8 (splat of 5) is the extreme opposite. A 1x8 broadcast-to-{2, 8} test with input {a, b, c, d, e, f, g, h} (all distinct) would give per-lane full-distinct-bytes coverage across two whole-word groups — worth considering as a follow-up but not blocking.
Copilot AI review
One comment, on the block comment above the tests:
"The new comment claims values vary within each 4-byte packed group, but some of the added cases (e.g., scalar broadcast) use a single repeated value. This makes the comment misleading relative to the actual coverage."
Copilot is technically right — Expand_3x3_uint8 (input {5}, output all 5s) doesn't help catch byte-position bugs. Copilot's proposed rewrite ("Some cases vary values within each packed u32 to catch byte-position bugs") is more accurate. Non-blocking comment nit, worth addressing in a squash-time tweak or in a follow-up.
Minor observations (non-blocking)
ComputeInternalcomment update. The old comment"// Check if either input is boolean"is now stale — the branch also handles uint8. The new comment"// bool and uint8 are 1-byte-per-element types that are not valid storage buffer types..."correctly replaces it. Good hygiene.- Cache-hint symmetry.
program.CacheHint(input_last_dim_divisible_by_4, output_last_dim_divisible_by_4)is insideif (is_packed_byte)and does NOT include a bit foris_boolvsis_uint8. Two different element types → two differentProgramVariableDataType::Boolx4/Uint8x4var_type values → two different program signatures → two different cache entries automatically (becausevar_typeis part of the program identity). So no need to hash the type discriminator into the cache hint explicitly. Correct. - Element-type in the shader.
input.GetByOffset(...)returnsu32for bothBoolx4andUint8x4(both alias u32 in the shader_variable table). Soget_elementreturns au32in both branches andpack_elementsproduces au32for uint8 (matchingoutput.SetByOffset(...)). Consistent. For bool,vec4<bool>(...)is produced and written through, using the vec4-of-bool storage convention that WebGPU-EP already handles for bool. splat_element(e)forbool.vec4<bool>(e)broadcasts a single bool component to a 4-component vector. WGSL supports this scalar-to-vec ctor. Preserved from the pre-refactor shader — no behavioral change.pack_elementsforboolstill usesvec4<bool>(e0, e1, e2, e3)— WGSL's 4-argument vec ctor, correct.- Merge choreography.
@hariharans29enabled/disabled auto-merge at the fetch time — the disable is probably to wait on the Copilot comment being addressed. Consider addressing that comment before re-enabling auto-merge.
Bottom line
Textbook uint8 support: shader-side byte arithmetic is correct little-endian, three shader branches all exercised by new tests, kernel-registration factory picks up uint8 for both opset ranges symmetrically, small diff. The only outstanding item is Copilot's one-comment nit about a slightly imprecise doc comment above the tests. Address that (or acknowledge it) and this is ready to merge. CI is 86/86 OK.
|
Please take a look at the copilot comment |
…ecific test The Boolx4 path works because GetByOffsetImpl has a Boolx4 special-case that returns vec4<bool>, enabling component indexing. Uint8x4 has no such special case - its value type is scalar u32, so component indexing is invalid WGSL. Fix: split pack_as_bytes into is_bool / is_uint8 branches: - is_bool: keep existing vec4<bool> component-access path - is_uint8: use bit-shift/mask to extract each byte from the packed u32, accumulate into output u32 via OR-shift (matching the standard Uint8x4 pattern established in PR microsoft#29854 for Expand) Also add a USE_WEBGPU-gated test that runs only on the WebGPU EP to ensure the shader path is actually exercised rather than silently falling back to CPU.
Done in 2e16d94, PTAL, thanks! @hariharans29 |
uint8 is a 1-byte-per-element type that cannot be stored directly in a WebGPU storage buffer, so it is packed 4-per-u32 just like bool. Extend the existing bool-only packed path to also handle uint8: route it through the same flatten/4-component/broadcast-indices setup, and generalize the shader so the three broadcast branches (whole-word copy, splat, and per-element pack) differ only in how a single element is extracted from and assembled into the packed word. bool keeps using vec4<bool>; uint8 uses byte shift/mask extraction and reassembly. Add uint8 to the Expand kernel type constraints and cover all three shader branches with tests.
Re-review: PR #29854 — [WebGPU EP] Support uint8 for Expand (head
|
The Uint8x4 SetByOffset expects the value as a vec4<u32> with one byte value per lane and packs it into the storage word itself. The Expand shader was instead handing it a pre-packed u32, which SetByOffset then broadcast across all four lanes and masked to the low byte, so every output element in a word collapsed to a single value. Build the 4-lane vector in every packed-byte branch (vec4<u32> for uint8, matching vec4<bool> for bool) and let SetByOffset do the packing. For the whole-word copy path, unpack the raw uint8 storage word with unpack4xU8 since its GetByOffset returns the packed u32 rather than the per-lane vector bool returns.
Head branch was pushed to by a user without write access
|
Tried to fix CI failure in latest commit 0fea358, PTAL, thanks! |
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
uint8 is a 1-byte-per-element type that cannot be stored directly in a WebGPU storage buffer, so it is packed 4-per-u32 just like bool. Extend the existing bool-only packed path to also handle uint8: route it through the same flatten/4-component/broadcast-indices setup, and generalize the shader so the three broadcast branches (whole-word copy, splat, and per-element pack) differ only in how a single element is extracted from and assembled into the packed word. bool keeps using vec4; uint8 uses byte shift/mask extraction and reassembly.
Add uint8 to the Expand kernel type constraints and cover all three shader branches with tests.
This is to fix part of #29756