Skip to content

[WebGPU] Support int8 data in Gather - #33790

Merged
adrastogi merged 1 commit into
mainfrom
adityar-microsoft-webgpu-gather-int8
Nov 4, 2026
Merged

adrastogi merged 1 commit into
mainfrom
adityar-microsoft-webgpu-gather-int8

Conversation

@adrastogi

@adrastogi adrastogi commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Description

Adds int8 data support to the WebGPU Gather kernel.

  • int8 is added to the T type constraint of every WebGPU Gather registration (opsets 1-10, 11-12, 13+).
  • int8 uses the packed-byte path that bool and uint8 already use (four 1-byte elements per u32). Gather only moves data, so int8 is copied as raw bytes. unpack4xU8 reads each byte, and the Int8x4 SetByOffset keeps each lane's low byte. The two's-complement bit pattern, and so the sign, is preserved.
  • bool, uint8 and int64 behavior is unchanged.

Tests in onnxruntime/test/providers/webgpu/tensor/gather_op_test.cc run only on the WebGPU EP with kOrtSessionOptionsDisableCPUEPFallback set, and skip if the EP is unavailable:

  • Gather_int8_axis0_webgpu / Gather_int8_axis1_webgpu: values include -128, -1, 0 and 127. They use int32 and int64 indices, negative indices, 2-D indices, and output sizes that do and don't fill the last packed word.
  • Gather_int8_scalar_indices_webgpu: 0-D indices, giving 0-D and 1-D outputs.
  • Gather_int8_all_opsets_webgpu: opsets 7, 11 and 13, one for each registration.
  • Gather_int8_all_byte_values_webgpu: all 256 int8 values, with 257 outputs that span two workgroups and end on a partial word.
  • Gather_int8_then_cast_to_float_webgpu: Gather(int8) -> Cast(float) graphs with runtime and initializer tables, 2-D and 0-D indices, and two Gathers on one table. This mirrors the WebNN conformance graphs and catches an int8 table that is read as unsigned.

The existing all-EP CPU tests Gather_axis1_indices2d_int8 and Gather_axis1_neg_indices2d_int8 now also run on the WebGPU kernel.

Motivation and Context

The new WPT page webnn/conformance_tests/gather-cast-int8-table.https.any.js fails 6 required subtests on the WebGPU EP. Each failing graph reads an int8 table with gather (constant or runtime table, followed by cast to float32 and sometimes mul/reshape, 0-D indices, or two gathers on one table). The WebGPU Gather kernel didn't register int8:

ep.cc:96 Ep::GetCapabilityImpl webgpu kernel not found in registries for Op type: Gather
This session contains graph nodes that are assigned to the default CPU EP, but fallback to CPU EP has been explicitly disabled

WebNN disables CPU fallback, so these graphs failed to build. The cast-then-gather variants already passed.

This is one of several WebGPU EP fixes for WebNN conformance.

Validation

  • Built locally with the WebGPU EP. This commit was merged with main 121e34d92b together with a separate Relu int32 change.
  • onnxruntime_provider_test --gtest_filter="WebGpuActivationTest.*:ActivationOpTest.Relu*:GatherOpTest.*:CastOpTest.*:GatherElementsOpTest.*:GatherNDOpTest.*": passed, 268/268, 0 skipped. This includes all six new Gather_int8_*_webgpu tests and the existing Gather_axis1_*_int8 tests, which now run on WebGPU.
  • Negative check: removing the type_constraints.push_back(DataTypeImpl::GetTensorType<int8_t>()); line makes all six Gather_int8_* tests fail at session init with fallback to CPU EP has been explicitly disabled. This is the same error the WPT page hits.
  • lintrunner -a and clang-format --dry-run -Werror (20.1.8) on the changed files: clean.
  • WPT in the browser. Setup: plugin EP built from main 121e34d92b, this PR and the Relu int32 PR (local merge 5c68bbb9f9); Edge Canary 159.0.4372.0; RTX Spark N1X / D3D12; ORT runtime 1.28.3.
    • webnn/conformance_tests/gather-cast-int8-table.https.any.js?gpu: PASS 9/9. Before this change, 6 required subtests failed with "Failed to build graph".
    • webnn/conformance_tests/gather.https.any.js?gpu: still PASS 39/39.

Register int8 for the WebGPU Gather T constraint (opsets 1-10, 11-12, 13+)
and route it through the existing packed-byte (4 per u32) path shared with
bool/uint8. Bytes are copied raw: unpack4xU8 reads each byte and the Int8x4
SetByOffset keeps each lane's low byte, so negative values keep their sign.

Without this, WebNN graphs that gather from an int8 table fail to build on
the WebGPU EP when CPU fallback is disabled (WPT
webnn/conformance_tests/gather-cast-int8-table.https.any.js).

Add WebGPU-pinned tests with CPU fallback disabled covering -128/127 and
other negatives, partial packed words, axis 0/1, int32/int64 and negative
indices, 0-D indices, every registered opset range, all 256 byte values,
and Gather(int8) -> Cast(float) graphs with runtime and initializer tables.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

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.

🟢 Approval recommended

There are no active review issues for the int8 WebGPU Gather change.

No open issues 🎉

What changed in this PR

Adds int8 support to ONNX Runtime’s WebGPU Gather kernel, allowing WebNN gather-and-cast graphs to run without CPU fallback.

Changes:

  • Registers int8 across all supported Gather opsets.
  • Reuses the packed-byte path to preserve int8 values without changing bool, uint8, or int64 behavior.
  • Adds WebGPU-only tests for signed values, scalar and multidimensional indices, partial packed words, and Gather-to-Cast graphs.
File Description
onnxruntime/​test/​providers/​webgpu/​tensor/​gather_op_test.cc Adds int8 Gather coverage with CPU fallback disabled.
onnxruntime/​core/​providers/​webgpu/​tensor/​gather.cc Enables int8 registration and packed-byte data movement.

Review effort: Balanced

@adrastogi
adrastogi marked this pull request as ready for review November 2, 2026 23:19
@adrastogi
adrastogi merged commit dd1a165 into main Nov 4, 2026
@adrastogi
adrastogi deleted the adityar-microsoft-webgpu-gather-int8 branch November 4, 2026 21:02
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