Skip to content

[WebGPU] Support concurrent use with explicit Session streams - #32587

Open
Xiaofei Han (xiaofeihan1) wants to merge 2 commits into
mainfrom
xiaofeihan/webgpu-explicit-stream-concurrency
Open

Xiaofei Han (xiaofeihan1) wants to merge 2 commits into
mainfrom
xiaofeihan/webgpu-explicit-stream-concurrency

Conversation

@xiaofeihan1

Copy link
Copy Markdown
Contributor

Summary

Publish the explicit Session-stream implementation at commit 2bef0d899383a72ed153926e833480e67ec2df24 as 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.

  • Isolate command recording per Session while sharing ordinary BufferManagers at the WebGpuContext level. Captured graphs retain per-graph BufferManagers.
  • Implement OrtEp::CreateSyncStreamForDevice so stream-bearing copies select their owning EP's recording explicitly, preserving ordering between deferred compute and copies without TLS or creation-thread lookup.
  • Expose OrtAllocator::AllocOnStream: Session-stream allocations defer cached-buffer clears; plain allocations submit clears before returning.
  • Keep stream-less environment transfers on private recording state, protect mutable recording/cache state, and use conservative completion waits for stream flush and notification activation.
  • Route graph Memcpy kernels through their owning EP because the generic single-tensor transfer wrapper in this version drops the stream argument. No new public EP API is introduced.

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

  • Requires an ORT build with stream support.
  • Concurrent graph capture, concurrent profiling, cross-device transfer, and arbitrary foreign stream overrides are not established by the historical tests.
  • Conservative waits prioritize correctness over overlap; performance has not been measured for this version.
  • This is the original historical branch, not a rebase onto current main. Integration conflicts or CI failures may need follow-up.

Copilot AI balanced review requested due to automatic review settings September 14, 2026 15:36

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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_dispatches is encoded. For a graph path that copies a value after a deferred compute dispatch, Flush will submit this copy before that compute, so the destination can contain stale data. Encode the deferred dispatches before recording CopyBufferToBuffer, matching Download and 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 public CreateSyncStreamForEpDevice API still unconditionally returns ORT_NOT_IMPLEMENTED. Internal session setup calls Ep::CreateSyncStreamForDeviceImpl and therefore hides this mismatch, while applications creating a stream from an OrtEpDevice will 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/PopErrorScope operate on the shared WGPUDevice error-scope stack. OnRunStart and OnRunEnd call 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

  • OnRunStart pushes a validation error scope on context_.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.

Comment thread onnxruntime/core/providers/webgpu/allocator.cc
Comment thread onnxruntime/core/providers/webgpu/buffer_manager.cc
Comment thread onnxruntime/core/providers/webgpu/compute_context.h
Comment thread onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc
@xiaofeihan1

Copy link
Copy Markdown
Contributor Author

/run bots

@xiaofeihan1
Xiaofei Han (xiaofeihan1) force-pushed the xiaofeihan/webgpu-explicit-stream-concurrency branch from 22fefef to c3eb238 Compare September 18, 2026 03:27
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>
@xiaofeihan1
Xiaofei Han (xiaofeihan1) force-pushed the xiaofeihan/webgpu-explicit-stream-concurrency branch from c3eb238 to be4e4e3 Compare September 18, 2026 03:37
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants