WebGPU: enable INT64 for Equal/Sub/Where/ReduceSum under enable_int64 flag to support Gemma4

[feich-ms]

Summary

These four ops were forcing CPU fallback when inputs were INT64, preventing WebGPU graph capture. Converted from static macro registration to the factory function pattern (same as Cast/Unsqueeze/Expand/Range) so INT64 type constraints are included when enable_int64=true.

enable_int64 is always true when enable_graph_capture=true — GetKernelRegistry calls RegisterKernels(true, true) unconditionally in that path — so this fix has no effect on the default (non-graph-capture) execution path.

Changed ops

Op	Versions
Equal	7–10, 11–12, 13–18, 19+
Sub	7–12, 13, 14+
Where	9–15, 16+
ReduceSum	1–10, 11–12, 13+

Approach

Each op follows the established factory function pattern:

• CreateXVersionedKernelInfo<Start, End>(bool enable_int64) / CreateXKernelInfo<Since>(bool enable_int64)
• Removed from static build_kernel_create_info_function_table[]
• Registered in RegisterKernels() alongside Cast/Unsqueeze/Expand/Range

ReduceSum version 13+ retains InputMemoryType(OrtMemTypeCPUInput, 1) (axes input must be on CPU) — unchanged from the original registration.

Where factory functions delegate to GetOpTypeConstraints(enable_int64, /*enable_bool=*/true) rather than a local duplicate type list. All three files (where.cc, binary_elementwise_ops.cc, reduction_ops.cc) use stack-allocated KernelDefBuilder() consistent with other WebGPU factory functions, and drop the unused webgpu_execution_provider.h include.

Runtime fixes for INT64 inference

BinaryElementwise (Sub, Equal) and Where required additional shader-level fixes to correctly handle INT64 tensors at inference time:

• Vectorization disabled for INT64: INT64 is stored as vec2<u32> (8 bytes/element) and has no vec4 representation. Vectorization is now disabled when either input is INT64, and component=1 is used for correct buffer binding.
• Correct dispatch size: vec_size is computed as element count (not rounded-up vec4 count) for INT64 outputs.
• Scalar input path: The scalar broadcast path no longer applies an extra .x dereference on top of GetByOffset, which already extracts the i32 element for INT64.
• Where non-broadcast path: Non-broadcast INT64 Where now uses the element-per-thread shader path instead of the vec4 fast path that would truncate values to i32.
• Cache key: is_int64 is included in the Where cache hint, preventing float and INT64 shaders from incorrectly sharing a cached entry.
• Shape indices for INT64: Shape indices are registered for non-broadcast INT64 Where so BroadcastedIndicesToOffset has the required helpers available.
• Equal INT64 with bool output: The element-wise path now explicitly reads 4 consecutive INT64 elements per thread and packs them into a vec4<bool> before writing to the bool output buffer. Previously, only element 0 was read and broadcast across all 4 positions of the comparison, giving wrong results.

Testing

• All 1174 nodes of a Gemma4 WebGPU decoder model assigned to WebGpuExecutionProvider (previously ~20 nodes fell back to CPU due to INT64 Equal/Sub/Where/ReduceSum)
• End-to-end graph capture inference: ~87 tok/s on Gemma4-E2B-it INT4
• USE_WEBGPU-guarded unit tests added for all four ops with kEnableInt64=1: Sub_int64_webgpu, Equal_int64_webgpu, Where_int64_webgpu, ReduceSum_int64_webgpu

Related

• Indirect dispatch fix (prerequisite for graph capture): WebGPU: Add indirect dispatch for flash attention graph capture to support Gemma4 #29236
• https://github.com/onnxruntime/mobius/pull/380

[Copilot]

Pull request overview

This PR updates the WebGPU EP kernel registration for Equal, Sub, Where, and ReduceSum so that INT64 type constraints are included when enable_int64=true, avoiding unintended CPU fallback (notably in graph-capture scenarios).

Changes:

• Converted Equal/Sub/Where/ReduceSum from static macro/table-based registration to factory-function registration, enabling conditional INT64 type constraints.
• Wired the new factory registrations into RegisterKernels(enable_graph_capture, enable_int64) alongside existing conditional-int64 kernels (e.g., Cast/Unsqueeze/Expand/Range).
• Preserved existing special-casing (e.g., ReduceSum v13+ axes input remains OrtMemTypeCPUInput).

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Removes static table entries for the 4 ops and registers them via RegisterKernels() with enable_int64 propagation.
onnxruntime/core/providers/webgpu/tensor/where.h	Declares Where kernel factory functions to support conditional INT64 registration.
onnxruntime/core/providers/webgpu/tensor/where.cc	Implements Where kernel factory functions and conditional type constraints.
onnxruntime/core/providers/webgpu/reduction/reduction_ops.h	Declares ReduceSum kernel factory functions to support conditional INT64 registration.
onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc	Implements ReduceSum kernel factory functions with conditional INT64 type constraints and preserves CPU axes input for v13+.
onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.h	Declares Equal/Sub kernel factory functions to support conditional INT64 registration.
onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc	Implements Equal/Sub kernel factory functions with conditional INT64 type constraints.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

[hariharans29]

Review: PR #29392 — WebGPU: enable INT64 for Equal/Sub/Where/ReduceSum under enable_int64 flag to support Gemma4

LGTM overall. The enable_int64 gating is scoped correctly (only reachable via graph capture, so no impact on the default execution path), the factory-function registration mirrors the Cast/Unsqueeze/Expand/Range pattern already in the EP, and the shader-level INT64 fixes are all necessary — the Equal-INT64→bool packing fix in particular corrects a real silent-wrong-answer that existed before this PR. A few non-blocking items below.

What's good

1. The factory-function refactor is the right shape. Matches Cast/Unsqueeze/Expand/Range exactly, including the explicit template instantiations at the bottom of each .cc and the forward declarations in the headers. The KERNEL_CREATE_INFO entries for these four ops are removed from the static table with explanatory comments, and the new registrations land in RegisterKernels(). No double-registration risk.
2. ReduceSum v13+ retains InputMemoryType(OrtMemTypeCPUInput, 1) — the axes input still stays on CPU. Easy to miss in a refactor, handled correctly.
3. The Equal INT64→bool fix is a real correctness fix, not just a new code path. The old scalar path (input_a_value_t(GetByOffset("0").x) applied component-wise to the vec4 comparison) was silently comparing only element 0 and broadcasting the result across all 4 output positions. The new code reads all 4 elements per thread and packs them properly, with element_count-guarded reads for lanes 1–3.
4. Vectorization gating is comprehensive: vectorize = !is_int64 && (...) on the outer flag, and the "consider enabling vectorize via divisible-by-4 check" branch is also !is_int64 gated, so INT64 tensors always take the element-per-thread path.
5. Cache-hint fix in Where (.CacheHint(is_broadcast, is_int64)) prevents INT64 and float shaders from being confused for one another in the shader cache. That would have been a nasty runtime bug otherwise.
6. Test coverage exists for all 4 ops via USE_WEBGPU-guarded unit tests that flip on kEnableInt64=1 — small tensors, known outputs, exercising the actual WebGPU path.
7. CI is green (86 / 86) on the latest commit after multiple iterations.

Non-blocking items

1. Missing #include <memory> — the Optional Lint C++ leg is flagging cpplint warnings in three files:

   • onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc:461
   • onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc:86
   • onnxruntime/core/providers/webgpu/tensor/where.cc:202

   All three: Add #include <memory> for make_unique<> [build/include_what_you_use]. Transitively included today so builds pass, but reviewdog is right — add explicit #include <memory> at the top of each .cc.

2. The i32-range limitation for INT64 deserves more than a code comment. The PR comments make this clear:

   "int64 arithmetic in the WebGPU shader operates on the low 32 bits only (i32 element type). Values outside the int32 range [-2^31, 2^31-1] will produce incorrect results."

   That's an out-of-spec silent-wrong-answer for values ≥ 2³¹ on Sub/Equal/Where/ReduceSum, only mitigated by the fact that enable_int64 currently only activates under graph capture and Gemma4-style workloads use INT64 tensors only for token positions / attention masks that fit in i32. Suggest documenting it near the kEnableInt64 config key so users toggling this flag directly (not via graph capture) know, or logging a one-time warning at kernel-selection time when enable_int64=true.

3. Style consistency nit in WhereProgram::GenerateShaderCode. The int64 branch uses raw storage-buffer accessors (c_data[...], a_data[...], output_data[global_idx] = ...) while the surrounding non-int64 code uses output.SetByOffset("global_idx", ...) / X.GetByOffset(...). Both work; the raw-access style is legitimate for the packed-bool byte extraction (c_data[offset_c / 4u] & (0xffu << (u32(offset_c % 4u) * 8u))), but a_data[offset_a] / output_data[global_idx] = ... could stay in the standard accessor form for uniformity.

4. Sanity-check the Equal INT64→bool dispatch bound-checking. The Equal INT64→bool path uses is_output_int64=false → vec_size = (size+3)/4 and reads inputs one element at a time with base = global_idx * 4u. Lanes 1–3 are guarded against base + K < n; lane 0 (base itself) is not. If the workgroup rounding dispatches threads with global_idx >= vec_size, base could exceed n. Under robust buffer access an OOB read returns 0 and 0 == 0 → true would be written via SetByOffset("global_idx", ...), which should silently drop OOB stores — so this is probably fine, but it's the same class of bug that Copilot-swe-agent already caught in commit 563d5f5. Worth confirming that either SetByOffset bounds-checks global_idx or that the shader has an explicit early-return if (global_idx >= uniforms.vec_size) { return; } at the top of this branch. If not, a leading if (base >= uniforms.element_count) { return; } is cheap insurance.

5. Prefer a WebGPU-EP-owner sign-off before merge. qjia7 has been requested but hasn't reviewed here yet. This PR touches more of the EP than WebGPU: Add indirect dispatch for flash attention graph capture to support Gemma4 #29236 did (kernel registration, shader packing math, cache-hint keys), so I'd wait for qjia7 to look before merging even after we approve.