[WebGPU] Zero-initialize writable device allocator buffers

[Jiajia Qin (qjia7)]

Description

• Zero-initialize writable WebGPU device-allocator buffers while leaving read-only initializer and internal WebGPU allocations unchanged.
• Clear reused pooled buffers directly in BufferManager::Create in command order, and enable Dawn lazy resource clearing for fresh buffers.
• Track session-run activity so allocator clears are batched during Session::Run, but submitted before allocation returns outside Run; shared allocator clears are always submitted immediately.
• Keep allocation clears out of captured graph commands because graph replay reuses buffers allocated before replay.
• Keep dispatch-window limit enforcement in WebGpuContext and add regressions for allocation submission policy, graph capture, and profiled dispatch batching.

Motivation and Context

Fixes #28974.

Reused device-allocator buffers can retain data from prior model execution. Writable session and shared allocations must therefore be zero-initialized when handed out. Reused buffers are explicitly cleared before use, while newly created buffers rely on Dawn's lazy-clear guarantee.

Clearing and submission are separate policies. During Session::Run, the clear remains ordered in the current command encoder and is submitted with the surrounding dispatch batch. Outside Run, including through the shared allocator exposed to users, the clear is submitted before allocation returns so subsequent queue work observes it.

Allocator clears are allocation-time initialization rather than graph work. They are encoded in normal command order but are not stored in CapturedCommandInfo or repeated by ReplayGraph(), whose inputs and outputs are allocated in advance.

Testing

• Incremental Release build of onnxruntime_provider_test
• WebGpuContextTest.*
• WebGpuDispatchBatchingTests.*
• GraphCaptureTests.TestReleaseCapturedGraph
• InferenceSessionTests.TestGraphCapture

All 16 focused tests pass.

[Jiajia Qin (qjia7)]

cc Wei Wang (@wangw-1991) Wanming Lin (@Honry)

[Copilot]

Pull request overview

This PR hardens the WebGPU EP’s device allocator by ensuring writable, model-visible device buffers are zero-initialized on allocation, preventing reuse of stale GPU memory contents across runs (addressing the info-disclosure risk described in #28974). It also adjusts WebGPU dispatch batching logic so allocator-triggered work draining doesn’t desynchronize profiling/pending-dispatch bookkeeping, and adds targeted regressions.

Changes:

• Add allocation-time zeroing for reused pooled device-allocator buffers via BufferManager::Create(..., initialize_to_zero=true), while leaving read-only/initializer and other internal allocations unchanged.
• Enable Dawn’s lazy clearing for fresh buffers by turning on the lazy_clear_resource_on_first_use device toggle.
• Tighten dispatch-window accounting (encoded + deferred) and add tests covering pooling, graph capture non-capture behavior, and profiled dispatch batching.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/providers/webgpu/webgpu_context_test.cc	Adds tests for pooled-buffer zero init, graph-capture non-capture of allocator clears, and lazy-clear toggle enablement.
onnxruntime/test/providers/graph_capture_test.cc	Adds regression to ensure profiling/dispatch batching remains bounded across allocator clears.
onnxruntime/core/providers/webgpu/webgpu_context.cc	Bounds combined encoded+deferred dispatch count and enables Dawn lazy clearing via device toggles.
onnxruntime/core/providers/webgpu/buffer_manager.h	Extends BufferManager::Create signature with an initialize_to_zero parameter (default false).
onnxruntime/core/providers/webgpu/buffer_manager.cc	Clears reused cached buffers in command order when initialize_to_zero is set, with flush safeguards around pending dispatch limits.
onnxruntime/core/providers/webgpu/allocator.h	Adds allocator state to control allocation-time zero initialization.
onnxruntime/core/providers/webgpu/allocator.cc	Enables zero initialization for writable device allocator buffers and wires the flag into BufferManager::Create.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Xiaofei Han (xiaofeihan1)]

Could we flush the command encoder after recording ClearBuffer() and before returning this buffer?

This path is also reachable through the public WebGPU shared allocator obtained via Env::GetSharedAllocator(). In that case, an ORT user can call Alloc() to create a device tensor and may expect the returned buffer to already be zero-initialized. Currently, ClearBuffer() is only recorded on ORT's command encoder, while Create() returns before the clear is submitted.

As a result, a caller using the returned buffer outside ORT's subsequent command sequence may observe that the buffer has not been cleared yet. Could we call context_.Flush(*this) here to establish the zero-initialization guarantee before returning?

I am not sure if the perf has regresion.

[Jiajia Qin (qjia7)]

Good catch. I think the perf impact is limited since it only impacts the writable device-allocator which usually are allocated in advance before the session.run. So I suppose it won't impact the run performance.

[Xiaofei Han (xiaofeihan1)]

Is this also used for intermediate tensors during model execution? Also, I noticed we're still using defer_size + dispatches_size elsewhere—do we still need that, or can it be removed?

[Jiajia Qin (qjia7)]

Yes, the session device allocator is also used for intermediate tensors. The latest implementation therefore avoids flushing for every allocation during Session::Run: the clear is encoded in command order and submitted by OnRunEnd(). Allocations outside a run, including the public shared allocator, submit the clear before Alloc() returns.

I also audited all GpuBufferAllocator construction sites. Read-only initializer allocators intentionally do not clear, while prepack allocations are submitted at the PrePack() boundary.

The combined num_pending_dispatches_ + deferred_dispatches_.size() check is still needed. These values represent already-encoded and not-yet-encoded dispatches respectively. Checking only the deferred count could exceed the dispatch/profiling window when encoded dispatches remain in the current command encoder. The allocation path itself no longer checks this limit; it flushes only when immediate zero-initialization submission is required.

[Hariharan Seshadri (hariharans29)]

Review — PR #32063: Zero-initialize writable WebGPU device allocator buffers

Summary of changes (10 files, +317/−14)

The core fix: when GpuBufferAllocator::Alloc is called on a writable device allocator and hands out a reused pooled buffer, BufferManager::Create now records command_encoder.ClearBuffer(buffer, 0, size) in command order before returning it. Fresh (uncached) buffers get zeroed by Dawn via the newly-enabled lazy_clear_resource_on_first_use toggle (moved from disabled → enabled in WebGpuContext). Read-only initializer allocators are untouched (initialize_to_zero_ = !is_read_only_allocator).

Submission policy is a separate axis, driven by a new should_submit_zero_initialize callback passed into CreateWebGpuAllocator:

• Inside Session::Run (WebGpuExecutionProvider tracks run_active_ via OnRunStart/OnRunEnd) → defer the clear into the current batch, so dispatch batching is preserved. OnRunEnd flushes.
• Outside Run (shared allocator, prepack, first alloc, etc.) → flush immediately via Flush(*this) before Create returns, so subsequent queue work — including work not issued by ORT — observes the zeroed buffer. This is the fix for the Env::GetSharedAllocator() case Xiaofei Han (@xiaofeihan1) correctly flagged.

Two additional hardening pieces travel with the fix:

1. Dispatch-window accounting in WebGpuContext::Run is tightened from deferred_dispatches_.size() >= max_num_pending_dispatches_ to num_pending_dispatches_ + deferred_dispatches_.size() >= max_num_pending_dispatches_. The old check ignored already-encoded-but-not-yet-submitted dispatches — with clears now being inserted mid-batch, that gap could let the encoded count exceed maxNumPendingDispatches (breaking profiler bookkeeping / validation). A matching ORT_ENFORCE was added to EncodeDeferredDispatches as a belt-and-suspenders check.
2. Graph-capture path stays clean: CapturedCommandInfo only tracks dispatches, not raw ClearBuffer encoder ops, so allocation-time clears are naturally excluded from replay. The description's rationale ("graph replay reuses buffers allocated before replay") is right — a captured clear replayed against a live user tensor would corrupt it.

Approach is sound and the security angle is real

Issue #28974 is a legitimate GPU-memory-disclosure bug: a pooled buffer that held model activations gets recycled and shipped to a subsequent tensor whose consumer reads before writing. WebGPU's spec-level guarantee (fresh buffers zeroed on first use) doesn't help pooled buffers, because from Dawn's perspective they're not fresh. Explicit ClearBuffer on reuse plus the Dawn toggle for fresh buffers is the right two-axis coverage.

The split-policy (encode vs. submit) is elegant — I'd been half-expecting a Flush per allocation, which would have serialized intermediate-tensor allocations against dispatch batching and tanked perf on shape-varying models. The run_active_ gate avoids that.

Correctness checks I did

• wgpu::Buffer::Acquire(buffer) + buffer_guard.MoveToCHandle() — RAII wrapper around the raw handle guarantees the buffer isn't leaked if EncodeDeferredDispatches / EndComputePass / Flush throw. ✓
• Ordering of EncodeDeferredDispatches → EndComputePass → ClearBuffer — deferred programs get encoded first (may open a compute pass), then the pass is closed (required before non-compute ops), then the clear is appended. ✓
• Empty callback default (should_submit_zero_initialize_() guarded by should_submit_zero_initialize_ &&) — safe. The WebGpuNoOpAllocator device-free path also doesn't wire it, which is fine since there's no real device to clear. ✓
• Test DoesNotCaptureDeviceAllocatorBufferClear — constructs GpuBufferAllocator with no should_submit_zero_initialize, verifies captured_commands stays empty across a Flush inside CaptureBegin/CaptureEnd. Assumes Flush never adds CapturedCommandInfo for raw ClearBuffer calls — that's true because capture entries are built by ProgramManager::Run, not by the encoder. Worth a brief comment in the test explaining why "flush inside capture" is the interesting property being tested.
• Thread safety of run_active_: WebGPU EP has ConcurrentRunSupported() == false (verified in session_buffer_pool.h — InferenceSession serializes Run under session_mutex_), so single-writer/single-reader chain for the same session's allocator lambda. The shared allocator uses []() { return true; } and doesn't touch run_active_. No race in current code paths.

Comments / suggestions

1. Make run_active_ std::atomic<bool> defensively. It's essentially free and prevents a subtle future footgun. Right now it's a plain bool written in OnRunStart/OnRunEnd and read from an allocator callback whose call site is typically on the same execution thread — but a shared-allocator user, or an IOBinding path that runs on a different thread, could read it while OnRunStart on another session's thread is writing (different sessions share the ORT process). Even under session_mutex, this is a happens-before boundary that TSan will flag when the WebGPU EP eventually gets thread-sanitized. std::atomic<bool> with default memory_order_seq_cst costs nothing observable and documents intent.

2. Comment the two-axis zero-init strategy on the initialize_to_zero_ field. Reading allocator.h alone, it's not obvious that the field only controls the cached path and that fresh buffers are covered by the Dawn toggle. Two lines above the member would save the next reader the archaeology:

   // Zero-initialize writable buffers on Alloc. Cached (reused) buffers are
   // cleared explicitly via BufferManager::Create -> ClearBuffer; fresh
   // buffers rely on the Dawn "lazy_clear_resource_on_first_use" toggle
   // enabled in WebGpuContext.
   bool initialize_to_zero_;

3. Consider an opt-out knob for callers who fully overwrite. The shared-allocator path now pays for a ClearBuffer + immediate flush on every reused buffer, even when the caller is about to MemCpy/Upload the entire range. That's the safe default and I wouldn't block on this, but for high-frequency IO-binding users a per-allocation opt-out (or a separate unsafe_shared_allocator factory) is worth filing as a follow-up. Related to already-filed [Feature Request] [WebGPU] Expose lazy_clear_resource_on_first_use option #32080 for exposing the Dawn toggle.

4. ORT_ENFORCE in EncodeDeferredDispatches is on a hot path. Cheap condition, fine as-is. But since the enclosing function returns Status, ORT_RETURN_IF_NOT would be more consistent with the file's style (the check right below uses ORT_RETURN_IF_ERROR). Nit only.

5. Test WebGpuDispatchBatchingTests.ProfilingDispatchCountStaysBoundedAcrossAllocatorClears relies on the fact that if the ORT_ENFORCE fires, Session::Run throws and the try/catch captures it as run_error. That's implicit — an explicit comment "the point of this test is that the enforce below does not fire" would make future breakers understand the intent. Also, the profiling-file cleanup uses gsl::finally — good pattern, but the try/catch doesn't rethrow after populating run_error, so if session.Run throws before writing warm_output, the second_output block is skipped and the two EXPECT_FLOAT_EQs over second_output run on a default-constructed empty vector — which passes the size-based assertion path how? Actually ASSERT_EQ(second_output.size(), expected_output.size()) will fail loudly. OK, safe.

6. Nit — commit history is noisy. 14 commits with a mix of "fix", "refactor", "test", "style" is fine for the review process, but please squash to one or two logical commits at merge (allocator zero-init + dispatch-window accounting) so the git log stays useful for future bisection.

Non-issues I verified

• is_read_only_allocator=true → initialize_to_zero_=false → initializers aren't cleared. Safe because initializer buffers are Upload-ed immediately after Alloc, fully overwriting any prior contents. ✓
• Factory site factory.cc line 249 (CreateAllocatorImpl — the public shared allocator) uses []() { return true; } → always submits. That's the correct answer for external users. ✓
• CreatePreferredAllocators in webgpu_execution_provider.cc and CreateEpImpl in factory.cc both use !IsRunActive() — correct dual-mode behavior. ✓
• The SessionAllocatorDefersReusedBufferClearDuringRun test proves the deferral end-to-end (Upload → Free → OnRunStart → Alloc-reuse → external readback sees stale → OnRunEnd → external readback sees zero). Exactly the right regression. ✓

Recommendation

Approve after (1) std::atomic<bool> for run_active_ and (2) the doc comment on initialize_to_zero_. (3) is a follow-up. Commit squash at merge. Nice fix — the split between encode-order and submit-policy is exactly the right factoring.