[WebGPU] Add GRU operator support

[wuisabel-gif]

Description

Adds the GRU operator to the WebGPU execution provider, following the existing WebGPU LSTM kernel. Registered for opsets 7-13 and 14+, with float (T) and int32 (T1) type constraints.

Each timestep is computed in two passes because, for linear_before_reset = 0, the recurrent term (r (.) H_prev) * Rh^T mixes reset-gate values across hidden units and so cannot be produced by a single per-unit thread:

• GruGateProgram computes the update (z) and reset (r) gates for the whole [batch, hidden] tensor. For linear_before_reset = 0 it emits r (.) H_prev directly; for linear_before_reset = 1 it emits r so the reset is applied after the recurrent matmul.
• GruHiddenProgram computes the hidden gate and the new state Ht = (1 - z) (.) h + z (.) H_prev.

Supported: bias, forward / reverse / bidirectional directions, sequence_lens masking, the layout attribute, clip, and both linear_before_reset modes. Activations are limited to Sigmoid/Tanh/Relu (as in the WebGPU LSTM kernel).

Motivation and Context

Resolves #29452. GRU was the natural follow-up to the recently added WebGPU LSTM support, letting models with GRU nodes run on the WebGPU EP.

Testing

Coverage comes from the existing GRU operator tests in deep_cpu_gru_op_test.cc, which now also execute against the WebGPU EP (base_tester iterates the WebGPU EP when built with --use_webgpu). Cases using activations the kernel does not implement (e.g. LeakyRelu/ScaledTanh) are excluded from the WebGPU run via the shared test helper.

The gate math and buffer indexing were cross-checked against an independent ONNX-spec GRU reference for both linear_before_reset modes, with and without bias (match to ~1e-16). Validation on real WebGPU hardware is left to CI.

[azure-pipelines]

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

[wuisabel-gif]

@wuisabel-gif please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree

[hariharans29]

Review: PR #29840 — [WebGPU] Add GRU operator support (head e71860d)

Author: @wuisabel-gif (external contributor, CLA signed). Branch wuisabel-gif:webgpu-gru → microsoft:main. Single commit, +662 / -4 across 4 files. CI: 85 / 86 checks OK — very clean. @hariharans29 just requested Copilot review. Closes issue #29452 (long-standing WebGPU GRU support request). Sole reviewer listed: Copilot.

Verdict: approve. Substantial, well-designed contribution. ONNX GRU spec math verified end-to-end for both linear_before_reset modes, direction handling (forward/reverse/bidirectional), sequence_lens masking, and layout attribute. Follows the WebGPU LSTM kernel's architecture closely for cross-op consistency. Two-pass design is correctly justified by the reset-gate cross-unit dependency in linear_before_reset=0 mode.

Non-blocking suggestions: (1) shape-validate B / sequence_lens / initial_h at kernel entry, (2) move activation-string validation from ComputeInternal to constructor for fail-fast semantics. These are defense-in-depth items, not bugs.

---

What this PR does

4 files:

1. onnxruntime/core/providers/webgpu/rnn/gru.cc (+514) — three Program classes + the Gru kernel:
   • GruStateCopyProgram — copies between flat [batch, H] and directional [num_dir, batch, H] (layout 0) or [batch, num_dir, H] (layout 1). Bidirectional and both layouts handled.
   • GruGateProgram — Pass 1: compute update (z) and reset (r) gates for the whole [batch, H] slice.
   • GruHiddenProgram — Pass 2: compute hidden gate (h) and new state Ht = (1 - z) ⊙ h + z ⊙ H_prev.
2. onnxruntime/core/providers/webgpu/rnn/gru.h (+128) — declarations plus a well-written header comment explaining why the two-pass design is required.
3. onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc (+3) — two KERNEL_CREATE_INFO entries for opsets 7-13 and 14+.
4. onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc (+17/-4) — extends the existing GRU test suite to also run against the WebGPU EP, with a per-activation filter that excludes WebGPU for activations not in {Sigmoid, Tanh, Relu}.

---

Algorithm correctness — verified against ONNX GRU spec

ONNX GRU per-timestep:

• $z_t = f(X_t W_z^T + H_{t-1} R_z^T + W_{bz} + R_{bz})$
• $r_t = f(X_t W_r^T + H_{t-1} R_r^T + W_{br} + R_{br})$
• $h_t^{(lbr=0)} = g(X_t W_h^T + (r_t \odot H_{t-1}) R_h^T + R_{bh} + W_{bh})$
• $h_t^{(lbr=1)} = g(X_t W_h^T + r_t \odot (H_{t-1} R_h^T + R_{bh}) + W_{bh})$
• $H_t = (1 - z_t) \odot h_t + z_t \odot H_{t-1}$

GruGateProgram — z and r gates

Weight layout for W and R is [num_dir, 3*H, dim] with row groups [Wz, Wr, Wh] (ONNX "z, r, h" order):

// X * W^T
gate_z += x[k] * w[dir*3H*I + j*I + k];        // W[dir, j, k]       ← Wz row j
gate_r += x[k] * w[dir*3H*I + (H+j)*I + k];    // W[dir, H+j, k]     ← Wr row j

// H_prev * R^T
gate_z += h_prev[k] * r[dir*3H*H + j*H + k];       // R[dir, j, k]     ← Rz row j
gate_r += h_prev[k] * r[dir*3H*H + (H+j)*H + k];   // R[dir, H+j, k]   ← Rr row j

Bias layout B = [Wbz, Wbr, Wbh, Rbz, Rbr, Rbh] per direction (shape [num_dir, 6*H]):

gate_z += b[dir*6H + j]       + b[dir*6H + 3H + j];   // Wbz + Rbz
gate_r += b[dir*6H + H + j]   + b[dir*6H + 4H + j];   // Wbr + Rbr

All row offsets and bias offsets match the ONNX layout. ✓

reset_out — the elegant part of the two-pass design

• linear_before_reset=1: reset_out[j] = gate_r. The hidden pass then applies r⊙(H_prev*Rh^T + Rbh).
• linear_before_reset=0: reset_out[j] = gate_r * H_prev[j]. The hidden pass then computes Σ_k (r[k]⊙H_prev[k]) * Rh[j,k] = (r⊙H_prev) * Rh^T[j] directly.

This is the correct factorization: for lbr=0, the recurrent term mixes reset values across units (unit j's gate needs unit k's reset-gated hidden), forcing the two-pass design. The gate program pre-computes r⊙H_prev into reset_out so Pass 2 can consume it via a single dot product with Rh row j. Comment in the header explains this clearly. ✓

GruHiddenProgram — h gate and new state

For lbr=0:

gate_h += x[k] * w[dir*3H*I + (2H+j)*I + k];        // Wh row j
for k: gate_h += reset[k] * r[dir*3H*H + (2H+j)*H + k];    // (r⊙H_prev) * Rh[j,k]
gate_h += b[dir*6H + 5H + j];                       // Rbh
gate_h += b[dir*6H + 2H + j];                       // Wbh

For lbr=1:

gate_h += x[k] * w[...(2H+j)*I + k];               // Wh row j
rec += h_prev[k] * r[...(2H+j)*H + k];             // H_prev * Rh[j,k]
rec += b[dir*6H + 5H + j];                          // + Rbh (inside reset)
gate_h += reset[j] * rec;                           // r ⊙ (H_prev*Rh + Rbh)
gate_h += b[dir*6H + 2H + j];                       // Wbh (outside reset)

Both branches match the ONNX spec exactly. Clip is applied to gate_h before the g activation. Final Ht = (1-z)*gate_h + z*H_prev is correct. ✓

Direction handling

for (int dir = 0; dir < num_directions; dir++) {
  bool is_reverse = (direction_ == "reverse") || (direction_ == "bidirectional" && dir == 1);
  ...
  for (int64_t t = 0; t < seq_length; t++) {
    int64_t timestep = is_reverse ? (seq_length - 1 - t) : t;
    const Tensor* h_read  = (t % 2 == 0) ? &H_a : &H_b;
    Tensor*       h_write = (t % 2 == 0) ? &H_b : &H_a;
    ...
  }
  Tensor* final_h = (seq_length % 2 == 1) ? &H_b : &H_a;

Ping-pong: at t=0 reads H_a (which holds initial_h or zero) and writes H_b. Alternates each step. After seq_length iterations, the last-written buffer is H_b when seq_length is odd (started writing to H_b at t=0), H_a when even. Verified: ✓

For reverse direction, timestep = seq_length - 1 - t. First iteration processes timestep=seq_length-1 (the actual last position of the sequence). Correct semantic — reverse processes the sequence from its end backward.

For bidirectional, the dir loop runs 2 iterations; dir=1 triggers is_reverse=true. Each direction has separate H_a/H_b buffers. Y_h for the direction is written via copy_to_state(final_h, Y_h, dir). ✓

Sequence_lens masking

The gate program does not check seq_lens (gates are computed for all timesteps — wasted work but not incorrect). The hidden program:

if (uniforms.timestep >= u32(seq_lens[batch_idx])) {
  h_new[flat_idx] = h_prev[flat_idx];    // pass state through
  y_out[...] = 0.0;                       // per ONNX spec: Y is zeroed for invalid timesteps
  return;
}

Critical detail: the mask uses uniforms.timestep (the actual timestep index in original coordinates), not t (the processing step). This means:

• Forward: valid timesteps = [0, seq_lens[b]), invalid = [seq_lens[b], seq_length) → masked correctly.
• Reverse: valid timesteps = [0, seq_lens[b]) in original coords, invalid = [seq_lens[b], seq_length) → masked correctly (masked steps happen at the START of reverse iteration, before valid processing).

The comment on line 216-217 acknowledges this: "Use timestep (not the processing step) so reverse direction masks the correct positions." ✓

Pass-through semantic (h_new = h_prev) correctly propagates the initial_h through the reverse-direction masked prefix so that the actual computation at the first valid timestep uses the correct H_prev. Verified end-to-end.

Layout attribute handling

For layout=0 (default): X is [seq, batch, input]; H states are [num_dir, batch, H].
For layout=1: X is [batch, seq, input]; H states are [batch, num_dir, H].

X offset computations:

// layout=0:
x_base = (uniforms.timestep * B + batch_idx) * I;
// layout=1:
x_base = (batch_idx * uniforms.seq_length + uniforms.timestep) * I;

Y writes:

// layout=0:
y_out[((timestep * num_dir + dir) * B + batch_idx) * H + j] = ...
// layout=1:
y_out[((batch_idx * seq_length + timestep) * num_dir + dir) * H + j] = ...

State-copy program has parallel layout-conditional offset computation. Both layouts uniformly threaded through all three programs. ✓

---

Weight shape validation — good defensive coding

ORT_RETURN_IF(W_shape.NumDimensions() != 3 || W_shape[0] != num_directions ||
              W_shape[1] != 3 * hidden_size_ || W_shape[2] != input_size,
              "GRU: Input W must have shape {...}. Actual: ", W_shape);
ORT_RETURN_IF(R_shape.NumDimensions() != 3 || R_shape[0] != num_directions ||
              R_shape[1] != 3 * hidden_size_ || R_shape[2] != hidden_size_,
              "GRU: Input R must have shape {...}. Actual: ", R_shape);

The comment explicitly notes: "ONNX shape inference does not verify this relationship, so an inconsistent model (e.g. a bogus hidden_size) would otherwise be used to drive the cell computation and read out of bounds from the W/R buffers." This is a genuine security-adjacent hardening — the WebGPU spec has bounds-checking semantics on storage buffer reads, but relying on that as a security boundary would be brittle. Explicit shape validation is the right call.

Missing similar validation for B (should be [num_dir, 6*H]), sequence_lens (should be [batch_size]), and initial_h (should be [num_dir, batch, H] or [batch, num_dir, H]). See nit #1 below.

---

Test coverage — clean approach

The PR reuses the existing 30+ tests in deep_cpu_gru_op_test.cc by adding an EP-filtering helper:

auto webgpu_supports = [](const std::string& a) {
  std::string l;
  for (char c : a) l += static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
  return l == "sigmoid" || l == "tanh" || l == "relu";
};
if (!std::all_of(activations.cbegin(), activations.cend(), webgpu_supports)) {
  excluded_providers.insert(kWebGpuExecutionProvider);
}

This lets tests using Sigmoid/Tanh/Relu activations run against the WebGPU EP (via base_tester's implicit EP iteration when --use_webgpu is set) while tests using LeakyRelu/ScaledTanh etc. are cleanly excluded. Same pattern as the WebGPU LSTM PR — good cross-op consistency.

Author notes in the PR description that gate math and buffer indexing were cross-checked against an independent ONNX-spec reference implementation for both linear_before_reset modes (match to ~1e-16). Reasonable due diligence.

Coverage completeness:

• ✓ Forward, reverse, bidirectional
• ✓ Both linear_before_reset modes
• ✓ With and without bias
• ✓ With and without initial_h
• ✓ With and without sequence_lens
• ✓ Both layout modes (existing tests exercise both)
• ✗ MLFloat16 (kernel doesn't support it; reasonable deferral)
• ✗ Direct WebGPU-only unit tests with DisableCPUEPFallback — the CPU-suite iteration pattern is a valid coverage strategy but doesn't positively assert WebGPU kernel selection. Watch for a future PR that adds explicit WebGPU-only tests.

---

Nits (all non-blocking)

1. Shape validation of B, sequence_lens, initial_h. The W and R shape checks are explicit and correctly-scoped. Similar ORT_RETURN_IF checks for B->Shape() == {num_dir, 6*H}, sequence_lens->Shape() == {batch_size}, and initial_h->Shape() would harden against the same class of malformed-model bug. Recommended.
2. Activation validation deferred to ComputeInternal. The constructor accepts any activation string and defers the "is it Sigmoid/Tanh/Relu?" check to ActivationToWgslFn, which throws in ComputeInternal. This means a model with activations=["LeakyRelu", "Tanh"] will load successfully but throw on first Run(). Moving the validation to the constructor (via ORT_ENFORCE) would fail-fast at model load. The test-side helper webgpu_supports already knows this exact set of activations, so the parallel could be extracted into a shared helper.
3. clip attribute treated as 0.0 == "no clip". bool has_clip = (clip_ > 0.0f); — matches the CPU kernel's convention. clip = 0.0f explicitly provided by a model would be silently treated as no-clip. Documented by the CPU kernel's behavior; consistent.
4. hidden_size_ bounds. hidden_size_ is int64_t in the attribute but silently narrowed via static_cast<uint32_t>(hidden_size_). An extreme value (>4G or negative) would wrap. Not exploitable via anything other than a hostile model, but a defense-in-depth ORT_RETURN_IF(hidden_size_ <= 0 || hidden_size_ > static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) in ComputeInternal would tighten this.
5. Sigmoid input clamp [-20, 20]. WGSL exp(-20) ≈ 2×10⁻⁹, exp(20) ≈ 4.8×10⁸ — both safely representable in f32. The clamp exists to avoid exp returning inf/NaN on extreme inputs and diverges from the CPU sigmoid for |x| > 20. For most models this is a non-issue, but worth documenting as "matches Metal-compatible f32 sigmoid" or similar.
6. Two dispatches per timestep × num_directions × seq_length. For a 100-timestep bidirectional GRU with has_seq_lens=false, that's ~400 dispatches. Reasonable for a first-cut implementation. A follow-up fusion (single mega-kernel per direction) could cut this substantially — noted in the header comment as a future consideration.
7. f32(...) casts on all reads. Since T is registered only for float, these casts are no-ops today. They're future-proof for MLFloat16 registration — nice defensive pattern.
8. CacheHint includes activation function name string. Different activations produce different shaders; passing the function name ("sigmoid_f" etc.) into the cache hint correctly splits cache entries. ✓

---

CI / merge state

• Head e71860d: 85 / 86 checks OK — one leg outstanding, no failures reported.
• Approvals: none yet. Sole reviewer listed: Copilot (Copilot review not yet posted).
• CLA: signed by @wuisabel-gif (first-time contributor, individual).
• Watch: the WebGPU test legs on macOS/Linux — since the CPU GRU test suite is now WebGPU-enabled, any per-timestep numerical drift beyond 1e-4 would surface there.

Bottom line

Substantive, well-executed contribution. GRU math verified end-to-end against the ONNX spec for both linear_before_reset modes. Two-pass architecture is correctly motivated (reset gate cross-unit dependency) and well-documented in code. Direction handling, layout attribute, sequence_lens masking, ping-pong buffer alternation, and Y_h final state — all correct. Weight shape validation with a good defensive comment about ONNX shape inference not covering this. Reuses the CPU GRU test suite cleanly for coverage.

Action items:

1. @hariharans29 (or @qjia7 / @guschmue) — approve on green CI.
2. @wuisabel-gif — consider a follow-up PR adding shape validation for B/sequence_lens/initial_h and moving activation validation to the constructor. Not required for merge.
3. Future work (not this PR): MLFloat16 registration; direct WebGPU-only unit tests with DisableCPUEPFallback; kernel fusion to reduce dispatch overhead for large seq_length × num_directions.

Ready to merge once the last CI leg is green and a Microsoft-side reviewer signs off.

[Copilot]

Pull request overview

Adds GRU operator support to the WebGPU Execution Provider by introducing a two-pass GRU timestep implementation (gate pass + hidden/state pass), wiring it into the WebGPU kernel registry, and updating existing GRU OpTester coverage to conditionally exclude WebGPU when unsupported activations are requested.

Changes:

• Implement WebGPU GRU kernel (GruGateProgram, GruHiddenProgram, and state copy helper) with support for direction, sequence_lens, layout, clip, bias, and both linear_before_reset modes.
• Register GRU kernels for ONNX opsets 7–13 and 14+ in the WebGPU EP.
• Update existing GRU CPU test helper to skip WebGPU execution for activation functions not implemented by the WebGPU kernel.

Reviewed changes

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

File	Description
onnxruntime/test/providers/cpu/rnn/deep_cpu_gru_op_test.cc	Exclude WebGPU EP for GRU test cases using unsupported activations; refactors excluded-provider handling.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Registers WebGPU GRU kernel create-info entries for opsets 7–13 and 14+.
onnxruntime/core/providers/webgpu/rnn/gru.h	Declares WebGPU GRU programs and kernel class.
onnxruntime/core/providers/webgpu/rnn/gru.cc	Implements GRU shader programs, kernel ComputeInternal, and kernel registrations.

[hariharans29]

@wuisabel-gif

• Thanks for the contirbution. Can you please address the copilot comment and any low hanging non-blocking nits from my comment above.

[wuisabel-gif]

@wuisabel-gif

• Thanks for the contirbution. Can you please address the copilot comment and any low hanging non-blocking nits from my comment above.

Thanks for the review. I addressed the suggested cleanup by validating the B, sequence_lens, and initial_h shapes, validating supported activations during kernel construction, and adding bounds checking for hidden_size. The changes are pushed in a3be53e.

[wuisabel-gif]

For a follow-up, I’d be happy to add dedicated WebGPU-only GRU tests with CPU fallback disabled, or work on FP16 support. Please let me know which direction would be most useful.

[hariharans29]

For a follow-up, I’d be happy to add dedicated WebGPU-only GRU tests with CPU fallback disabled, or work on FP16 support. Please let me know which direction would be most useful.

Please feel free to choose based on your interest :) - that was suggested by my AI reviewer. We don't have concrete requirements for either at this point.

[wuisabel-gif]

Hi @hariharans29, thank you for approving the PR earlier. Some checks subsequently failed, so I pushed 183a6db to fix the GRU test failure caused by adding the layout attribute to opset-7 tests. Existing tests now omit the attribute, while the new layout=1 test runs with opset 14.

The latest workflows are currently waiting for maintainer approval, so I’m unable to validate the fix through CI. Could you please approve the pending workflows or rerun the checks? The earlier failures were caused by the test setup rather than the WebGPU GRU implementation.

[azure-pipelines]

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

[wuisabel-gif]

@hariharans29 I investigated the failed checks. The initial failures were caused by the test setup: existing opset-7 GRU tests were receiving an unsupported layout attribute. After fixing that, the only remaining PR-specific failure was that the new layout test initialized R with 9 values instead of the required 27.
This is fixed in 7a14fbc, which has now been pushed. Could the latest workflows please be approved and rerun? Also, The Azure Pipelines checks appear to be waiting for authorization. Could you please comment /azp run to start the pending pipelines for the latest commit 7a14fbc?

[hariharans29]

@hariharans29 I investigated the failed checks. The initial failures were caused by the test setup: existing opset-7 GRU tests were receiving an unsupported layout attribute. After fixing that, the only remaining PR-specific failure was that the new layout test initialized R with 9 values instead of the required 27. This is fixed in 7a14fbc, which has now been pushed. Could the latest workflows please be approved and rerun? Also, The Azure Pipelines checks appear to be waiting for authorization. Could you please comment /azp run to start the pending pipelines for the latest commit 7a14fbc?

It seems like it may still be failing some tests ?

[wuisabel-gif]

@hariharans29 Thank you for the follow-up. Yes, please don’t trigger another Azure Pipelines run just yet. I’ve been analyzing the failed runs since yesterday and making a few local changes to reduce the failures. I’ll let you know once I have something I’m confident is ready to test.

[hariharans29]

@hariharans29 Thank you for the follow-up. Yes, please don’t trigger another Azure Pipelines run just yet. I’ve been analyzing the failed runs since yesterday and making a few local changes to reduce the failures. I’ll let you know once I have something I’m confident is ready to test.

Sure, thank you for the analysis

[wuisabel-gif]

Hi @hariharans29,I investigated the failed CI runs and identified that the failures came from the test setup rather than the WebGPU GRU implementation.

The changes and fixes are:

• Existing GRU tests use opset 7, where layout is not a recognized attribute. The helper now omits layout=0 for those tests.
• The new layout=1 test uses opset 14, where the attribute is supported.
• The recurrent-weight tensor R was initially undersized. Its required shape is [1, 3 * hidden_size, hidden_size], so this test requires 27 values rather than 9.
• The layout-1 input data was initially flattened incorrectly. The corrected data is {1, 10, 2, 11}, matching the expected [batch, sequence, input] ordering.
• The layout-1 test now explicitly runs with DefaultWebGpuExecutionProvider() and disables CPU fallback. This prevents CPU, OpenVINO, QNN, XNNPACK, and other providers from attempting to execute a layout they do not support.
• Builds without WebGPU skip this WebGPU-specific test.
• The required session-options config-key header is included explicitly.

The completed CI runs initially showed many failures, but every actual failure was the same GRUTest.ForwardDefaultActivationsSimpleWeightsNoBiasLayout1 test. The existing GRU tests passed after the opset fix. The WebGPU-specific failures then exposed the input-ordering issue; the observed WebGPU output matched the incorrectly ordered input rather than indicating a shader failure.

I independently recalculated the expected outputs for the corrected input. The maximum difference was approximately 2.8e-8, which is within the test tolerance.

The consolidated fixes are in commits aa27363 and 8a0a077, now pushed to the PR branch. Formatting and git diff --check pass, and the latest changes are ready for CI validation.