[WebGPU] Support concurrent use with explicit Session streams - #32587
Xiaofei Han (xiaofeihan1) wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
Critical and moderate review findings remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds per-Session WebGPU command recording and stream-aware allocation/transfer handling while synchronizing shared context resources.
Changes:
- Isolates Session recording and graph-capture state.
- Adds explicit sync streams and stream-aware allocator behavior.
- Adds native and AutoEP concurrency coverage.
Unresolved findings affect allocator exception safety, dispatch ordering, synchronization, stream registration, error-scope ownership, and graph-copy routing.
File summaries
| File | Summary |
|---|---|
onnxruntime/test/providers/webgpu/webgpu_context_test.cc |
Updates context tests for explicit recordings. |
onnxruntime/test/providers/webgpu/concurrent_context_test.cc |
Adds shared-context concurrency coverage. |
onnxruntime/test/autoep/test_webgpu_concurrency.cc |
Adds AutoEP concurrency coverage. |
onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc |
Adds plugin transfer and allocator integration. |
onnxruntime/core/providers/webgpu/webgpu_kernel.cc |
Integrates kernel recording and graph-copy routing. |
onnxruntime/core/providers/webgpu/webgpu_execution_provider.h |
Defines Session recording ownership. |
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc |
Handles EP lifecycle, recording, flushing, and Memcpy routing. |
onnxruntime/core/providers/webgpu/webgpu_context.h |
Defines recording-aware context state. |
onnxruntime/core/providers/webgpu/webgpu_context.cc |
Implements recording, flushing, and capture handling. |
onnxruntime/core/providers/webgpu/session_buffer_pool.cc |
Handles Session buffer-pool management. |
onnxruntime/core/providers/webgpu/program_manager.h |
Updates pipeline-cache interfaces. |
onnxruntime/core/providers/webgpu/program_manager.cc |
Synchronizes shared pipeline-cache access. |
onnxruntime/core/providers/webgpu/ep/README.md |
Documents Session streams and limitations. |
onnxruntime/core/providers/webgpu/ep/factory.cc |
Registers plugin stream integration. |
onnxruntime/core/providers/webgpu/ep/ep.h |
Declares EP sync-stream interfaces. |
onnxruntime/core/providers/webgpu/ep/ep.cc |
Implements sync-stream creation. |
onnxruntime/core/providers/webgpu/data_transfer.h |
Defines stream-bound transfer interfaces. |
onnxruntime/core/providers/webgpu/data_transfer.cc |
Implements stream and notification handling. |
onnxruntime/core/providers/webgpu/compute_context.h |
Updates recording-aware compute helpers. |
onnxruntime/core/providers/webgpu/compute_context.cc |
Supports per-recording compute encoders. |
onnxruntime/core/providers/webgpu/buffer_manager.h |
Defines recording-scoped buffer APIs and synchronization. |
onnxruntime/core/providers/webgpu/buffer_manager.cc |
Synchronizes shared buffer-cache operations. |
onnxruntime/core/providers/webgpu/allocator.h |
Declares stream-aware allocator interfaces. |
onnxruntime/core/providers/webgpu/allocator.cc |
Implements recording-aware and stream-aware allocation. |
Review details
Suppressed comments (4)
onnxruntime/core/providers/webgpu/buffer_manager.cc:545
- The GPU-to-GPU copy is recorded before
recording.deferred_dispatchesis encoded. For a graph path that copies a value after a deferred compute dispatch,Flushwill submit this copy before that compute, so the destination can contain stale data. Encode the deferred dispatches before recordingCopyBufferToBuffer, matchingDownloadand cached-buffer clearing.
auto& command_encoder = context_.GetCommandEncoder(recording);
context_.EndComputePass(recording);
command_encoder.CopyBufferToBuffer(src, 0, dst, 0, copy_size);
onnxruntime/core/providers/webgpu/ep/factory.cc:296
IsStreamAwareImpl()now advertises factory-level stream support, but the callback used by the publicCreateSyncStreamForEpDeviceAPI still unconditionally returnsORT_NOT_IMPLEMENTED. Internal session setup callsEp::CreateSyncStreamForDeviceImpland therefore hides this mismatch, while applications creating a stream from anOrtEpDevicewill fail despite the advertised capability. Please implement the factory callback (or update the awareness/registration contract so this public capability is not advertised while EP-level session streams remain registered).
return true;
}
OrtStatus* ORT_API_CALL Factory::CreateSyncStreamForDeviceImpl(
OrtEpFactory* /*this_ptr*/,
const OrtMemoryDevice* /*memory_device*/,
const OrtKeyValuePairs* /*stream_options*/,
OrtSyncStreamImpl** stream) noexcept {
EXCEPTION_TO_RETURNED_STATUS_BEGIN
*stream = nullptr;
return Api().ort.CreateStatus(ORT_NOT_IMPLEMENTED,
"CreateSyncStreamForDevice is not implemented for this EP factory.");
onnxruntime/core/providers/webgpu/webgpu_context.h:114
- This mutex only serializes one Session's recording, while
WebGpuContext::PushErrorScope/PopErrorScopeoperate on the sharedWGPUDeviceerror-scope stack.OnRunStartandOnRunEndcall them for every Session, so interleaved runs can pop another Session's scope, misattribute validation errors, or fail to pop. Protect the device-wide scope lifetime (or disable per-Session scopes when a device is shared); locking this per-recording mutex alone cannot establish the required LIFO pairing.
std::recursive_mutex mutex;
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc:880
OnRunStartpushes a validation error scope oncontext_.Device()and this path later pops it, but that scope stack is shared by every Session using context 0. Two runs can interleave A.Push, B.Push, A.Pop, B.Pop, causing each run to consume the other's validation errors and report a failure to the wrong caller; implicit device synchronization does not make the nesting per-session. The per-recording lock here cannot fix that, so serialize the complete error-scope lifetime or disable per-run scopes for shared concurrent contexts.
Status flush_status = context_.Flush(BufferManager(), *recording_);
- Files reviewed: 24/24 changed files
- Comments generated: 4
- Review effort level: Lite (auto)
Note
Copilot is running an experiment and ran this review at Lite.
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
/run bots |
22fefef to
c3eb238
Compare
Isolate command recording per Session while sharing synchronized WebGPU buffer management. Route copies and allocations through explicit Session streams and preserve deferred-dispatch ordering. Contain Session allocator Free exceptions with diagnostics, serialize FillZero recording, and explicitly reject unsupported factory-owned streams. Add focused regression coverage and keep the tests compatible with native/plugin and Dawn build configurations. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
c3eb238 to
be4e4e3
Compare
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Summary
Publish the explicit Session-stream implementation at commit
2bef0d899383a72ed153926e833480e67ec2df24as a separate alternative to #29851. This branch preserves that historical version and its two prerequisite WebGPU concurrency commits, without the later creation-thread-based transfer routing or scope reductions.OrtEp::CreateSyncStreamForDeviceso stream-bearing copies select their owning EP's recording explicitly, preserving ordering between deferred compute and copies without TLS or creation-thread lookup.OrtAllocator::AllocOnStream: Session-stream allocations defer cached-buffer clears; plain allocations submit clears before returning.Validation
Historical validation recorded for this implementation: 12 AutoEP tests passed in three consecutive rounds, including 16-thread same-Session mixed API use, CPU I/O, graph-internal CPU/GPU copies, and graph capture. This is not a claim of parallel Run execution on one Session.
These tests have not been rerun locally for this PR. Fresh PR CI is requested to validate integration with current main.
Limitations