[WebGPU] Separate command recording from buffer reuse tracking by qjia7 · Pull Request #32874 · microsoft/onnxruntime

Jiajia Qin (qjia7) · 2026-09-28T14:54:03Z

Description

Remove the mutex and pending-buffer list from CommandRecordingState. Each Session owns its command recording, while plugin streamless allocations and transfers use independent recordings. BufferManager owns deferred buffer releases and protects them with its existing cache mutex.

Plain plugin allocations submit cached-buffer clears before returning. Streamless GPU copies also submit before returning, so later Session work on the shared queue observes them. Matching explicit-stream operations continue to use the Session recording.
Reusable buffers remain pending until their recording submits. Pending entries are grouped by recording; another recording's submission cannot return them to the pool. An atomic work flag lets concurrent frees distinguish pending from submitted work without inspecting encoders or capture state. Non-reusing caches bypass this tracking.
Flush retires entries in both the active manager and the context-shared manager. Graph cache policy applies only to the active manager, so frees during replay and uniforms used during prepacking return to the correct pools. Failure and Session destruction discard abandoned entries.
Keep concurrency support scoped to plugin WebGPU. Preserve the standalone shader-dump concurrency test and sequential buffer-pool memory diagnostic while removing tests that require built-in inference concurrency.

Follow-up to #32587. Initial-capture coverage includes independent Sessions and serialized Runs for different graph IDs within one Session, using preallocated bindings.

Dependency

Builds on #32902, now merged into main, which serializes graph-state checks and replay with normal Session execution. This PR targets main and contains only the WebGPU changes. Standard PR CI runs automatically on updates.

Validation

Native Windows Release validation before rebasing onto the latest main:

Plugin tests: 43 passed across suite and focused runs; two existing tests disabled.
Built-in WebGPU suite: 37 tests passed; one configuration-specific adapter test skipped.
Core/session suite: 78 tests passed.
Plugin concurrency stress: 12 tests repeated 20 times with shuffled order, all 240 executions passed.
Initial capture and replay: three tests repeated 20 times with shuffled order, all 60 executions passed.
Restored shader-dump test and sequential memory diagnostic: both passed. The diagnostic was explicitly enabled with two Sessions, 1 MB tensors, two iterations and no hold time; it remains disabled by default.
Lint and whitespace checks passed.

Rebase verification: all five WebGPU patches are unchanged according to git range-diff, and git diff --check passed.

### Description Captured-graph replay in `InferenceSession::Run()` bypasses the existing session lock even when an EP returns false from `ConcurrentRunSupported()`. Replay and graph-state queries can therefore overlap another replay, ordinary execution, or `ReleaseCapturedGraph()` on the same session. Move the existing lock acquisition to the outer `Run()` wrapper so it covers graph-state checks, replay, normal execution, and internal capture retries. `RunImpl()` retries retain the lock without acquiring it recursively. EPs that support concurrent runs continue to bypass the lock. Add CPU-only regression tests using a capture-capable test EP without an internal lock. They cover replay serialization with other operations, graph-state queries during release, capture retries and recapture, replay failure cleanup, and preservation of concurrency for EPs that support it. ### Motivation and Context This is a core prerequisite for #32874, which removes WebGPU's command-recording mutex. Session serialization must cover captured-graph replay so EPs that disallow concurrent runs can rely on the core contract. ### Validation - Windows Release builds of `onnxruntime_test_all` and `onnxruntime_provider_test` passed. - All six new serialization cases failed against the unfixed core and passed with this change. - 111 selected core tests passed, including all nine new tests and the existing session and plugin EP interface tests. - 37 plugin WebGPU tests passed, covering graph capture, concurrency, allocators, and dispatch batching; two existing tests remain disabled. - `lintrunner -a` and `git diff --check` passed.

Retire submitted buffers in both active and shared managers while preserving graph cache policy. Cover concurrent allocator frees, initial graph capture and replay, and failed flush cleanup. Keep SupportsBufferReuse as the cache capability name.

The base branch was changed.

Copilot

Copilot review overview

🔵 Needs a closer look

The cross-session recording, graph-capture, and buffer-lifetime concurrency changes warrant final human validation despite comprehensive tests.

Review effort: Balanced
Findings: None

What changed in this PR

Separates WebGPU command recording from thread-safe buffer reuse tracking, primarily enabling plugin EP concurrency without per-recording mutexes.

Changes:

Moves deferred buffer ownership and synchronization into BufferManager.
Uses independent recordings for plugin streamless allocations and transfers.
Expands plugin concurrency, graph-capture, and buffer-reuse coverage while removing unsupported built-in concurrency tests.

No actionable defects were identified.

File	Description
`onnxruntime/test/providers/webgpu/webgpu_stream_regression_test.cc`	Replaces mutex probing with clear correctness coverage.
`onnxruntime/test/providers/webgpu/webgpu_context_test.cc`	Adds recording-aware buffer lifecycle tests.
`onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc`	Expands plugin concurrency and capture tests.
`onnxruntime/test/providers/webgpu/concurrent_context_test.cc`	Removes unsupported built-in concurrency cases.
`onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc`	Creates per-call plugin transfer recordings.
`onnxruntime/core/providers/webgpu/webgpu_execution_provider.h`	Clarifies same-Session serialization.
`onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc`	Cleans abandoned buffers and tightens graph routing.
`onnxruntime/core/providers/webgpu/webgpu_context.h`	Removes recording mutex and adds atomic work state.
`onnxruntime/core/providers/webgpu/webgpu_context.cc`	Retires buffers per recording and manager.
`onnxruntime/core/providers/webgpu/ep/sync_stream.cc`	Removes obsolete recording locking.
`onnxruntime/core/providers/webgpu/ep/README.md`	Documents concurrency and ordering guarantees.
`onnxruntime/core/providers/webgpu/ep/factory.cc`	Configures independent plugin allocation recording.
`onnxruntime/core/providers/webgpu/ep/allocator.h`	Updates allocator ownership documentation.
`onnxruntime/core/providers/webgpu/ep/allocator.cc`	Routes stream allocations to Session recording.
`onnxruntime/core/providers/webgpu/data_transfer.h`	Documents call-local transfer recording.
`onnxruntime/core/providers/webgpu/data_transfer.cc`	Removes obsolete recording locking.
`onnxruntime/core/providers/webgpu/compute_context.h`	Removes obsolete FillZero locking.
`onnxruntime/core/providers/webgpu/buffer_manager.h`	Defines recording-grouped pending releases.
`onnxruntime/core/providers/webgpu/buffer_manager.cc`	Implements synchronized deferred reuse and discard.
`onnxruntime/core/providers/webgpu/allocator.h`	Adds atomic statistics and explicit recording allocation.
`onnxruntime/core/providers/webgpu/allocator.cc`	Implements independent plugin allocation recordings.

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Xiaofei Han (xiaofeihan1) · 2026-10-02T14:13:36Z

LGTM, thanks for your help!

…and serialized execution (#33082) ### Description  Restore WebGPU plugin compatibility starting with **ORT 1.24.4**, while retaining **Session-owned command recording in both execution modes**. The additional real-host matrix below found an unresolved **ORT 1.26.0 host-side rejection**; this is not a claim that every intervening release currently works. - Each Session exclusively owns its recording through `unique_ptr`; there is no context-owned shared legacy recording. - Select the older-host compatibility policy with `UseSerializedExecutionMode()` and `ShouldUseSerializedExecutionMode()`, now isolated in `ep/runtime_compatibility.h/.cc` rather than the stream implementation. - In serialized mode, Session `Alloc` and matching `AllocOnStream` use `!IsRunActive()` to submit cached-buffer clears outside Run and defer them during Run. Kernel scratch uses ordinary Tensor allocation through the same allocator, without an `AllocForKernel` API or the unavailable old-host `KernelContext_GetSyncStream`. - Env allocations and concurrent-mode ordinary Session allocations retain independent, immediately submitted clears. Cached-buffer clearing and kernel dispatch batching remain enabled. - For streamless framework copies, submit pending active-Session work before recording the copy independently. The context borrows a non-owning pointer to that recording; it does not own or share the command state. Session teardown discards only that Session's pending entries. - Remove the context-wide Run gate and rely on the required caller-serialization contract. `SerialRunsAfterRejectedStreamOverride` covers a rejected stream override followed by successful serial Runs on both the affected Session and another Session. | Host ORT version | Default execution mode | | --- | --- | | 1.26.0 | Currently blocked by the host's plugin-version check; see validation below | | 1.28.x starting at 1.28.3 | Concurrent independent Sessions | | 1.30.x starting at 1.30.1 | Concurrent independent Sessions | | 1.31 and later | Concurrent independent Sessions | | Other supported versions, including all 1.29.x | Serialized compatibility | Set `ORT_WEBGPU_EP_FORCE_LEGACY=1` **before plugin registration** to force serialized compatibility. The existing override and `onnxruntime_webgpu_legacy_test` CTest name are retained for compatibility; neither implies shared recording. The override cannot enable concurrency on an unsupported host. **Serialized mode requires callers to serialize all WebGPU operations on the same device**, including Session creation/destruction, Run, allocation, I/O binding, and Env transfers. Use sequential graph execution. This is not a fixed-CPU-thread requirement, and the plugin neither detects nor serializes overlapping operations. Callback reentry remains unsupported. #### Input/output copy ordering ORT 1.24.4 forwards supplied streams through batched `CopyTensors`; its separate single-tensor wrapper drops the stream. `BindInput` uses a single-input copy, and mixed source/destination device pairs can make a batch fall back to individual copies. | Path | Ordering mechanism | | --- | --- | | `BindInput` outside Run, CPU to GPU | Allocation submits any cached-buffer clear before returning, including through `AllocOnStream`. A subsequent streamless upload is ordered after it on the device queue; no active Run pointer is needed. | | Input copy during Run, CPU to GPU | With a stream, the clear and upload use the same Session recording. Without a stream, flush the active Session first, then upload using a local recording. | | Output copy during Run, GPU to CPU | Submit the producer before readback, either through its explicit stream or the active-Session pre-copy flush. This does not depend on `OnRunEnd`, which occurs after output copies. | | Mixed-direction inputs/outputs | Each streamless copy follows the same submission boundary if the old host falls back to individual copies. | | Streamless GPU-to-GPU copy within one context | Submit preceding active Session work, then submit the local copy before returning. Cross-context copies remain unsupported. | These are submission boundaries, not per-kernel CPU waits. Readback retains its existing completion wait. Compared with a shared recording, a streamless copy with pending Session work can require an additional queue submission; performance equivalence has not been established. ### Motivation and Context  Allow independently distributed WebGPU plugins to support older applications without requiring the host-side single-copy stream-forwarding fix in #32666. Builds on #32874. Following review, the implementation now keeps Session-owned recording rather than introducing a second ownership model for old hosts. A real ORT 1.24.4 mixed-output repro demonstrated the missing submission boundary: an independent GPU Neg branch fetched to CPU returned 0 instead of -1 when the Session's pending computation was not submitted before a streamless readback. The other independent CPU branch was fetched to GPU, triggering mixed-direction single-copy fallback. Both targets and the GPU input were preallocated. Adding the boundary makes both output orderings pass; this shows the need for ordering coordination, not shared recording ownership. The alternative developed in #33228 has been incorporated here, and that duplicate PR is closed. This revision includes upstream `main` at `e19a9ae831`. **Scope:** Removing the context-wide gate avoids permanently blocking other Sessions when a host error skips `OnRunEnd`. It does not guarantee recovery of all unfinished Run/capture state. Clarifying the host's `OnRunStart`/`OnRunEnd` error-path contract and complete failed-replay recovery remain separate follow-up work. ### Validation Validated commit **`20ff4e2f63`** on Windows x64 Release with the shared WebGPU plugin and EP API adapters. Current-host tests use ORT **1.32.0**. `W:` was a temporary mapping to this worktree and was removed afterward. **Build: passed.** ```powershell cmake --build W:\build\WebGpuCompat\Release --config Release --target onnxruntime_providers_webgpu onnxruntime_provider_test --parallel 8 -- /nologo /verbosity:minimal ``` **Serialized compatibility: 41 executed tests passed, no failures or skips.** CTest sets `ORT_WEBGPU_EP_FORCE_LEGACY=1` before loading the plugin and selects serial-safe cases. ```powershell Set-Location W:\build\WebGpuCompat\Release ctest -C Release -R '^onnxruntime_webgpu_legacy_test$' --output-on-failure ``` **Concurrent-mode selection: 57 executed tests passed, no failures or skips.** Two pre-existing disabled diagnostic tests were not enabled in either selection. ```powershell Set-Location W:\build\WebGpuCompat\Release\Release $env:ORT_UNIT_TEST_MAIN_DYNAMIC_PLUGIN_EP_CONFIG_JSON='{"ep_library_registration_name":"WebGPU_PluginEP","ep_library_path":"onnxruntime_providers_webgpu.dll","selected_ep_name":"WebGpuExecutionProvider"}' $env:ORT_WEBGPU_EP_FORCE_LEGACY='0' .\onnxruntime_provider_test.exe '--gtest_filter=WebGpuPluginRecordingModeTest.*:WebGpuPluginSharedAllocatorRegistrationTest.*:WebGpuPluginSharedAllocatorTest.*:WebGpuSessionAllocatorDeathTest.*:PluginEpWebGpuConcurrency.*' ``` Coverage includes version boundaries, allocation-clear submission, scratch batching, serial and concurrent capture/replay, mixed feeds/fetches, and serial Runs after a rejected stream override. **Actual ORT 1.24.4: all 9 local-harness scenarios passed, 20 iterations each.** The standalone API-24 C++ probe loaded the official 1.24.4 DLL and the plugin rebuilt from this commit; profiles confirmed CPU/WebGPU node placement. Cases cover CPU I/O, dirty-buffer BindInput, mixed-direction inputs and outputs in both orders, GPU-to-CPU-to-GPU graph execution, scratch reuse, and capture-enabled execution. The following invocation used the existing local `recording-ownership-probe` directory (`$probe`); this harness is not part of the repository or CI: ```powershell $variant='serialized-policy-20ff4e2f63' $dest=Join-Path $probe $variant $env:ORT_WEBGPU_EP_FORCE_LEGACY='1' python "$probe\run_probe.py" $variant "$dest\onnxruntime_providers_webgpu.dll" ``` **Capture limitation:** the actual 1.24.4 run exercised initial capture and its internal replay. Its subsequent Runs did not exercise host-driven replay, because that host lacks the plugin capture/replay callbacks. Repeated host-driven replay was covered by the current-host tests above. Other operating systems and performance benchmarks were not tested in the above `20ff4e2f63` validation. #### Additional local validation at `62e3504c2e` Removed only the redundant pre-capture `OnRunStart` flush. `CaptureBegin` still flushes before entering capture; streamless-copy, Run-end, and standalone replay submission boundaries are unchanged. The plugin rebuilt successfully, and the three focused serialized capture tests passed: ```powershell $env:ORT_UNIT_TEST_MAIN_DYNAMIC_PLUGIN_EP_CONFIG_JSON='{"ep_library_registration_name":"WebGPU_PluginEP","ep_library_path":"onnxruntime_providers_webgpu.dll","selected_ep_name":"WebGpuExecutionProvider"}' $env:ORT_WEBGPU_EP_FORCE_LEGACY='1' .\onnxruntime_provider_test.exe '--gtest_filter=PluginEpWebGpuConcurrency.SerialGraphCaptureWithCpuFeedsAndFetches:PluginEpWebGpuConcurrency.SerialSingleSessionMultipleGraphCaptureIds:PluginEpWebGpuConcurrency.GraphCaptureReplayInterleavedWithIdleSessionCpuBindInput' ``` A separate local matrix used **one unchanged plugin DLL** with official Windows x64 cores 1.24.4, 1.25.1, 1.26.0, 1.27.1, 1.28.3, 1.29.1, and 1.30.0, plus the local 1.32.0 Release core. Loaded module paths, actual core versions, and plugin hashes were checked. - **Native copy/allocation/capture probe: 63 scenarios passed, 9 failed**, with 20 iterations per scenario. All nine scenarios passed on every tested core except 1.26.0. That core rejects `OrtEp::ort_version_supported=32` because its kernel-creation check requires a value no greater than 26. This host restriction remains unresolved; forcing serialized mode does not address it. - **Capture support differs from execution-mode eligibility.** In the simple native capture probe, 1.24.4 and 1.25.1 performed one initial internal replay; the other passing cores performed repeated host-driven replay. - **Local Phi-4 workload, capture OFF:** all seven usable cores passed fresh-generator repeatability, serially interleaved-generator checks, multi-turn checks, and a repeated exact 1024-input/300-output-token workload. An isolated GenAI build compiled against API 24 was held fixed across cores; the normal installed package requires API 29. Performance measurements used two warmups and five measured runs, generator reuse and shared KV buffers, with profiling disabled. - **Capture ON is not valid for this model graph.** It contains an `If` node. Tested cores from 1.26 onward reject capture at Session initialization. On 1.24.4 and 1.25.1, initialization succeeds but generated tokens differ after rewinding and repeating the same prompt. This also reproduces without profiling and with the pre-cleanup `20ff4e2f63` plugin, so it was not introduced by removing the redundant flush. No capture-ON throughput is reported for this model. - **Performance interpretation:** default core graph optimizations were retained; newer cores use additional MLP fusion. Cross-version throughput is therefore not an isolated measurement of serialized-mode overhead. Generated tokens were repeatable within each core but not identical between the older and newer optimizer groups. The matrix scripts and model are session-local, not checked into the repository or run by CI. `$matrix` below refers to the isolated `phi4-core-matrix` artifact directory: ```powershell python "$matrix\run_native.py" # failed overall: 63 passed, 9 failed on 1.26.0 python "$matrix\run_phi4.py" --phase smoke # failed overall: 7 passed, 9 unsupported/failing combinations python "$matrix\run_phi4.py" --phase correctness # passed: 21 checks; failing smoke combinations excluded python "$matrix\run_phi4.py" --phase benchmark # completed: 7 capture-OFF configurations; failing smoke combinations excluded ``` --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b6fbaacb-9846-4330-8fb3-2f0b13737d48 Copilot-Session: 8ca4e5fc-8c08-46bd-82a9-5cff5108efed

Jiajia Qin (qjia7) requested a review from Xiaofei Han (xiaofeihan1) September 29, 2026 05:26

Jiajia Qin (qjia7) marked this pull request as ready for review September 29, 2026 05:26

Jiajia Qin (qjia7) requested a review from Edward Chen (edgchen1) September 29, 2026 05:27

Jiajia Qin (qjia7) mentioned this pull request Sep 29, 2026

[Core] Serialize captured graph replay with session runs #32902

Merged

Edward Chen (edgchen1) reviewed Sep 29, 2026

View reviewed changes

Comment thread onnxruntime/core/providers/webgpu/allocator.cc Outdated

Comment thread onnxruntime/core/providers/webgpu/buffer_manager.cc

Jiajia Qin (qjia7) marked this pull request as draft September 30, 2026 05:44

Jiajia Qin (qjia7) changed the base branch from main to users/qjia7/pr-32902-base September 30, 2026 12:14

Jiajia Qin (qjia7) force-pushed the refactor/webgpu-session-command-recording branch from 0af48bf to 71d1970 Compare September 30, 2026 12:14

Jiajia Qin (qjia7) marked this pull request as ready for review September 30, 2026 12:24

Jiajia Qin (qjia7) requested a review from Edward Chen (edgchen1) September 30, 2026 14:38

Edward Chen (edgchen1) previously approved these changes Sep 30, 2026

View reviewed changes

Jiajia Qin (qjia7) added 5 commits October 1, 2026 09:44

Refactor WebGPU command recording and buffer reuse tracking

34e0783

Track WebGPU pending buffers on release

1df1754

Avoid session capture state reads in WebGPU streamless allocation

0fd1626

Restore WebGPU shared-helper tests and clarify allocator policy

0493b97

Jiajia Qin (qjia7) changed the base branch from users/qjia7/pr-32902-base to main October 1, 2026 01:45

Jiajia Qin (qjia7) force-pushed the refactor/webgpu-session-command-recording branch from 71d1970 to 0493b97 Compare October 1, 2026 01:45

Merge main into WebGPU command recording refactor

cbb697b

Copilot AI balanced review requested due to automatic review settings October 1, 2026 13:29

Copilot started reviewing on behalf of Jiajia Qin (qjia7) October 1, 2026 13:29

View session

Copilot AI reviewed Oct 1, 2026

View reviewed changes

Edward Chen (edgchen1) approved these changes Oct 1, 2026

View reviewed changes

Xiaofei Han (xiaofeihan1) reviewed Oct 1, 2026

View reviewed changes

Comment thread onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc

Comment thread onnxruntime/core/providers/webgpu/allocator.cc

Xiaofei Han (xiaofeihan1) approved these changes Oct 2, 2026

View reviewed changes

Jiajia Qin (qjia7) merged commit 4ae10e3 into microsoft:main Oct 3, 2026
96 checks passed

Jiajia Qin (qjia7) mentioned this pull request Oct 3, 2026

Support plugin EPs in the C++ model benchmark microsoft/onnxruntime-genai#2670

Merged

Xiaofei Han (xiaofeihan1) mentioned this pull request Oct 3, 2026

[WebGPU] Restore ORT 1.24.4 compatibility with per-Session recording and serialized execution #33082

Merged

Jiajia Qin (qjia7) mentioned this pull request Oct 7, 2026

[WebGPU] Accelerated model loading using Win32 and D3D12 #33055

Merged

Navigation Menu

Repository navigation

[WebGPU] Separate command recording from buffer reuse tracking - #32874

[WebGPU] Separate command recording from buffer reuse tracking#32874
Jiajia Qin (qjia7) merged 6 commits into
microsoft:mainfrom
qjia7:refactor/webgpu-session-command-recording

Jiajia Qin (qjia7) commented Sep 28, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Xiaofei Han (xiaofeihan1) commented Oct 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

Jiajia Qin (qjia7) commented Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Dependency

Validation

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🔵 Needs a closer look

Uh oh!

Uh oh!

Uh oh!

Xiaofei Han (xiaofeihan1) commented Oct 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Jiajia Qin (qjia7) commented Sep 28, 2026 •

edited

Loading