[WebGPU EP] Support Cast to and from uint8

[mingmingtasd]

Add uint8 support to the WebGPU Cast in both directions, so uint8 Cast nodes stop falling back to the CPU EP. uint8 was in neither the input (T1) nor output (T2) type constraint, so both casts to uint8 (e.g. a terminal bool->uint8 at a graph output) and casts from uint8 (back to int32/float/bool) fell back.

WebGPU stores uint8 packed as Uint8x4 (4 per u32, lane 0 -> low byte), the same layout already used for bool.

• to uint8 (cast.cc, shader_variable.cc): add uint8 to the output (T2) type constraint, a to-uint8 shader expression, and a Uint8x4 packing path in SetByOffset.

• from uint8 (cast.cc, shader_variable.cc): add uint8 to the input (T1) type constraint, a Uint8x4 unpacking path in GetByOffset (reads the 4 bytes into a vec4, inverse of the packing), and include uint8 in is_from_unsigned so uint8->int64 zero-extends.

This is a split CL from #29682 to fix part of #29756

[azure-pipelines]

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

[mingmingtasd]

@adrastogi Please help trigger the copilot review and testing workflows, thanks!
@hariharans29, PTAL, thanks!
cc/ @huningxin @ibelem @xhcao

[Copilot]

Pull request overview

This PR extends the WebGPU Execution Provider’s Cast kernel to support uint8 as both an input (T1) and output (T2) type, eliminating unnecessary CPU EP fallback for uint8 cast nodes (notably in WebNN-exported graphs). It also implements the correct packed Uint8x4 load/store behavior in shader variable helpers and adds targeted WebGPU EP tests.

Changes:

• Add uint8 to WebGPU Cast kernel type constraints for both inputs and outputs.
• Implement Uint8x4 unpacking (GetByOffset) and packing (SetByOffset) paths so uint8 tensors use the intended packed storage layout.
• Add WebGPU EP unit tests covering bool -> uint8 and uint8 -> {int32,float,bool} casts, including non-multiple-of-4 sizes.

Show a summary per file

File	Description
onnxruntime/test/providers/cpu/tensor/cast_op_test.cc	Adds explicit WebGPU EP tests validating uint8 cast behavior and packed tail handling.
onnxruntime/core/providers/webgpu/tensor/cast.cc	Enables uint8 in Cast type constraints and ensures uint8 participates in unsigned-to-int64 zero-extension logic.
onnxruntime/core/providers/webgpu/shader_variable.cc	Adds Uint8x4 unpack/pack implementations matching the packed byte layout used by the WebGPU EP.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0
• Review effort level: Low

[huningxin]

Thanks for adding uint8 Cast support to the WebGPU EP!

This is also important for WebNN ORT implementation when using WebGPU EP as its backend. The WebNN logical operators use uint8 to represent boolean values, whereas ONNX uses bool. When the Chromium WebNN ORT backend maps these ops, it inserts Cast nodes to convert between bool and uint8 at the boundaries of WebNN subgraphs. Without uint8 Cast support in the WebGPU EP, these inserted Cast nodes would fall all the way back to the CPU EP, forcing expensive GPU→CPU→GPU data transfers and breaking the end-to-end GPU execution pipeline.

[hariharans29]

Review: PR #29839 — [WebGPU EP] Support Cast to and from uint8 (head b6024f6)

Author: @mingmingtasd. 1 commit. CI: 77 / 86 checks OK. 3 files changed. Copilot AI review: 0 comments (Low effort). Positive endorsement from @huningxin confirming the concrete motivation (Chromium's WebNN ORT backend inserts bool↔uint8 casts at subgraph boundaries because WebNN uses uint8 for booleans; without this PR those casts fall back to CPU and force GPU↔CPU↔GPU thrashing). Split from PR #29682, closes part of #29756.

Verdict: approve. Small, tight extension of the existing Cast op. Correctness verifiable by direct comparison to the pre-existing Boolx4 packing path (uint8 uses the same 4-bytes-per-u32 storage layout), with one subtle-but-necessary difference — an explicit per-lane byte mask on the pack that Boolx4 didn't need.

---

What it does

Three coordinated changes in onnxruntime/core/providers/webgpu/tensor/cast.cc, onnxruntime/core/providers/webgpu/shader_variable.cc, and onnxruntime/test/providers/cpu/tensor/cast_op_test.cc:

1. T1 / T2 type constraints. Adds uint8_t to both the input (T1) and output (T2) constraint lists in CreateCastKernelInfo. Prior state: uint8_t was in neither, so uint8→anything and anything→uint8 both fell back to the CPU EP.
2. Uint8x4 pack / unpack in ShaderVariableHelper::SetByOffsetImpl and GetByOffsetImpl.
3. is_from_unsigned extended so uint8→int64 zero-extends (rather than sign-extends).

Plus a shader-code case for UINT8 output that widens the source to vec4<u32> — the actual byte packing is delegated to SetByOffsetImpl(Uint8x4).

---

Correctness — pack / unpack against the Boolx4 baseline

Pack (SetByOffsetImpl):

// Boolx4 (existing):
name[offset] = dot(vec4<u32>(0x1, 0x100, 0x10000, 0x1000000), vec4<u32>(value));

// Uint8x4 (new):
name[offset] = dot(vec4<u32>(0x1u, 0x100u, 0x10000u, 0x1000000u),
                   (vec4<u32>(value) & vec4<u32>(0xFFu)));

For a vec4<u32> value = (v0, v1, v2, v3):

• After & 0xFFu: each lane is truncated to a byte.
• Dot product: v0 + (v1 << 8) + (v2 << 16) + (v3 << 24) — little-endian byte pack, lane 0 → low byte.

The extra & 0xFFu mask is load-bearing. Boolx4 gets away without it because bool values are 0/1 by construction, so no lane's dot-product contribution can overflow into the next byte's field. For real uint8 values (e.g. 128, 255), the v0*1 term for v0 = 255 produces 255 (fine, fits in byte 0), but v1 = 255 produces 255*256 = 0xFF00 (also fine, byte 1 exactly). Actually, without the mask, v1 = 256 would spill into byte 2 — and while ONNX uint8 values are in [0, 255] by contract, the vec4 lanes could carry values > 255 depending on how the input expression widens (vec4<u32>(a) where a is a wider type). The mask makes the pack defensive against any lane value that isn't already byte-limited, which is the right posture given the shader-code case just does expression = "vec4<u32>(a)" and leaves the byte-truncation semantics to the setter. Comment on the new case documents this rationale explicitly. Good.

Unpack (GetByOffsetImpl):

// Boolx4 (existing):
vec4<bool>(
  bool(name[offset] & 0xFFu),
  bool(name[offset] & 0xFF00u),
  bool(name[offset] & 0xFF0000u),
  bool(name[offset] & 0xFF000000u))

// Uint8x4 (new):
vec4<u32>(
  name[offset] & 0xFFu,
  (name[offset] >> 8u) & 0xFFu,
  (name[offset] >> 16u) & 0xFFu,
  (name[offset] >> 24u) & 0xFFu)

Boolx4 gets to short-cut with bool(byte_slice_word) because "any nonzero → true". Uint8x4 needs the actual byte value, so shift-then-mask is the right form. Correct inverse of the pack.

---

is_from_unsigned extension

bool is_from_unsigned = input_tensor->DataType() == DataTypeImpl::GetType<uint32_t>() ||
                        input_tensor->DataType() == DataTypeImpl::GetType<uint8_t>() ||
                        input_tensor->DataType() == DataTypeImpl::GetType<bool>();

Governs whether the widening-to-int64 shader emits a zero-extension or a sign-extension. For uint8 (unsigned) source, zero-extension is the correct semantic (a uint8 value of 255 should become int64(255), not int64(-1)). Correctly grouped with uint32_t and bool. No test in the PR exercises uint8→int64 directly (see coverage note below), but the code path is trivially symmetric with uint32_t→int64 which is exercised elsewhere.

---

Shader-code case for UINT8 output

case ONNX_NAMESPACE::TensorProto_DataType_UINT8:
  // Output uint8 is stored packed (4 bytes per u32); the SetByOffset for Uint8x4 does the
  // packing, so here we just widen to a vec4<u32> of per-lane values (e.g. bool 0/1 -> uint8).
  expression = "vec4<u32>(a)";
  break;

Identical form to the pre-existing UINT32 case — the shader-code side is a plain widen to vec4<u32>, and the Uint8x4 setter takes care of the byte packing. Clean separation of concerns. The comment on the SetByOffsetImpl case explicitly notes the mask handles values > 1, so this widen-without-mask at the shader level is safe.

---

Type-constraint mechanics

// before (kept const-ref):
const auto& t1_constraints = GetOpTypeConstraints(/*enable_int64=*/enable_int64, /*enable_bool=*/true);

// after (takes a copy to append):
std::vector<MLDataType> t1_constraints = GetOpTypeConstraints(...);
t1_constraints.push_back(DataTypeImpl::GetTensorType<uint8_t>());

Assuming GetOpTypeConstraints returned a const reference to a static-lifetime vector previously (which the const auto& binding permitted), the new code copies once at kernel-registration time. That's session-init cost, not per-inference cost — negligible. If in the future the same pattern gets used for many more added types, factoring GetOpTypeConstraints to accept an "extras" list would be cleaner, but not worth doing for one type.

---

Test coverage

Four new tests, all guarded with GTEST_SKIP() when the WebGPU EP isn't in the build, all using the standard explicit std::vector<std::unique_ptr<IExecutionProvider>> + push_back + test.Run(..., &execution_providers) pattern:

Test	Direction	Notes
BoolToUint8_WebGpu	bool → uint8	Exercises the Uint8x4 packing path from a bool source.
Uint8ToInt32_WebGpu	uint8 → int32	Full byte range {0, 1, 128, 255, 42, 7}. Exercises Uint8x4 unpacking + widening.
Uint8ToFloat_WebGpu	uint8 → float	Same input values, exercises the float-conversion arm with unpacked uint8.
Uint8ToBool_WebGpu	uint8 → bool	Inverse of test 1. Exercises Uint8x4 unpack and Boolx4 output packing, plus the "any nonzero → true" cast-to-bool semantic.

All shapes are {2, 3} (6 elements), deliberately not a multiple of 4 — the comment on each test documents that this exercises the final partial packed word. This is the crucial edge case for a pack-4-into-1 storage layout.

Not covered (minor gaps, not blocking):

• uint8 → int64 — the code path is enabled by the is_from_unsigned addition, but no test pins the zero-extension semantic. If a regression flipped it to sign-extension, existing uint32→int64 coverage would catch the general pattern but not the specific uint8 slice.
• uint8 → uint8 — unusual (identity cast) but legal.
• Multiple-of-4 shape — the current 6-element shape does one full vec4 + a 2-element tail, so the pure-vectorized loop is exercised transitively, but a dedicated 4- or 8-element case would be tighter.

Coverage-as-is is adequate for the load-bearing behaviors (pack, unpack, tail).

---

Copilot review disposition

0 comments generated. Nothing to reconcile.

---

CI status

77 / 86 checks OK on the single commit. The 9 not-OK are likely still pending; the code doesn't look like a plausible cause of new-regression failures at first pass. Worth confirming they land green before merge.

---

Minor observations (none blocking)

1. The pack's & 0xFFu mask is a per-lane bitwise-AND on a vec4 before the dot product — one extra SIMD op per Uint8x4 write. Trivial cost, well worth the safety.
2. The comment "lane 0 → low byte" is called out on both the pack and unpack sites, which makes the byte-endianness of the storage explicit and testable. Good discipline.
3. The test comments deliberately explain why the 6-element shape was chosen (partial packed word) — this makes future maintainers less likely to "simplify" the tests to convenient 4- or 8-element shapes and accidentally erase the tail-coverage.
4. Fifth-order nit: the GetByOffsetImpl for Boolx4 uses inconsistent-width masks (0xFFu, 0xFF00u, 0xFF0000u, 0xFF000000u — implicit u32) whereas the new Uint8x4 uses >> 8u & 0xFFu etc. Different styles for adjacent cases in the same switch, but each is internally consistent and the Uint8x4 form is genuinely different math. Not worth reconciling.

---

Bottom line

Clean, small, well-motivated PR that closes a real fallback path with concrete downstream impact (Chromium's WebNN ORT backend, confirmed by @huningxin). Pack/unpack correctness is verifiable by direct comparison to the Boolx4 baseline plus the necessary byte-mask on the pack. Test coverage exercises the load-bearing byte-packing edge cases including the partial-word tail. Approve after CI settles.

[mingmingtasd]

@hariharans29
Pushed ee81e5c to try to fix CI issues, PTAL again, thanks!

[hariharans29]

Re-review: PR #29839 — [WebGPU EP] Support Cast to and from uint8 (head ee81e5c)

Author: @mingmingtasd. Force-pushed once since the last review: b6024f6 → ee81e5c ("try to fix CI issues"). CI: 85 / 86 checks OK (up from 77/86 on the previous head — improvement, not regression). Copilot AI still shows 0 comments. Endorsements from @huningxin (confirms the concrete Chromium WebNN-ORT motivation) and @hariharans29 (reviewer engagement).

Verdict: approve. The re-fetch shows one substantive design shift since the last pass, plus a plumbing addition — both changes are net improvements and align the Uint8x4 read convention with what the rest of the WebGPU EP is doing.

---

What changed since the previous head (b6024f6 → ee81e5c)

1. Unpack moved from ShaderVariableHelper::GetByOffsetImpl to the caller (via WGSL unpack4xU8).

The previous head had GetByOffsetImpl for Uint8x4 emit the four-shift-and-mask unpack:

vec4<u32>(name[i] & 0xFFu, (name[i] >> 8u) & 0xFFu, (name[i] >> 16u) & 0xFFu, (name[i] >> 24u) & 0xFFu)

The new head has shader_variable.cc's Uint8x4 fall through to the default case:

default:
  // Uint8x4 falls through here intentionally: GetByOffset returns the raw packed u32 storage
  // word, matching the convention other kernels rely on for byte-packed uint8 tensors (they
  // unpack sub-byte fields themselves, e.g. via unpack4xU8). Callers that want the 4 unpacked
  // byte values apply unpack4xU8 at the use site.
  ss << name_ << "[" << offset << "]";

Cast then wraps the raw word with unpack4xU8(...) at both call sites (int64-output and generic-output paths) in cast.cc:

const std::string load_a = is_from_uint8_ ? "unpack4xU8(" + input.GetByOffset("global_idx") + ")"
                                          : input.GetByOffset("global_idx");

2. New is_from_uint8 flag on CastProgram (declared in cast.h) plumbed through the constructor to gate the unpack4xU8 wrap.

Everything else on both files — the T1/T2 constraint appends, the is_from_unsigned extension, the UINT8 output shader-code case (vec4<u32>(a)), the SetByOffsetImpl pack with the load-bearing & 0xFFu mask, and all four tests — is unchanged from the previous head.

---

Why this design shift is an improvement

Alignment with the other Uint8x4 consumers in the WebGPU EP. PR #29854 (Expand for uint8) explicitly avoided the GetByOffset → vec4 route and instead shifted-and-masked the raw u32 inline, exactly because Expand's per-lane logic only wants one byte at a time. The previous head of this PR would have made Uint8x4 an inconsistent case — auto-vec4 for Cast, raw-u32 for Expand. The new convention is uniform: GetByOffset always returns the raw storage word for byte-packed types, and callers unpack how they need (unpack4xU8 for the whole vec4, or (x >> (k*8u)) & 0xFFu for a single byte). This also unblocks a clean fix for PR #29861 (Gather uint8), which had a shader-correctness bug in an earlier revision precisely because it assumed GetByOffset would return a vec4 for Uint8x4.

WGSL builtin over manual unpack. unpack4xU8 is a single WGSL 1.1 builtin (defined in the "Data packing built-in functions" section of the WGSL spec) that maps to platform-optimal byte-decompress paths on the underlying HAL — Dawn lowers it to native SIMD-friendly ops where available, rather than four separate shift-and-mask operations. Semantically identical to the old four-shift-and-mask code (unpack4xU8(x)[k] == (x >> (k*8u)) & 0xFFu for all k ∈ [0, 3], lane 0 → low byte), so the byte-endianness contract is preserved. Also cleaner shader source (one call instead of a four-element vec4 constructor with four different bit expressions).

Portability footnote. unpack4xU8 requires a WGSL implementation that has landed the 1.1 data-packing builtins — Dawn added support around Chrome 122/123 timeframe (early 2024). The ORT WebGPU EP bundles Dawn at a pinned version, and the CI is 85/86 OK, so the pinned Dawn has it. Any downstream consumer that swaps in an older Dawn or a non-Dawn WGSL implementation would need to check. Not a real concern today — noting it for the record.

---

Correctness — re-verified against the design shift

Pack (unchanged): shader_variable.cc's SetByOffsetImpl for Uint8x4 still does:

name[offset] = dot(vec4<u32>(0x1u, 0x100u, 0x10000u, 0x1000000u), (vec4<u32>(value) & vec4<u32>(0xFFu)));

Lane 0 → low byte, byte-mask on every lane defends against source expressions where a lane isn't already ≤ 255. Same analysis as the previous review; still correct.

Unpack (new path): unpack4xU8(raw_u32) returns vec4<u32>(raw & 0xFFu, (raw >> 8u) & 0xFFu, (raw >> 16u) & 0xFFu, (raw >> 24u) & 0xFFu). Same byte order (lane 0 → low byte), same zero-extension into u32 (uint8 semantics). Perfect inverse of the pack.

Cast expression composition:

• is_from_uint8_ = true → let a = unpack4xU8(input[global_idx]); → a is vec4<u32>.
• Non-int64-output branch: applies expression (e.g. vec4<f32>(a), vec4<bool>(a), vec4<u32>(a)) — all valid WGSL conversions from vec4<u32> and produce the correct per-lane semantics (u32 → float promotion for float, nonzero-check for bool, identity for u32).
• int64-output branch: reads a.x / a.y / a.z / a.w — component access on vec4<u32>, valid.

For non-uint8 inputs, GetByOffset returns whatever the pre-existing convention was (typically vec4<T> for the input's element type at 4 components), so a is already vec4<T> and no change is needed. The is_from_uint8_ branch is the only new code path.

Bit-identity with the previous head: since unpack4xU8 is definitionally equivalent to the old manual four-shift-and-mask, the emitted shader semantics are exactly the same across the two heads. The test outputs are unchanged, confirming this.

---

Registration & flag plumbing (verified)

• CreateCastKernelInfo<S, E> still copies GetOpTypeConstraints(enable_int64, /*enable_bool=*/true) and push_back(uint8_t) for both t1_constraints and t2_constraints. Comments clearly justify each append.
• is_from_uint8 computed from input_tensor->DataType() == DataTypeImpl::GetType<uint8_t>().
• is_from_uint8 also OR'd into is_from_unsigned — so uint8 → int64 will zero-extend, correct.
• CastProgram constructor updated in cast.h to take/store is_from_uint8_. All four flags (is_from_int64_, is_from_float_, is_from_unsigned_, is_from_uint8_) initialized in the mem-init list in the same order they appear as constructor parameters. Clean.

---

Test coverage — unchanged, still solid

Four tests preserved:

• BoolToUint8_WebGpu — bool → uint8, exercises Uint8x4 pack path.
• Uint8ToInt32_WebGpu — uint8 → int32, values {0, 1, 128, 255, 42, 7} span the full byte range; exercises unpack4xU8 + widening.
• Uint8ToFloat_WebGpu — uint8 → float, same values on the float-conversion arm.
• Uint8ToBool_WebGpu — uint8 → bool, exercises unpack4xU8 + Boolx4 output packing + the "any nonzero → true" cast-to-bool semantic.

All use the 2×3 = 6 non-multiple-of-4 shape to exercise the final partial packed word. Test comments explicitly document the reason for that choice, which helps future maintainers avoid "simplifying" to a 4-aligned shape.

Non-blocking gaps unchanged from previous review:

• No uint8 → int64 test (code path enabled but not pinned; is_from_unsigned grouping means it should behave identically to uint32 → int64 which is exercised elsewhere).
• No dedicated multiple-of-4 shape (the 6-element shape covers one full vec4 + a 2-tail, so the pure-vectorized loop is exercised transitively).

---

Copilot / CI status

• Copilot AI review: 0 comments (Low effort) — unchanged from previous head.
• CI: 85 / 86 checks OK on ee81e5c, up from 77 / 86 on b6024f6. The one still-not-OK leg is likely pending or a known-unrelated Azure Pipelines leg (per the standard ORT pattern where 2 of the 86 checks are Azure-Pipelines-not-run without a /azp run). Nothing in the change surface (shader-variable pack/unpack + Cast plumbing + tests) is a plausible break source for the wider CI matrix.

---

Minor observations (non-blocking)

1. unpack4xU8 WGSL feature footprint. As noted above, requires a Dawn version that landed the WGSL 1.1 data-packing builtins. ORT's pinned Dawn has it. If any downstream consumer swaps Dawn, they'll need to verify. Consider capturing this as a comment in cast.cc near the unpack4xU8 call sites so a future porter to a stripped-down Dawn / non-Dawn WGSL doesn't get a puzzling shader-compile failure — but that's a future-proofing nit, not a blocker.
2. Uint8x4 fallthrough comment is on the default case. Perfectly clear inline, but if Int8x4 ever ships (see the STORAGE_TYPE_ARRAY placeholder in shader_variable.cc), consider bumping the comment out into a case-block that explicitly names both types. Cosmetic.
3. is_from_uint8_ is used at two nearly-identical let a = ... construction sites (int64-output and generic paths). Fine as-is; if a third site ever needs it, factoring load_a into a helper would be cleaner.

---

Bottom line

The re-fetch confirms the PR is in a strictly-better state than the previous approved head. The design shift from auto-unpack in GetByOffsetImpl to caller-side unpack4xU8 (a) aligns the WebGPU EP's Uint8x4 read convention with what Expand (and soon Gather) do, (b) uses a WGSL builtin instead of four manual bit operations, and (c) leaves the pack side and all tests exactly as they were. Correctness is preserved by construction (unpack4xU8 is definitionally equivalent to the removed manual unpack). CI moved from 77/86 to 85/86 OK. Approve — ready to merge once the last check settles.