WebGPU EP: Add uint8 support to Gather kernel - #29861
Conversation
- Add SupportedNumberBoolAndUint8Types and WebGpuSupportedNumberBoolAndUint8Types() to webgpu_supported_types.h, scoped to uint8 only for Gather - Generalize is_bool to pack_as_bytes in gather.cc to cover both bool and uint8 using the existing 4-elements-per-u32 packing path - Update Gather kernel registrations to use WebGpuSupportedNumberBoolAndUint8Types() - Add Gather_axis1_indices2d_uint8 unit test
…d alias Instead of adding SupportedNumberBoolAndUint8Types / WebGpuSupportedNumberBoolAndUint8Types() to the shared header (a one-off type list only used by Gather), inline the type list directly in the three Gather macro calls, matching the existing Tind inline pattern.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
@hariharans29, PTAL, thanks! |
hariharans29
left a comment
There was a problem hiding this comment.
Review: PR #29861 — WebGPU EP: Add uint8 support to Gather kernel (head 6a2a965)
Author: @huningxin, 3 commits by copilot-swe-agent: 86989d4 (plan), f192a58 (implementation), 6a2a965 (registration simplification — inline the type list). CI: 3/71 OK at fetch time — still spinning up, the WebGPU test legs have not yet reported. No reviewers assigned; @hariharans29 and @adrastogi pinged. References issue #29756.
Verdict: request changes. Motivation and registration wiring are fine (WebNN maps ONNX bool → WebNN uint8, so a Gather-on-bool in models like detr-resnet-50 lowers to WebGPU Gather on uint8, and adding uint8 to the type constraint is the right fix), but the shader code as-written is invalid WGSL for uint8 — it reuses the vec4<bool> component-indexing pattern in a place where the value/element type for Uint8x4 is scalar u32. This should be caught by the WebGPU test legs of CI once they run.
What it does
- Shader (onnxruntime/core/providers/webgpu/tensor/gather.cc) — renames is_bool to
pack_as_bytesand generalizes the "one-of-four" check so both ProgramVariableDataType::Boolx4 and ProgramVariableDataType::Uint8x4 follow the loop-of-4 packed path. - ComputeInternal — sets
pack_as_bytes = (input_dt == bool) || (input_dt == uint8)and gates thecomponents = 4/ data_size = (data_size + 3) / 4 reshape onpack_as_bytes. The AddInputs / AddOutput calls now pass(pack_as_bytes ? 4 : 1)as the component count for both bool and uint8. - Kernel registration — for opset ranges
[1, 10],[11, 12], and[13+], replaces TypeConstraint("T", WebGpuSupportedNumberAndBoolTypes()) with an explicit inline TypeList<float, MLFloat16, int32_t, uint32_t, bool, uint8_t>. Correct approach: doesn't leak the uint8 constraint into every other kernel that uses WebGpuSupportedNumberAndBoolTypes(). - Test — one new TEST(GatherOpTest, Gather_axis1_indices2d_uint8) in onnxruntime/test/providers/cpu/tensor/gather_op_test.cc, excluding TensorRT/OpenVINO.
🔴 Blocking: the shader is invalid WGSL for Uint8x4
Look at the code inside the loop:
if (pack_as_bytes) {
shader.MainFunctionBody() << " value[" << comp << "] = "
<< data.GetByOffset("data_offset / 4")
<< "[data_offset % 4];\n";
}This works for Boolx4 because of two Boolx4-specific specializations that are visible in onnxruntime/core/providers/webgpu/shader_variable.cc:
- VALUE_TYPE_ARRAY[Boolx4] = "vec4", so
output_value_t(used to declarevar value : output_value_t) isvec4<bool>, and value[comp] = bool_expr is legal indexed assignment. - GetByOffsetImpl has a Boolx4 special case that unpacks the storage
u32into avec4<bool>via bit-mask casts:So data.GetByOffset(...) for Boolx4 is acase ProgramVariableDataType::Boolx4: ss << "vec4<bool>(bool(name[off] & 0xFFu), bool(name[off] & 0xFF00u), ...)";
vec4<bool>expression, and[data_offset % 4]is legal component access. - SetByOffsetImpl similarly has a Boolx4 special case that packs the four bools back into the storage
u32via a dot product.
None of these specializations exist for Uint8x4. For Uint8x4:
- VALUE_TYPE_ARRAY[Uint8x4] = "u32" (line 61) →
output_value_tis scalaru32, so the emittedvar value : u32;combined with value[comp] = ...; is invalid WGSL — WGSL's[]component access is defined only for vectors, matrices, and arrays, not for scalars. - GetByOffsetImpl falls through to default: ss << name << "[" << offset << "]";. data[data_offset / 4] is scalar
u32, andu32[data_offset % 4]is also invalid WGSL. - SetByOffsetImpl also falls through to default name[offset] = value;, which writes an unpacked u32 (whose bit pattern is undefined here because value[comp] = ... never validly executed).
The correct pattern for Uint8x4 — established just a few days ago in PR #29854 for the WebGPU EP Expand op — is bit-shift + mask on the read side and OR-shift assembly on the write side:
// Extract byte k from packed u32:
let byte_k = (data[i] >> (k * 8u)) & 0xFFu;
// Assemble 4 bytes into packed u32:
let packed = e0 | (e1 << 8u) | (e2 << 16u) | (e3 << 24u);Concretely, the fix for pack_as_bytes needs to branch on is_bool vs is_uint8:
if (is_bool) {
// Existing code:
// var value : vec4<bool>;
// value[comp] = data.GetByOffset(offset/4)[offset % 4];
// output.SetByOffset(global_idx, value); // uses Boolx4 pack in SetByOffsetImpl
} else if (is_uint8) {
// Assemble a packed u32 across the 4 iterations, e.g. accumulate:
// var value : u32 = 0u;
// value |= ((data[offset/4] >> ((offset % 4) * 8u)) & 0xFFu) << (comp * 8u);
// and at the end:
// output[global_idx] = value;
}Or, if the ORT WebGPU shader-machinery is going to grow more Uint8x4 consumers (this PR + Expand + potential future ops), it may be cleaner to add matching Uint8x4 specializations to GetByOffsetImpl / SetByOffsetImpl in onnxruntime/core/providers/webgpu/shader_variable.cc (returning a vec4<u32> on the read side and packing back with the OR-shift on write) so that the Gather codegen can stay a single if (pack_as_bytes) { value[comp] = data.GetByOffset(...)[...]; } block. That's a wider refactor, though — an in-op branch would suffice here.
CI check: the WebGPU test legs haven't reported yet (3/71 OK is very early). Once they do, the shader compilation is expected to fail on the new Gather_axis1_indices2d_uint8 test. If for some reason it does pass, there's likely a Dawn/naga leniency that shouldn't be relied on — the WGSL spec (https://www.w3.org/TR/WGSL/#component-access) is explicit that component access is a vector/matrix operation, not a scalar operation.
Even hypothetically, if invalid indexing produced some result on a specific driver, the value semantics wouldn't work either: the Boolx4 SetByOffsetImpl packs values assuming each byte is 0 or 1 (dot(vec4(0x1, 0x100, 0x10000, 0x1000000), vec4(value))). For uint8 you need full-byte packing, not 1-bit packing. The Uint8x4 default SetByOffsetImpl (plain name[offset] = value) does no packing at all, so it just writes whatever bits happened to end up in the value u32.
Test coverage gap (secondary)
The one added test uses default test.Run(...) which tries every enabled EP. That means:
- On WebGPU-disabled builds, it silently runs on CPU (which has uint8 Gather natively) and passes.
- On WebGPU-enabled builds, it would try the WebGPU kernel — this is what the shader bug would break.
Two suggestions worth applying either way:
- Positively assert the WebGPU kernel runs. Follow the pattern established in PR #29847 (Tile int32/uint32) — either test.ConfigExcludeEps({kCpuExecutionProvider}) or a WebGPU-only single-EP config via test.ConfigEp(std::move(webgpu_provider)).RunWithConfig(). Otherwise a future WebGPU regression would still silently pass on CPU fallback and the test wouldn't catch it.
- Add value diversity across the packed 4-byte group. The current test uses
{10, 20, 30, 40, 50, 60, 70, 80, 90}— nine distinct bytes — which is already good coverage; each output row of the 3×2×2 result exercises 4-byte packed groups that span distinct source bytes. Fine as-is.
A byte-diverse test like this, once the shader is correct, would also catch endianness / shift-amount bugs in the fix.
Registration change looks fine
Switching from WebGpuSupportedNumberAndBoolTypes() (a shared helper in onnxruntime/core/providers/webgpu/webgpu_supported_types.h) to the inline TypeList<float, MLFloat16, int32_t, uint32_t, bool, uint8_t> for all three opset ranges is a defensible choice: it keeps the uint8 broadening scoped to Gather rather than accidentally broadening every other kernel that uses the shared helper. Compared to the PR #29854 (Expand) approach of push_back(uint8_t) on a copied GetOpTypeConstraints(...) output, the inline approach here is more explicit. Both patterns are valid; consistency across the codebase is a longer-term hygiene item, not a blocker for this PR.
Nits (address alongside the fix)
-
Cache hint doesn't include the packed-type discriminator. program.CacheHint(std::to_string(axis)) is unchanged. This is OK today because different var_type values produce different program identities on their own — but if you go with the in-op
if (is_bool) / else if (is_uint8)branch, the two shaders differ semantically and the automatic program-identity split will handle it. No action needed if you follow that route. -
Comment update lags the code. The existing comment
// Shader will pack four bools into one uint, so we consider the types of input and output as vec4<bool>.was renamed to// Shader will pack four bytes into one uint, so we consider the types of input and output as vec4.This is inaccurate for uint8 — the output is not avec4, it's a scalaru32(there is novec4<u32>value type for Uint8x4). Once you fix the shader, please make the comment match reality. -
data.GetByOffset(...) string-fetched twice per iteration. Both the read and the potential fix would benefit from a locally computed
let packed = data[data_offset / 4];and then extract-per-comp. Minor; codegen readability.
Bottom line
- Registration wiring: ✅ correct.
- Motivation: ✅ closes another slice of the WebNN-demo umbrella.
- Shader for bool: ✅ unchanged, still correct.
- Shader for uint8: ❌ invalid WGSL — needs bit-shift/mask extraction and OR-shift assembly on writes, mirroring the pattern PR #29854 (Expand) established.
- Test: passes on CPU, should fail on WebGPU CI legs once they run. Please also switch the test to a WebGPU-only run to positively assert kernel selection.
Please rework the packed-uint8 shader path (either in-op branching or by adding Uint8x4 specializations to GetByOffsetImpl / SetByOffsetImpl), then re-request review.
…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.
|
Addressed review comments in 1324e91, please take another look, thanks! |
The uint8 path accumulates output bytes via `value = value | (...)` but `value` was declared without an initializer. For partial (last) threads where output_size is not a multiple of 4, un-guarded byte positions kept undefined contents, and the OR-accumulator relied on uninitialized memory in general. This produced incorrect results and failed GatherOpTest.Gather_axis1_indices2d_uint8 and GatherOpTest.Gather_axis1_indices2d_uint8_webgpu. Zero-initialize `value` for the uint8 packed-byte path so the accumulator is well-defined.
Head branch was pushed to by a user without write access
|
Fixed the CI failure. |
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
…-tests WebGPU Gather: align uint8 packed-byte handling with Uint8x4 helpers
Head branch was pushed to by a user without write access
|
After the rework, this PR follows the same packed-byte conventions as the merged #29854, @hariharans29, please take another look. |
…d non-multiple-of-4 test
Co-authored-by: huningxin <1005673+huningxin@users.noreply.github.com>
|
Still work on the CI failures. |
Co-authored-by: huningxin <1005673+huningxin@users.noreply.github.com>
…-support WebGPU Gather: fix uint8 packed-byte kernel producing wrong output
|
Push a commit to fix the ci failures. Please help trigger the CI build. Thanks! |
|
There are still some failing checks. I also have a minor suggestion - If you have more PRs that are on the similar theme (i.e.) uint8 type support and int64 type support for ops, can you please combine a few op changes into a single PR - maybe we can cap it at 5 ops per PR ? I think this is a good balance point between PR size/change and PR content density. It will lower the E2E time it takes to get all the 5 changes into main and it will not bottleneck CI resources and reduce reviewer distraction. Please let me know if there is any concern on this approach. Thanks. |
|
@hariharans29, thanks for your feedback. I am fixing the ci build failures.
Sounds good. We'll follow that. |
Thank you. Much appreciated. Meanwhile, I hope to tackle most of the open PRs later today. |
|
345b7b6 should fix the previous ci test failures. I verified it locally on Windows and remotely on Mac with a job in my own fork https://github.com/huningxin/onnxruntime/actions/runs/30493889593/job/90718101966 Please take another look, thanks for your patient review. @hariharans29 |
1ffca78 to
345b7b6
Compare
|
/azp run Test Linux CUDA x64 Release, Test Linux TensorRT x64 Release, web_Debug / build_onnxruntime_web, web_Debug / build_onnxruntime_web |
|
No pipelines are associated with this pull request. |
|
I am surprised that the CI build keep failing, e.g. MacOS CI Pipeline / webgpu / build-and-test (arm64, arm64, Release) (pull_request) But the same pipeline on my own fork is fine Did I miss anything? |
The root cause is an interaction with #29839 ("Support Cast to and from uint8"), which landed after this PR was written. That PR added a |
Gather on uint8 fell back to the CPU EP because uint8 was not in the T type constraint. The WebNN EP converts ONNX bool to WebNN uint8, so with WebGPU as the WebNN backend, Gather on a bool tensor becomes Gather on uint8 - which models such as detr-resnet-50 need end to end. uint8 is stored packed as Uint8x4 (four elements per u32, lane 0 -> low byte), the same layout bool already uses, so the kernel handles four output elements per thread. Each element's index math runs in its own guarded block, since the last word of a non-multiple-of-4 output is only partially populated. The accumulator holds one element per lane and is passed to SetByOffset unpacked: since microsoft#29839 SetByOffset has a Uint8x4 case that performs the byte packing itself. Pre-packing into a u32 and passing that scalar is wrong - SetByOffset would splat it across all four lanes, mask each to its low byte and re-pack, so every output word would come back as element 0's byte repeated four times: expected: [20,10,30,20 | 50,40,60,50 | 80,70,90,80] got: [20,20,20,20 | 50,50,50,50 | 80,80,80,80] uint8 needs an explicit vec4<u32> accumulator because output_value_t for Uint8x4 is the packed u32 storage type, unlike Boolx4's vec4<bool>. Note the kernel is therefore only correct in a tree containing microsoft#29839, so the branch must be built merged with main rather than standalone. Tests: Gather_axis1_indices2d_uint8 exercises the CPU EP path and, in a WebGPU build, the WebGPU EP as well. Gather_axis1_indices2d_uint8_webgpu pins the packed-byte shader explicitly, and Gather_axis0_uint8_non_multiple_of_4_webgpu covers an output size that is not a multiple of 4, where the final word holds one valid byte. Verified with main merged in on Intel Arc A770 and NVIDIA RTX A2000 over D3D12 and on Vulkan: all three tests fail before this change and pass after it on every configuration, and the full GatherOpTest suite (22 tests, covering the bool, float, int8, string and scalar-indices paths) passes.
345b7b6 to
554a7b2
Compare
|
@hariharans29, thanks much for your patient review and insightful guidance! Very helpful. |
Description
The WebNN EP maps ONNX BOOL → WebNN uint8. When WebNN uses WebGPU as its backend, a Gather on a bool tensor (e.g., in detr-resnet-50) is lowered to WebGPU Gather on uint8, which previously had no matching kernel.
Motivation and Context
Models such as detr-resnet-50 use Gather on bool tensors. The WebNN EP converts bool to uint8 before dispatching to the WebGPU backend, so the WebGPU Gather kernel must accept uint8 to handle this path end-to-end.