Skip to content

[WebGPU EP] Support uint8 for Expand - #29854

Merged
hariharans29 merged 2 commits into
microsoft:mainfrom
mingmingtasd:expand
Jul 29, 2026
Merged

[WebGPU EP] Support uint8 for Expand#29854
hariharans29 merged 2 commits into
microsoft:mainfrom
mingmingtasd:expand

Conversation

@mingmingtasd

Copy link
Copy Markdown
Contributor

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

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

@mingmingtasd

mingmingtasd commented Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

@hariharans29, PTAL, thanks!
cc/ @huningxin @ibelem @adrastogi

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Expand shader path to support both bool (vec4<bool>) and uint8 (byte shift/mask + reassembly into u32).
  • Add uint8 to WebGPU Expand kernel type constraints.
  • Add uint8 Expand test 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

Comment thread onnxruntime/test/providers/cpu/tensor/expand_test.cc Outdated
@hariharans29

Copy link
Copy Markdown
Member

Review: PR #29854 — [WebGPU EP] Support uint8 for Expand (head b305de1)

Author: @mingmingtasd. 1 commit. CI: 86/86 checks OK. Copilot AI reviewer: 1 comment (nit on a test-block comment). @hariharans29 briefly enabled auto-merge and then disabled it — presumably waiting on the Copilot nit or a second look. Referenced from issue #29756 (WebNN demo umbrella).

Verdict: approve. Clean refactor that generalizes the existing bool-only "packed 1-byte type" path in the WebGPU EP Expand kernel to also handle uint8_t. The bit-shift/mask arithmetic is correct little-endian byte packing, all three shader branches are covered by tests, and the code change is small (~40 net lines across two files).


What it does

Two-file change in onnxruntime/core/providers/webgpu/tensor/expand.cc and onnxruntime/test/providers/cpu/tensor/expand_test.cc:

  1. Kernel shader (ExpandProgram::GenerateShaderCode) — introduces three C++ lambdas that abstract the "single-element extract / splat / 4-element pack" primitives, keyed on is_bool vs is_uint8:

    const bool is_bool  = Inputs()[0].var_type == ProgramVariableDataType::Boolx4;
    const bool is_uint8 = Inputs()[0].var_type == ProgramVariableDataType::Uint8x4;
    if (is_bool || is_uint8) {
      auto get_element   = [&](off)          { /* vec4<bool>[k] for bool ; shift+mask for uint8 */ };
      auto splat_element = [&](e)            { /* vec4<bool>(e) for bool ; e * 0x01010101u for uint8 */ };
      auto pack_elements = [&](e0,e1,e2,e3)  { /* vec4<bool>(e0..e3) for bool ; OR-shift for uint8 */ };
      // then the same three branches (whole-word / splat / per-element) as before,
      // rewritten to call these helpers instead of hardcoding vec4<bool>.
    }
  2. Kernel dispatch (Expand::ComputeInternal) — adds is_uint8 alongside is_bool and introduces bool is_packed_byte = is_bool || is_uint8. All existing if (is_bool) gates become if (is_packed_byte). The components_i / components_o computation, the program.AddInputs(..., ProgramInput::Flatten, ...) path, the AddIndices(input_shape) guard, and the CacheHint on input_last_dim_divisible_by_4 / output_last_dim_divisible_by_4 all extend to uint8 by construction — no separate code path.

  3. Kernel-info factories (CreateExpandVersionedKernelInfo, CreateExpandKernelInfo) — copies the shared GetOpTypeConstraints(enable_int64, /*enable_bool=*/true) static vector into a local std::vector<MLDataType> and appends DataTypeImpl::GetTensorType<uint8_t>(). Registered for both opset ranges <8, 12> and <13>.

  4. Tests — 5 new TEST(ExpandOpTest, Expand_*_uint8) cases, all no-EP-restriction test.Run() (works fine because Expand is well-supported on CPU and this PR is additive on WebGPU).


Correctness — the packed-uint8 arithmetic

The uint8 lane-arithmetic helpers use standard little-endian byte-in-word packing that matches the way ORT already packs host-side uint8_t[] into storage-buffer u32s (the ProgramVariableDataType::Uint8x4 mapping in program.cc and the u32 element type in shader_variable.cc).

  • get_element(off)((packed[off/4] >> ((off % 4) * 8u)) & 0xFFu). Byte k of a u32 lives at bits k*8 .. k*8+7. Right-shift by k*8, mask with 0xFF. Textbook. ✓
  • splat_element(e)(e * 0x01010101u). For e ∈ [0, 255], produces 0xEEEEEEEE where each byte lane equals e. Standard 4-byte splat idiom, no branch. ✓
  • pack_elements(e0, e1, e2, e3)(e0 | (e1 << 8u) | (e2 << 16u) | (e3 << 24u)). Little-endian little-to-big lane assembly. Each eN is already 0xFF-masked by get_element, so no cross-lane bleeding. ✓

Cross-checked against the ProgramVariableDataType::Uint8x4 machinery in the WebGPU EP — that already treats u8x4 as a u32 at the shader level and expects the same byte-in-word layout for uploads.

Splat verification for Expand_1x4_uint8: input {11, 22, 33} broadcast to {4, 4} output; each output row is a splat of one value. For value 22: 22 * 0x01010101 = 0x16161616. Extracting any lane yields 22. Test's expected output matches: {22, 22, 22, 22}. ✓

Whole-word copy for Expand_4x1_uint8: input {10, 20, 30, 40} at word-offset 0 = 0x281E140A little-endian. Whole-word path writes this u32 as-is into each output row. Test's expected row {10, 20, 30, 40} reads back as bytes 0/1/2/3 of that same word. ✓


is_packed_byte refactor in ComputeInternal

The single-flag consolidation (is_bool || is_uint8) is applied uniformly:

  • components_i = (is_packed_byte || input_last_dim_divisible_by_4) ? 4 : 1; — for a packed-byte type the shader always operates on 4-per-u32 packed words, so components must be 4 unconditionally.
  • components_o = ... (same).
  • if (is_packed_byte) selects the ProgramInput::Flatten + AddIndices(input_shape) branch that the shader relies on for input_indices / output_indices variables.
  • if (is_packed_byte || components_i != components_o) program.AddIndices(std::move(output_shape)); — same as the bool branch's need for output_indices in the packed shader.

All the shape-arithmetic guards that were already is_bool-aware (e.g., the input_shape.IsScalar() / is_int64 predicates for the input_last_dim_divisible_by_4 flag) are unchanged and continue to apply because uint8 tensors — like bool tensors — are never marked is_int64 and can be scalars just as bool can. Consistent.


Kernel-registration change

std::vector<MLDataType> type_constraints = GetOpTypeConstraints(enable_int64, true);
type_constraints.push_back(DataTypeImpl::GetTensorType<uint8_t>());

Two subtle-but-correct choices:

  1. Copy, not reference. GetOpTypeConstraints returns const std::vector<MLDataType>& to a static inside the function — the assignment triggers a copy, so push_back(uint8_t) mutates the local copy and never contaminates the shared static across other Expand-like ops. Correct.
  2. No duplicate uint8_t. webgpu_supported_types.h's base set is {fp16, fp32, int32, uint32} and the enable_bool=true branch adds bool (plus int64 if enable_int64). None of them include uint8_t, so the push_back doesn't create a duplicate.

Both factory functions (CreateExpandVersionedKernelInfo for opset 8-12 and CreateExpandKernelInfo for opset 13+) get the identical change. Symmetry preserved with the Sub/Equal/Expand conditional-registration pattern.


Test coverage — hits all three shader branches

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)

  1. ComputeInternal comment 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.
  2. Cache-hint symmetry. program.CacheHint(input_last_dim_divisible_by_4, output_last_dim_divisible_by_4) is inside if (is_packed_byte) and does NOT include a bit for is_bool vs is_uint8. Two different element types → two different ProgramVariableDataType::Boolx4 / Uint8x4 var_type values → two different program signatures → two different cache entries automatically (because var_type is part of the program identity). So no need to hash the type discriminator into the cache hint explicitly. Correct.
  3. Element-type in the shader. input.GetByOffset(...) returns u32 for both Boolx4 and Uint8x4 (both alias u32 in the shader_variable table). So get_element returns a u32 in both branches and pack_elements produces a u32 for uint8 (matching output.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.
  4. splat_element(e) for bool. 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.
  5. pack_elements for bool still uses vec4<bool>(e0, e1, e2, e3) — WGSL's 4-argument vec ctor, correct.
  6. Merge choreography. @hariharans29 enabled/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.

@hariharans29

Copy link
Copy Markdown
Member

Please take a look at the copilot comment

Copilot AI added a commit to huningxin/onnxruntime that referenced this pull request Jul 24, 2026
…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.
@mingmingtasd

Copy link
Copy Markdown
Contributor Author

Please take a look at the copilot comment

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.
@hariharans29

Copy link
Copy Markdown
Member

Re-review: PR #29854 — [WebGPU EP] Support uint8 for Expand (head 2e16d94)

Author: @mingmingtasd. Force-pushed once since my previous approval: b305de12e16d94 ("Done in 2e16d94" per author, addressing the Copilot AI comment). CI: 1 / 1 OK on the new head (still spinning up, but the earlier b305de1 was 86/86 OK — the delta doesn't touch anything that could plausibly regress). 3 participants: @mingmingtasd, @hariharans29, Copilot. Auto-merge was enabled 3 days ago and then disabled, presumably to wait on the Copilot comment being addressed — which it now is.

Verdict: approve. The only substantive change since the last pass is a two-word tweak to the doc comment above the uint8 tests, exactly what Copilot AI asked for. Everything else is byte-identical to the previously-approved b305de1 — the shader arithmetic, the is_packed_byte refactor in Expand::ComputeInternal, the kernel-registration factory changes, and all five tests are unchanged.


Delta from b305de12e16d94: just the doc comment

Copilot AI's one comment on the previous head was:

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.

That comment previously read:

// These cases cover the three packed-byte shader paths: per-element assembly (neither last dim
// divisible by 4), splat (output last dim divisible by 4, input last dim 1), and whole-word copy
// (input last dim divisible by 4). Values vary within each 4-byte group to catch byte-position bugs.

In 2e16d94 it now reads:

// These cases cover the three packed-byte shader paths: per-element assembly (output last dim not
// divisible by 4), splat (output last dim divisible by 4, input last dim 1), and whole-word copy
// (input last dim divisible by 4). Some cases vary values within each packed u32 to catch
// byte-position bugs.

Two tightening changes:

  1. "neither last dim divisible by 4""output last dim not divisible by 4" — matches the actual branch condition. The old wording was ambiguous (could be read as "neither input nor output") and the pack path actually fires whenever the output-last-dim isn't divisible by 4, regardless of the input-last-dim.
  2. "Values vary within each 4-byte group to catch byte-position bugs""Some cases vary values within each packed u32 to catch byte-position bugs" — matches reality (Expand_3x3_uint8 uses {5} broadcast, so its 4-byte groups are all {5, 5, 5, 5}; only Expand_1x4_uint8 with distinct splat values {11, 22, 33} and Expand_4x1_uint8 with the whole {10, 20, 30, 40} word actually exercise distinct byte lanes).

Copilot's exact suggested phrasing was ("Some cases vary values within each packed u32 to catch byte-position bugs") — adopted verbatim. Clean resolution.


Everything from the pass-1 approval stands

Re-verified in the fetched diff against 2e16d94 — no changes to any of:

  • Pack in expand.cc: uint8 splat_element(e) = (e * 0x01010101u), pack_elements(e0..e3) = (e0 | (e1 << 8u) | (e2 << 16u) | (e3 << 24u)), get_element(off) = ((data[off/4] >> ((off % 4) * 8u)) & 0xFFu). Little-endian byte packing. ✓
  • is_packed_byte refactor in Expand::ComputeInternalis_bool || is_uint8 gating replaces all pre-existing is_bool gates. ✓
  • Kernel-info factoriesCreateExpandVersionedKernelInfo<8, 12> and CreateExpandKernelInfo<13> both copy GetOpTypeConstraints(...) and push_back(uint8_t). No duplicate uint8 (the base helper doesn't include it). ✓
  • Tests — same 5 _uint8 tests exercising per-element / splat / whole-word branches on shapes {3×3}, {3×1}, {1×3}, {1×4}, {4×1}. Test bodies unchanged. ✓
  • Cache-hint symmetryinput_last_dim_divisible_by_4 / output_last_dim_divisible_by_4 unchanged, and the automatic ProgramVariableDataType component of program identity still separates the Boolx4 and Uint8x4 shader binaries in the cache. ✓

Correctness of 2e16d94 inheriting from the byte-identity of b305de1

Since the shader-code emitter and the kernel-registration changes are byte-identical between b305de1 and 2e16d94, and the emitted WGSL is a function of only the C++ source (the doc comment above the tests doesn't reach the shader), the runtime behavior of 2e16d94 is identical to b305de1 on every input for every WebGPU adapter. b305de1 was 86/86 CI OK, so 2e16d94 should reach the same green state once its own CI legs finish.


Merge-state note

  • Auto-merge was disabled 3 days ago while the Copilot comment was open.
  • With the comment now addressed at 2e16d94, @hariharans29 can safely re-enable auto-merge (or approve + merge directly).
  • No new review comments have been posted since @mingmingtasd pushed the fix 16 hours ago.

Bottom line

2e16d94 is a strictly-more-accurate version of the previously-approved b305de1 — the only substantive change is a two-word doc-comment tweak that exactly resolves Copilot AI's one blocking-adjacent nit. Ready for @hariharans29 to re-enable auto-merge and land.

@hariharans29
hariharans29 enabled auto-merge (squash) July 27, 2026 18:01
hariharans29
hariharans29 previously approved these changes Jul 27, 2026
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.
auto-merge was automatically disabled July 28, 2026 05:10

Head branch was pushed to by a user without write access

@mingmingtasd

Copy link
Copy Markdown
Contributor Author

Tried to fix CI failure in latest commit 0fea358, PTAL, thanks!

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

@hariharans29
hariharans29 enabled auto-merge (squash) July 28, 2026 21:06
@hariharans29
hariharans29 merged commit 90cc3aa into microsoft:main Jul 29, 2026
176 of 255 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants