From 34e07836a0a46e9c148a8eb42f3d6080b84a69f5 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 28 Sep 2026 22:53:18 +0800 Subject: [PATCH 1/5] Refactor WebGPU command recording and buffer reuse tracking --- .../core/providers/webgpu/allocator.cc | 40 +- onnxruntime/core/providers/webgpu/allocator.h | 17 +- .../core/providers/webgpu/buffer_manager.cc | 67 +- .../core/providers/webgpu/buffer_manager.h | 15 +- .../core/providers/webgpu/compute_context.h | 3 +- .../core/providers/webgpu/data_transfer.cc | 3 - .../core/providers/webgpu/data_transfer.h | 10 +- .../core/providers/webgpu/ep/README.md | 51 +- .../core/providers/webgpu/ep/allocator.cc | 2 +- .../core/providers/webgpu/ep/allocator.h | 2 +- .../core/providers/webgpu/ep/factory.cc | 12 +- .../core/providers/webgpu/ep/sync_stream.cc | 3 - .../core/providers/webgpu/webgpu_context.cc | 15 +- .../core/providers/webgpu/webgpu_context.h | 28 +- .../webgpu/webgpu_execution_provider.cc | 7 +- .../webgpu/webgpu_execution_provider.h | 3 +- .../webgpu/webgpu_provider_factory.cc | 41 +- .../webgpu/concurrent_context_test.cc | 795 ------------------ .../webgpu/plugin/test_webgpu_concurrency.cc | 15 +- .../providers/webgpu/webgpu_context_test.cc | 67 +- .../webgpu/webgpu_stream_regression_test.cc | 65 +- 21 files changed, 273 insertions(+), 988 deletions(-) delete mode 100644 onnxruntime/test/providers/webgpu/concurrent_context_test.cc diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index bafd54b18b8ab..483a4bcabcb90 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -2,7 +2,6 @@ // Licensed under the MIT License. #include -#include #include #include "core/providers/webgpu/allocator.h" @@ -33,42 +32,45 @@ GpuBufferAllocator::GpuBufferAllocator( // Streamless allocation, e.g., application CreateTensor/Alloc APIs using a Session allocator, // or framework allocations without a stream, including during Run. The plugin's writable device // allocator submits cached clears before returning: the consumer may use a different recording. -// Other allocator roles/native-EP callers can supply a different submission policy. +// Built-in callers can supply a different submission policy. void* GpuBufferAllocator::Alloc(size_t size) { - auto& recording = recording_getter_(); - std::lock_guard lock{recording.mutex}; - return Allocate(size, should_submit_zero_initialize_ && should_submit_zero_initialize_()); +#if defined(ORT_USE_EP_API_ADAPTERS) + CommandRecordingState recording; + if (recording_getter_) { + // Preserve the per-graph cache's release policy without borrowing the Session's commands. + recording.graph_capture_state = recording_getter_().graph_capture_state; + } + return Allocate(size, recording, true); +#else + return Allocate(size, recording_getter_(), should_submit_zero_initialize_ && should_submit_zero_initialize_()); +#endif } -void* GpuBufferAllocator::Allocate(size_t size, bool submit_zero_initialize) { +void* GpuBufferAllocator::Allocate(size_t size, CommandRecordingState& recording, bool submit_zero_initialize) { if (size == 0) { return nullptr; } - auto& recording = recording_getter_(); - std::lock_guard lock{recording.mutex}; - stats_.num_allocs++; - wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; - return buffer_manager_getter_().Create(recording, size, usage, initialize_to_zero_, - submit_zero_initialize); + auto buffer = buffer_manager_getter_().Create(recording, size, usage, initialize_to_zero_, + submit_zero_initialize); + num_allocs_.fetch_add(1, std::memory_order_relaxed); + return buffer; } void GpuBufferAllocator::Free(void* p) { if (p != nullptr) { - auto& recording = recording_getter_(); - std::lock_guard lock{recording.mutex}; - buffer_manager_getter_().Release(recording, static_cast(p)); - stats_.num_allocs--; + buffer_manager_getter_().Release(static_cast(p), + recording_getter_ ? &recording_getter_() : nullptr); + num_allocs_.fetch_sub(1, std::memory_order_relaxed); } } void GpuBufferAllocator::GetStats(AllocatorStats* stats) { - auto& recording = recording_getter_(); - std::lock_guard lock{recording.mutex}; - *stats = stats_; + *stats = AllocatorStats{}; + stats->num_allocs = num_allocs_.load(std::memory_order_relaxed); } WebGpuNoOpAllocator::WebGpuNoOpAllocator(bool is_read_only_allocator) diff --git a/onnxruntime/core/providers/webgpu/allocator.h b/onnxruntime/core/providers/webgpu/allocator.h index d619175d8c687..44824f5624e3f 100644 --- a/onnxruntime/core/providers/webgpu/allocator.h +++ b/onnxruntime/core/providers/webgpu/allocator.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include "core/framework/allocator.h" @@ -20,16 +21,18 @@ inline constexpr OrtDevice WebGpuDevice{OrtDevice::GPU, 0}; // Shared allocation implementation for native and plugin builds. Session getters borrow the EP; -// plugin Env getters retain a context and an independent recording instead. The returned objects -// must remain alive throughout allocator use and tensor frees. Plugin ABI wrappers live in ep/allocator.h. +// plugin Env allocators have no Session recording. The returned objects +// must remain alive throughout allocator use and tensor frees. BufferManager synchronizes deferred +// releases when plugin Session allocators are used concurrently. Plugin ABI wrappers live in ep/allocator.h. class GpuBufferAllocator : public IAllocator { public: // Calls buffer_manager_getter on every Alloc/Free to obtain the current // BufferManager. This allows the EP to route allocations to different // buffer managers (e.g., per-graph) without explicit refresh calls. // Read-only initializers skip cached-buffer clears and can be mapped at creation on UMA. - // should_submit_zero_initialize controls plain Alloc; a matching plugin AllocOnStream - // instead defers clears on the supplied Session stream. + // should_submit_zero_initialize controls built-in Alloc. Plugin Alloc submits independent clears; + // a matching plugin AllocOnStream instead defers clears on the supplied Session stream. + // Only the plugin Env allocator omits recording_getter, as it never uses a Session stream. GpuBufferAllocator(std::function buffer_manager_getter, std::function recording_getter, bool is_read_only_allocator, @@ -40,13 +43,13 @@ class GpuBufferAllocator : public IAllocator { void GetStats(AllocatorStats* stats) override; #if defined(ORT_USE_EP_API_ADAPTERS) - bool IsStreamAware() const override { return true; } + bool IsStreamAware() const override { return static_cast(recording_getter_); } void* AllocOnStream(size_t size, Stream* stream) override; #endif private: - void* Allocate(size_t size, bool submit_zero_initialize); - AllocatorStats stats_; + void* Allocate(size_t size, CommandRecordingState& recording, bool submit_zero_initialize); + std::atomic num_allocs_{0}; std::function buffer_manager_getter_; std::function recording_getter_; std::function should_submit_zero_initialize_; diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 147c31d617bfb..58c562c6a320b 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -27,6 +27,8 @@ void EnforceBufferUnmapped(WebGpuContext& context, WGPUBuffer buffer) { } // namespace class DisabledCacheManager : public IBufferCacheManager { + bool SupportsBufferReuse() const override { return false; } + size_t CalculateBufferSize(size_t request_size) override { return NormalizeBufferSize(request_size); } @@ -48,6 +50,8 @@ class DisabledCacheManager : public IBufferCacheManager { }; class LazyReleaseCacheManager : public IBufferCacheManager { + bool SupportsBufferReuse() const override { return false; } + size_t CalculateBufferSize(size_t request_size) override { return NormalizeBufferSize(request_size); } @@ -492,7 +496,21 @@ BufferManager::BufferManager(WebGpuContext& context, BufferCacheMode storage_buf storage_cache_{CreateBufferCacheManager(storage_buffer_cache_mode)}, uniform_cache_{CreateBufferCacheManager(uniform_buffer_cache_mode)}, query_resolve_cache_{CreateBufferCacheManager(query_resolve_buffer_cache_mode)}, - default_cache_{CreateBufferCacheManager(default_buffer_cache_mode)} { + default_cache_{CreateBufferCacheManager(default_buffer_cache_mode)}, + supports_buffer_reuse_{storage_cache_->SupportsBufferReuse() || uniform_cache_->SupportsBufferReuse() || + query_resolve_cache_->SupportsBufferReuse() || default_cache_->SupportsBufferReuse()} { +} + +void BufferManager::BeginRecording(const CommandRecordingState& recording) const { + if (supports_buffer_reuse_) { + std::lock_guard lock{mutex_}; + pending_buffers_.try_emplace(&recording); + } +} + +const wgpu::CommandEncoder& BufferManager::GetCommandEncoder(CommandRecordingState& recording) const { + BeginRecording(recording); + return context_.GetCommandEncoder(recording); } void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const { @@ -523,7 +541,7 @@ void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuff staging_buffer.Unmap(); ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); - auto& command_encoder = context_.GetCommandEncoder(recording); + auto& command_encoder = GetCommandEncoder(recording); context_.EndComputePass(recording); command_encoder.CopyBufferToBuffer(staging_buffer, 0, dst, 0, copy_size); ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); @@ -542,7 +560,7 @@ void BufferManager::MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGP src_size, ", dst_size=", dst_size, ", copy_size=", copy_size, "."); ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); - auto& command_encoder = context_.GetCommandEncoder(recording); + auto& command_encoder = GetCommandEncoder(recording); context_.EndComputePass(recording); command_encoder.CopyBufferToBuffer(src, 0, dst, 0, copy_size); } @@ -560,14 +578,13 @@ WGPUBuffer BufferManager::Create(CommandRecordingState& recording, size_t size, } if (buffer) { if (initialize_to_zero) { - // initialize_to_zero controls whether a cached buffer is cleared; submit_zero_initialize - // separately controls submission before Create returns. Plugin allocations on an explicit - // Session stream defer clears; plain/null-stream and Env allocations submit them immediately, - // even during Run. Flush submits the whole recording, not just this buffer's clear. + // Plugin plain allocations pass call-local recording and submit only their own clear. + // AllocOnStream defers the clear on the Session timeline. Built-in allocations use the + // Session recording and retain their caller-supplied submission policy. auto buffer_guard = wgpu::Buffer::Acquire(buffer); ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); context_.EndComputePass(recording); - context_.GetCommandEncoder(recording).ClearBuffer(buffer, 0, buffer_size); + GetCommandEncoder(recording).ClearBuffer(buffer, 0, buffer_size); if (submit_zero_initialize) { ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); } @@ -603,15 +620,18 @@ bool BufferManager::SupportsUMA() const { #endif // !defined(__wasm__) } -void BufferManager::Release(CommandRecordingState& recording, WGPUBuffer buffer) const { +void BufferManager::Release(WGPUBuffer buffer, const CommandRecordingState* recording) const { EnforceBufferUnmapped(context_, buffer); - if (recording.has_unsubmitted_work) { - recording.pending_buffers.emplace_back(wgpu::Buffer::Acquire(buffer)); - return; + std::lock_guard lock{mutex_}; + auto& cache = GetCacheManager(buffer); + if (cache.SupportsBufferReuse()) { + if (auto it = pending_buffers_.find(recording); it != pending_buffers_.end()) { + it->second.emplace_back(wgpu::Buffer::Acquire(buffer)); + return; + } } - std::lock_guard lock{mutex_}; - GetCacheManager(buffer).ReleaseBuffer(buffer); + cache.ReleaseBuffer(buffer); } void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, void* dst, size_t size) const { @@ -627,7 +647,7 @@ void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, v desc.usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead; auto staging_buffer = context_.Device().CreateBuffer(&desc); - auto& command_encoder = context_.GetCommandEncoder(recording); + auto& command_encoder = GetCommandEncoder(recording); context_.EndComputePass(recording); command_encoder.CopyBufferToBuffer(src, 0, staging_buffer, 0, buffer_size); ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); @@ -664,10 +684,16 @@ void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, v void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording) const { std::lock_guard lock{mutex_}; - for (auto& buffer : recording.pending_buffers) { - GetCacheManager(buffer.Get()).ReleaseBuffer(buffer.MoveToCHandle()); + if (auto it = pending_buffers_.find(&recording); it != pending_buffers_.end()) { + for (auto& buffer : it->second) { + GetCacheManager(buffer.Get()).ReleaseBuffer(buffer.MoveToCHandle()); + } + pending_buffers_.erase(it); + } + + if (recording.graph_capture_state == GraphCaptureState::Replaying) { + return; } - recording.pending_buffers.clear(); storage_cache_->OnRefresh(recording.graph_capture_state); uniform_cache_->OnRefresh(recording.graph_capture_state); @@ -675,6 +701,11 @@ void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording) cons default_cache_->OnRefresh(recording.graph_capture_state); } +void BufferManager::DiscardPendingBuffers(const CommandRecordingState& recording) const { + std::lock_guard lock{mutex_}; + pending_buffers_.erase(&recording); +} + std::vector> BufferManager::ExtractCachedBuffers(wgpu::BufferUsage usage) { std::lock_guard lock{mutex_}; return GetCacheManager(usage).ExtractCachedBuffers(); diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index 2a17e4dc2d834..8eb644481208c 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/onnxruntime/core/providers/webgpu/buffer_manager.h @@ -10,6 +10,7 @@ #include "core/providers/webgpu/webgpu_external_header.h" +#include "core/common/inlined_containers.h" #include "core/framework/execution_provider.h" namespace onnxruntime { @@ -50,6 +51,9 @@ class IBufferCacheManager { public: virtual ~IBufferCacheManager() = default; + // Reusable buffers must stay out of the cache until their recording submits. + virtual bool SupportsBufferReuse() const { return true; } + // calculate actual buffer size to allocate based on the requested size. virtual size_t CalculateBufferSize(size_t request_size) = 0; @@ -90,15 +94,21 @@ class IBufferCacheManager { class BufferManager { public: BufferManager(WebGpuContext& context, BufferCacheMode storage_buffer_cache_mode, BufferCacheMode uniform_buffer_cache_mode, BufferCacheMode query_resolve_buffer_cache_mode, BufferCacheMode default_buffer_cache_mode); + // Register deferred work before releasing any buffers it uses. Submission or abandonment ends it. + void BeginRecording(const CommandRecordingState& recording) const; + const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) const; void Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const; void MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const; WGPUBuffer Create(CommandRecordingState& recording, size_t size, wgpu::BufferUsage usage, bool initialize_to_zero = false, bool submit_zero_initialize = false) const; bool SupportsUMA() const; // Check if CreateUMA is supported (i.e., the device has BufferMapExtendedUsages feature) - void Release(CommandRecordingState& recording, WGPUBuffer buffer) const; + // A null recording releases directly to the cache (e.g., plugin Env allocations). + void Release(WGPUBuffer buffer, const CommandRecordingState* recording = nullptr) const; void Download(CommandRecordingState& recording, WGPUBuffer src, void* dst, size_t size) const; void RefreshPendingBuffers(CommandRecordingState& recording) const; + // Drop retained references when a recording is abandoned instead of submitted. + void DiscardPendingBuffers(const CommandRecordingState& recording) const; std::vector> ExtractCachedBuffers(wgpu::BufferUsage usage); void AbsorbCachedBuffers(wgpu::BufferUsage usage, @@ -113,6 +123,9 @@ class BufferManager { std::unique_ptr uniform_cache_; std::unique_ptr query_resolve_cache_; std::unique_ptr default_cache_; + const bool supports_buffer_reuse_; + // An entry exists only while that recording has unsubmitted work, even if no buffers were freed. + mutable InlinedHashMap> pending_buffers_; }; class BufferManagerFactory { diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index edac52acf75d6..44547d2efdacb 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/onnxruntime/core/providers/webgpu/compute_context.h @@ -234,10 +234,9 @@ class ComputeContext final : public ComputeContextBase { // inline void FillZero(Tensor& dst) { auto& recording = ep_.Recording(); - std::lock_guard lock{recording.mutex}; ORT_THROW_IF_ERROR(webgpu_context_.EncodeDeferredDispatches(recording)); webgpu_context_.EndComputePass(recording); - auto& command_encoder = webgpu_context_.GetCommandEncoder(recording); + auto& command_encoder = ep_.BufferManager().GetCommandEncoder(recording); WGPUBuffer buffer = reinterpret_cast(dst.MutableDataRaw()); command_encoder.ClearBuffer(buffer, 0, dst.SizeInBytes()); } diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index e669621e2fe0f..3766478429155 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -3,8 +3,6 @@ #include "core/providers/webgpu/data_transfer.h" -#include - #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/webgpu_context.h" @@ -17,7 +15,6 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data, bool dst_is_gpu, size_t bytes) const { auto& command_state = recording_; - std::lock_guard recording_lock{command_state.mutex}; if (bytes > 0) { if (dst_is_gpu) { if (src_is_gpu) { diff --git a/onnxruntime/core/providers/webgpu/data_transfer.h b/onnxruntime/core/providers/webgpu/data_transfer.h index c7f7cea9d353f..0d1a27b90c98b 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.h +++ b/onnxruntime/core/providers/webgpu/data_transfer.h @@ -19,7 +19,7 @@ struct CommandRecordingState; // - Built-in Env, plugin Env, and plugin framework Session: // onnxruntime::plugin_ep::DataTransfer (core/framework/plugin_data_transfer.h) // -> onnxruntime::WebGpuDataTransferImpl (derived from OrtDataTransferImpl) -// -> onnxruntime::webgpu::DataTransferImpl (data_transfer_, lazily allocated). +// -> onnxruntime::webgpu::DataTransferImpl (local to each CopyTensors call). // The outer wrapper holds the C API object through an OrtDataTransferImpl& // and releases it via its Release callback. OrtDataTransferImpl is the C API // base/function table, not an additional, separately allocated wrapper. @@ -45,12 +45,12 @@ struct CommandRecordingState; // for its own adapter DataTransferManager. Adapter kernels use this object, // bypassing the factory's C API wrapper. // -// WebGpuDataTransferImpl, defined in webgpu_provider_factory.cc, owns private -// recording state for streamless copies. Session-level creation alone does not -// bind it to the requesting EP. In plugin builds, an explicit-stream copy uses +// WebGpuDataTransferImpl, defined in webgpu_provider_factory.cc, uses per-call recording +// for plugin streamless copies and object-owned recording in built-in builds. Session-level +// creation alone does not bind it to the requesting EP. In plugin builds, an explicit-stream copy uses // ep/sync_stream.cc to construct a temporary webgpu::DataTransferImpl with the // owning EP's BufferManager and recording, rather than the C API wrapper's -// cached data_transfer_. +// call-local recording. // // DataTransferImpl is the shared raw-pointer copy implementation used by both // DataTransfer (IDataTransfer subclass) and the C API wrapper. diff --git a/onnxruntime/core/providers/webgpu/ep/README.md b/onnxruntime/core/providers/webgpu/ep/README.md index df01ed064a436..695ba2a6253af 100644 --- a/onnxruntime/core/providers/webgpu/ep/README.md +++ b/onnxruntime/core/providers/webgpu/ep/README.md @@ -24,6 +24,12 @@ To ensure both static library and dynamic library builds work, we need to make a ### Session streams +Concurrent use of independent Sessions is supported only by the plugin WebGPU EP. +Built-in WebGPU does not support concurrent use. Each Session has one command recording +timeline without an internal mutex. ORT serializes Run calls on the same Session; callers +must serialize explicit stream and I/O binding operations with that Session's Run. +Streamless allocation, tensor release, and Env transfers can run concurrently on independent tensors. + The WebGPU plugin implements `OrtEp::CreateSyncStreamForDevice`. Each stream references its owning EP's command state; kernels and stream-bearing data transfers use that same state. Graph Memcpy kernels access their owning EP directly. The generic single-tensor transfer @@ -38,27 +44,29 @@ Session allocators expose the existing `OrtAllocator::AllocOnStream` callback an that the stream belongs to the same Session. Allocations with a matching stream defer cached-buffer clears. Plugin kernel scratch tensors created through `CreateGPUTensor` use the kernel's explicit sync stream, so cached-buffer clears stay ordered with kernel work without submitting each scratch -allocation. Plain `Alloc` and null-stream allocations submit clears before returning, including -during Run. This keeps CPU-produced outputs bound to GPU ordered with fallback uploads without +allocation. Plain `Alloc` and null-stream allocations use a fresh encoder to submit cached-buffer +clears before returning, including during Run, without recording or flushing Session work. +This keeps CPU-produced outputs bound to GPU ordered with fallback uploads without additional core stream creation. The policy depends on the allocation's stream, not `IsRunActive()`. Uploads, readbacks, stream synchronization, dispatch batch limits, and Run/capture boundaries can still submit work. -Env and Session allocators reuse the shared `GpuBufferAllocator` implementation. Session getters -borrow their owning EP; Env getters retain a context and a separate command recording, so Env -allocation does not require a Session. The Env implementation is still created lazily by -`Factory::CreateAllocatorImpl`, submits cached-buffer clears before returning, and uses the -streamless `adapter::Allocator` ABI wrapper. It does not expose `OrtAllocator::AllocOnStream`, -even though the underlying implementation supports that method in plugin builds. - -Multiple threads may use one Session allocator, including while that Session or other Sessions run, -provided they operate on independent tensors. A dedicated small Session can also allocate inputs -for concurrent inference Sessions on the same WebGPU device/context. Keep the allocator Session -and allocator alive until their tensors are released, and synchronize writes before consuming a -shared tensor. Plain allocations can flush pending Session work, so correctness isolation does -not guarantee freedom from contention. - -Environment transfers with no stream use their private command state. GPU-to-GPU copies +Env and Session allocators reuse the `GpuBufferAllocator` implementation. Session getters borrow +their owning EP for stream allocation and deferred buffer release. Env allocators retain a context +without a Session recording and track allocation counts atomically. They are created lazily by +`Factory::CreateAllocatorImpl` and use the streamless `adapter::Allocator` ABI wrapper without +exposing `OrtAllocator::AllocOnStream`. Independent allocation encoders and per-call transfer +recordings are specific to plugin builds; built-in WebGPU retains its existing recording policy. + +Multiple threads may use an Env or Session allocator on independent tensors, including while the +Session runs. Keep a Session allocator and its owning Session alive until their tensors are released, +and synchronize writes before another Session consumes a shared tensor. BufferManager owns pending +releases for reusable buffers, grouped by recording and protected by its cache lock. Submission +returns only that recording's buffers to the pool; caches without reuse need no such tracking. +Recording starts, submission, and abandonment update this tracking under the same cache lock. +CommandRecordingState contains only command recording data. Context-shared program caches retain their locks. + +Environment transfers with no stream use command state local to each CopyTensors call. GPU-to-GPU copies without a stream submit before returning, but do not wait for GPU completion. A subsequent Session Run uses a different recording, so the transfer must submit its copy first; the shared device queue orders it before later Session work. @@ -72,11 +80,12 @@ general host-completion barrier. In particular, returning GPU-backed outputs fro not guarantee that GPU execution has completed. The implementation requires an ORT build with stream support. CPU I/O, graph-internal CPU/GPU -copies, mixed feed copies, CPU outputs bound to GPU, concurrent Sessions, and same-Session and -dedicated-Session allocator concurrency are covered by the tests in +copies, mixed feed copies, CPU outputs bound to GPU, concurrent independent Sessions, serialized +same-Session Runs, same-Session and dedicated-Session allocator concurrency, and shared Env allocation/transfers +are covered by the tests in `onnxruntime/test/providers/webgpu/plugin`, built into `onnxruntime_provider_test`. -Concurrent graph capture, concurrent profiling, cross-device transfer, and arbitrary foreign -stream overrides are not established by these tests. Performance must be measured separately. +Concurrent graph capture and concurrent profiling are unsupported. Cross-device transfer and +arbitrary foreign stream overrides are not established by these tests. Performance must be measured separately. ### Missing parts diff --git a/onnxruntime/core/providers/webgpu/ep/allocator.cc b/onnxruntime/core/providers/webgpu/ep/allocator.cc index 6e1168554c0b9..92963ae27dcec 100644 --- a/onnxruntime/core/providers/webgpu/ep/allocator.cc +++ b/onnxruntime/core/providers/webgpu/ep/allocator.cc @@ -21,7 +21,7 @@ void* GpuBufferAllocator::AllocOnStream(size_t size, Stream* stream) { } ORT_ENFORCE(&ep::GetWebGpuStreamCommandState(reinterpret_cast(stream)) == &recording_getter_(), "WebGPU allocator and stream belong to different Sessions."); - return Allocate(size, false); + return Allocate(size, recording_getter_(), false); } namespace ep { diff --git a/onnxruntime/core/providers/webgpu/ep/allocator.h b/onnxruntime/core/providers/webgpu/ep/allocator.h index 107f2860c5e3e..d2ae375d66ca5 100644 --- a/onnxruntime/core/providers/webgpu/ep/allocator.h +++ b/onnxruntime/core/providers/webgpu/ep/allocator.h @@ -25,7 +25,7 @@ namespace ep { // | Internal use | Not used for EP kernel scratch | Run input/intermediate/output and scratch | // | Buffer manager | Context-shared BufferManager | Context-shared BufferManager | // | Capture route | Still context-shared default | Session-owned per-graph manager during Run | -// | Recording | Getter retains independent state | Getter borrows the owning EP's Recording() | +// | Recording | Per-call encoder; no Session recording | AllocOnStream/Free borrow EP Recording() | // | Lifetime | Getter retains Context; no Session | EP must outlive allocator use/tensor frees | // | Alloc | Submit cached clear before return | Submit cached clear, even during Run | // | AllocOnStream | Not exposed by the C API wrapper | Matching Session stream: defer cached clear | diff --git a/onnxruntime/core/providers/webgpu/ep/factory.cc b/onnxruntime/core/providers/webgpu/ep/factory.cc index 945dfc920f3db..5fd7e68853f2a 100644 --- a/onnxruntime/core/providers/webgpu/ep/factory.cc +++ b/onnxruntime/core/providers/webgpu/ep/factory.cc @@ -221,8 +221,7 @@ OrtStatus* ORT_API_CALL Factory::CreateEpImpl( device_free, [webgpu_ep_ptr]() -> const webgpu::BufferManager& { return webgpu_ep_ptr->BufferManager(); }, [webgpu_ep_ptr]() -> webgpu::CommandRecordingState& { return webgpu_ep_ptr->Recording(); }, - false, - []() { return true; }); + false); Ep::Config webgpu_ep_config{ CPUAllocator::DefaultInstance(), // CPU allocator device_alloc, // also retained by the EP adapter as the kernel temp-space allocator @@ -258,20 +257,17 @@ OrtStatus* ORT_API_CALL Factory::CreateAllocatorImpl( "Unsupported memory info for shared allocator."); } - // Env getters retain the context and an independent recording, not a Session EP. - // Keep the streamless ABI wrapper even though the shared implementation supports AllocOnStream. + // Env allocations submit their own clears and never borrow a Session recording. *allocator = new onnxruntime::ep::adapter::Allocator( memory_info, [](const OrtMemoryInfo&) -> AllocatorPtr { auto context = std::shared_ptr( &WebGpuContextFactory::DefaultContext(), [](WebGpuContext*) { WebGpuContextFactory::ReleaseContext(0); }); - auto recording = std::make_shared(); return std::make_shared( [context = std::move(context)]() -> const BufferManager& { return context->BufferManager(); }, - [recording = std::move(recording)]() -> CommandRecordingState& { return *recording; }, - false, - []() { return true; }); + std::function{}, + false); }); return nullptr; EXCEPTION_TO_RETURNED_STATUS_END diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc index 71b9d15a4adfd..f7831672483b3 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc +++ b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc @@ -3,8 +3,6 @@ #include "sync_stream.h" -#include - #include "core/providers/webgpu/data_transfer.h" #include "core/providers/webgpu/webgpu_context.h" #include "core/providers/webgpu/webgpu_execution_provider.h" @@ -34,7 +32,6 @@ struct WebGpuSyncStream final : OrtSyncStreamImpl { EXCEPTION_TO_RETURNED_STATUS_BEGIN auto& ep = static_cast(stream)->ep_; auto& context = WebGpuContextFactory::GetContext(ep.GetDeviceId()); - std::lock_guard lock{ep.Recording().mutex}; // Submit before streamless readbacks (e.g., node dumps) use a different recording. // Readbacks wait for completion; GPU consumers rely on ordering on the shared queue. ORT_THROW_IF_ERROR(context.Flush(ep.BufferManager(), ep.Recording())); diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index b21f618228577..79c88af4d7e7b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -436,7 +436,6 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra const webgpu::BufferManager& buffer_mgr = ComputeContextBase::BufferManagerAccessor::Get(context); CommandRecordingState& recording = ComputeContextBase::BufferManagerAccessor::GetRecording(context); - std::lock_guard lock{recording.mutex}; // validate inputs and outputs are on WebGPU buffers if (ValidationMode() >= ValidationMode::Basic) { @@ -741,10 +740,10 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra command.bind_group = CreateBindGroup(bind_buffers, bind_buffers_segments, *bind_group_layout, program.Name()); command.pending_build = std::move(pending_build); - recording.has_unsubmitted_work = true; + buffer_mgr.BeginRecording(recording); if (uniform_buffer) { // The bind group owns a reference now, so return the allocator's reference immediately. - buffer_mgr.Release(recording, uniform_buffer); + buffer_mgr.Release(uniform_buffer, &recording); } command.dispatch_group = {x, y, z}; if (program.IndirectDispatchTensor() != nullptr) { @@ -1086,7 +1085,6 @@ Status WebGpuContext::PopErrorScope() { Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; Status status = EncodeDeferredDispatches(recording); if (!recording.command_encoder) { return status; @@ -1131,12 +1129,9 @@ Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, } auto command_buffer = recording.command_encoder.Finish(); device_queue_.Submit(1, &command_buffer); - if (recording.graph_capture_state != GraphCaptureState::Replaying) { - buffer_mgr.RefreshPendingBuffers(recording); - } + buffer_mgr.RefreshPendingBuffers(recording); recording.command_encoder = nullptr; recording.num_pending_dispatches = 0; - recording.has_unsubmitted_work = false; return status; } @@ -1209,7 +1204,6 @@ void WebGpuContext::DispatchCommand(const webgpu::CapturedCommandInfo& command, void WebGpuContext::CaptureBegin(std::vector* captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "CaptureBegin with external storage"; // Flush any pending commands before we change the status ORT_THROW_IF_ERROR(Flush(buffer_manager, recording)); @@ -1226,7 +1220,6 @@ void WebGpuContext::CaptureBegin(std::vector* captu void WebGpuContext::Replay(const std::vector& captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "Replay with external storage"; recording.graph_capture_state = GraphCaptureState::Replaying; // Replay all captured commands from the provided vector @@ -1244,6 +1237,7 @@ void WebGpuContext::Replay(const std::vector& captu recording.pending_kernels.emplace_back(*command.pending_kernel_info); } + buffer_manager.BeginRecording(recording); DispatchCommand(command, recording); if (recording.num_pending_dispatches >= max_num_pending_dispatches_) { ORT_THROW_IF_ERROR(Flush(buffer_manager, recording)); @@ -1257,7 +1251,6 @@ void WebGpuContext::Replay(const std::vector& captu } void WebGpuContext::CaptureEnd(CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "CaptureEnd"; recording.graph_capture_state = GraphCaptureState::Default; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 90183c0e14f10..df20d6619d56a 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -108,17 +108,16 @@ struct CapturedCommandInfo { std::optional pending_kernel_info; }; -// State for one session's command recording timeline. WebGpuContext is shared across sessions, -// but Dawn command encoders and deferred dispatch windows must not be. +// A command recording timeline has one caller at a time. ORT serializes Run on a Session; +// plugin streamless operations use independent encoders. BufferManager synchronizes buffer +// releases from concurrent allocator callers separately from command recording. struct CommandRecordingState { - std::recursive_mutex mutex; + CommandRecordingState() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CommandRecordingState); + wgpu::CommandEncoder command_encoder; wgpu::ComputePassEncoder compute_pass_encoder; uint32_t num_pending_dispatches = 0; - // TODO: Make BufferManager defer reuse of buffers belonging to unsubmitted batches, including - // across Sessions, so has_unsubmitted_work and pending_buffers can be removed from this state. - bool has_unsubmitted_work = false; - std::vector pending_buffers; std::vector deferred_dispatches; std::vector pending_kernels; GraphCaptureState graph_capture_state{GraphCaptureState::Default}; @@ -257,14 +256,6 @@ class WebGpuContext final { bool DeviceHasFeature(wgpu::FeatureName feature) const { return device_features_.contains(feature); } const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return subgroup_matrix_configs_; } - const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) { - if (!recording.command_encoder) { - recording.command_encoder = device_.CreateCommandEncoder(); - recording.has_unsubmitted_work = true; - } - return recording.command_encoder; - } - const wgpu::ComputePassEncoder& GetComputePassEncoder(CommandRecordingState& recording) { if (!recording.compute_pass_encoder) { auto& command_encoder = GetCommandEncoder(recording); @@ -397,6 +388,13 @@ class WebGpuContext final { void Initialize(const WebGpuContextConfig& config); + const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) { + if (!recording.command_encoder) { + recording.command_encoder = device_.CreateCommandEncoder(); + } + return recording.command_encoder; + } + wgpu::BindGroup CreateBindGroup(const std::vector& bind_buffers, const std::vector& bind_buffers_segments, const wgpu::BindGroupLayout& bind_group_layout, diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 57df6dd0417fa..ea7cf7a71d9fd 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -816,6 +816,11 @@ WebGpuExecutionProvider::~WebGpuExecutionProvider() { prepack_allocator_.reset(); session_buffer_pool_.reset(); + if (context_.Device()) { + // A failed Run may leave an unsubmitted recording in the context-shared pools. + context_.BufferManager().DiscardPendingBuffers(*recording_); + context_.InitializerBufferManager().DiscardPendingBuffers(*recording_); + } recording_.reset(); #if defined(ENABLE_PIX_FOR_WEBGPU_EP) pix_frame_generator_.reset(); @@ -996,7 +1001,7 @@ Status WebGpuExecutionProvider::ReleaseCapturedGraph(int graph_annotation_id) { } webgpu::BufferManager& WebGpuExecutionProvider::BufferManager() const { - if (graph_buffer_mgr_active_) { + if (IsGraphCaptureEnabled() && graph_buffer_mgr_active_) { auto it = per_graph_buffer_mgrs_.find(current_graph_annotation_id_); if (it != per_graph_buffer_mgrs_.end()) { return *it->second; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index 823d3ee416b4c..507b426cb1ec1 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -99,8 +99,7 @@ class WebGpuExecutionProvider : public IExecutionProvider { FusionStyle GetFusionStyle() const override { return FusionStyle::FilteredGraphViewer; } - // WebGPU EP disallow concurrent run because actual implementation (eg. WebGPU backend) relies on global states to - // work, and concurrent run with async function may mess up the states and cause undefined behavior. + // A Session owns one command recording timeline, so Run calls on that Session must be serialized. bool ConcurrentRunSupported() const override { return false; } std::vector CreatePreferredAllocators() override; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index b68392874aeec..7a28b197bd681 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -410,8 +410,7 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl { WebGpuDataTransferImpl(const OrtApi& ort_api_in, int context_id) : ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()}, - recording_{}, - data_transfer_{nullptr}, + context_{nullptr}, context_id_{context_id}, init_mutex_{} { ort_version_supported = ORT_API_VERSION; @@ -493,18 +492,24 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl { { std::lock_guard lock(impl.init_mutex_); - if (impl.data_transfer_ == nullptr) { + if (impl.context_ == nullptr) { // Always create a new context with context_id 0 if (impl.context_id_ != 0) { return OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, "Shared data transfer can only be created for the default device (0)."); } - auto& context = WebGpuContextFactory::DefaultContext(); - - impl.data_transfer_ = std::make_unique(context.BufferManager(), impl.recording_); + impl.context_ = &WebGpuContextFactory::DefaultContext(); } } +#if defined(ORT_USE_EP_API_ADAPTERS) + // Plugin streamless calls may overlap. Each call submits its own recording before returning; + // explicit streams use their Session's recording instead. + CommandRecordingState recording; +#else + auto& recording = impl.recording_; +#endif + DataTransferImpl data_transfer{impl.context_->BufferManager(), recording}; for (size_t idx = 0; idx < num_tensors; ++idx) { #if defined(ORT_USE_EP_API_ADAPTERS) Ort::ConstValue src_value{src_tensors[idx]}; @@ -526,22 +531,22 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl { bool dst_is_gpu = dst_tensor.Location().device.Type() == OrtDevice::GPU; #endif #if defined(ORT_USE_EP_API_ADAPTERS) - auto status = streams != nullptr && streams[idx] != nullptr + const bool has_session_stream = streams != nullptr && streams[idx] != nullptr; + auto status = has_session_stream ? webgpu::ep::CopyTensorOnWebGpuStream(streams[idx], src_data, src_is_gpu, dst_data, dst_is_gpu, size) - : impl.data_transfer_->CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, size); + : data_transfer.CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, size); #else ORT_UNUSED_PARAMETER(streams); - auto status = impl.data_transfer_->CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, size); + constexpr bool has_session_stream = false; + auto status = data_transfer.CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, size); #endif if (!status.IsOK()) { return OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, status.ErrorMessage().c_str()); } - if (src_is_gpu && dst_is_gpu && (streams == nullptr || streams[idx] == nullptr)) { + if (src_is_gpu && dst_is_gpu && !has_session_stream) { // Env copies use a separate recording: a subsequent Session::Run cannot submit this copy. // Flush here so later Session work on the same queue is ordered after it, without a CPU wait. - auto& context = WebGpuContextFactory::GetContext(impl.context_id_); - std::lock_guard lock{impl.recording_.mutex}; - ORT_THROW_IF_ERROR(context.Flush(context.BufferManager(), impl.recording_)); + ORT_THROW_IF_ERROR(impl.context_->Flush(impl.context_->BufferManager(), recording)); } } return nullptr; @@ -554,7 +559,7 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl { bool data_transfer_initialized = false; { std::lock_guard lock(p_impl->init_mutex_); - data_transfer_initialized = (p_impl->data_transfer_ != nullptr); + data_transfer_initialized = (p_impl->context_ != nullptr); } delete p_impl; if (data_transfer_initialized) { @@ -564,10 +569,12 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl { const OrtApi& ort_api; const OrtEpApi& ep_api; +#if !defined(ORT_USE_EP_API_ADAPTERS) CommandRecordingState recording_; - std::unique_ptr data_transfer_; // Lazy-initialized - int context_id_; // Track which context we're using - std::mutex init_mutex_; // Protects lazy initialization +#endif + WebGpuContext* context_; // Lazily retained until ReleaseImpl. + int context_id_; // Track which context we're using + std::mutex init_mutex_; // Protects lazy initialization }; OrtDataTransferImpl* OrtWebGpuCreateDataTransfer(int context_id /* = 0 */) { diff --git a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc deleted file mode 100644 index 8ae14b9f0b0ab..0000000000000 --- a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc +++ /dev/null @@ -1,795 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -// Concurrency regression tests for the shared default WebGpuContext (context_id=0) used by -// multiple InferenceSessions on different threads. -// -// Background: -// InferenceSession only serializes a single session's Run via its own session_mutex_ (the -// WebGPU EP reports ConcurrentRunSupported() == false). It does NOT serialize across -// sessions. Multiple sessions with the default WebGPU provider share one WebGpuContext, so -// their Run / allocation / initializer-upload paths run concurrently and mutate the -// context's single command encoder (current_command_encoder_ / current_compute_pass_encoder_ -// / num_pending_dispatches_) AND the shared BufferManager cache maps. -// -// Before the fix this produced a data race and Dawn errors such as: -// "[CommandEncoder] is already finished. While encoding CopyBufferToBuffer(...)" -// "WebGPU validation failed. Command encoding already finished." -// a corrupted buffer cache -> "[Device] is lost", or - worst of all - a silently wrong result -// when a buffer was recycled before the work referencing it had been submitted. -// -// Context-level BufferManagers protect only their shared buffer-cache containers. Command -// recording remains per session and is passed independently to context operations. Graph capture -// retains a separate manager per graph so captured resources remain isolated. -// -// The tests cover several distinct multithreaded shapes: -// A. one session, run() concurrently from many threads -// B. many threads, each with its own pre-created session, running concurrently -// C. mixed: some threads create+Initialize+run new sessions while others run existing ones -// D. churn: many threads each repeatedly create + run + destroy their own session -// E. a cold and a warm session remain correct while running concurrently - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "gtest/gtest.h" - -#include "core/graph/onnx_protobuf.h" -#include "core/graph/model.h" -#include "core/platform/env.h" -#include "core/providers/webgpu/allocator.h" -#include "core/providers/webgpu/data_transfer.h" -#include "core/providers/webgpu/program_manager.h" -#include "core/providers/webgpu/webgpu_context.h" -#include "core/providers/webgpu/webgpu_external_header.h" -#include "core/providers/webgpu/webgpu_provider_options.h" -#include "core/session/inference_session.h" - -#include "test/test_environment.h" -#include "test/unittest_util/framework_test_utils.h" -#include "test/util/include/asserts.h" -#include "test/util/include/default_providers.h" -#include "test/util/include/temp_dir.h" - -namespace onnxruntime { -namespace test { - -namespace { - -// Builds an in-memory model: Y = ((X + W0) + W1) + ... + W(chain_len-1) -// where each Wi is a constant initializer of `num_elements` floats. The constant initializers -// exercise the BufferManager::Upload (CopyBufferToBuffer) path during Initialize, and each Add -// node produces a compute dispatch during Run. -void BuildAddChainModel(int chain_len, int64_t num_elements, std::string& model_bytes) { - const std::unordered_map domain_to_version{{"", 13}}; - Model model("webgpu_concurrent_ctx", false, ModelMetaData(), PathString(), - IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, - std::vector(), - DefaultLoggingManager().DefaultLogger()); - Graph& graph = model.MainGraph(); - - ONNX_NAMESPACE::TypeProto float_1d; - float_1d.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - float_1d.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(num_elements); - - const std::vector weight_values(static_cast(num_elements), 0.5f); - - NodeArg* prev = &graph.GetOrCreateNodeArg("X", &float_1d); - for (int i = 0; i < chain_len; ++i) { - const std::string w_name = "W" + std::to_string(i); - ONNX_NAMESPACE::TensorProto w_tensor; - w_tensor.set_name(w_name); - w_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - w_tensor.add_dims(num_elements); - w_tensor.set_raw_data(weight_values.data(), weight_values.size() * sizeof(float)); - graph.AddInitializedTensor(w_tensor); - - NodeArg* w_arg = &graph.GetOrCreateNodeArg(w_name, &float_1d); - const std::string out_name = (i == chain_len - 1) ? "Y" : ("H" + std::to_string(i)); - NodeArg* out_arg = &graph.GetOrCreateNodeArg(out_name, &float_1d); - std::vector inputs{prev, w_arg}; - std::vector outputs{out_arg}; - graph.AddNode("add" + std::to_string(i), "Add", "", inputs, outputs); - prev = out_arg; - } - - graph.SetOutputs(std::vector{prev}); - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_TRUE(model.ToProto().SerializeToString(&model_bytes)); -} - -// The distinct unary op types used by BuildUnaryFanOutModel. Each op type maps to its own WebGPU -// program, so a session using this model has to compile one pipeline per entry. -constexpr const char* kUnaryOps[] = { - "Abs", "Neg", "Floor", "Ceil", "Reciprocal", "Sqrt", "Exp", "Erf", "Sigmoid", - "Sin", "Cos", "Tan", "Atan", "Sinh", "Cosh", "Tanh", "HardSigmoid", "HardSwish"}; - -// Builds an in-memory model that fans one input out to every op in kUnaryOps: -// Y0 = Abs(X), Y1 = Neg(X), ... -// -// Distinct op types mean distinct programs, so the first Run of a session built from this model -// compiles std::size(kUnaryOps) pipelines. A fan-out rather than a chain keeps every op inside -// its valid input domain no matter how many are used. -void BuildUnaryFanOutModel(int64_t num_elements, std::string& model_bytes) { - const std::unordered_map domain_to_version{{"", 14}}; - Model model("webgpu_concurrent_ctx_cold", false, ModelMetaData(), PathString(), - IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, - std::vector(), - DefaultLoggingManager().DefaultLogger()); - Graph& graph = model.MainGraph(); - - ONNX_NAMESPACE::TypeProto float_1d; - float_1d.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - float_1d.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(num_elements); - - NodeArg* x_arg = &graph.GetOrCreateNodeArg("X", &float_1d); - std::vector outputs; - outputs.reserve(std::size(kUnaryOps)); - - for (size_t i = 0; i < std::size(kUnaryOps); ++i) { - NodeArg* out_arg = &graph.GetOrCreateNodeArg("Y" + std::to_string(i), &float_1d); - std::vector node_inputs{x_arg}; - std::vector node_outputs{out_arg}; - graph.AddNode("op" + std::to_string(i), kUnaryOps[i], "", node_inputs, node_outputs); - outputs.push_back(out_arg); - } - - graph.SetOutputs(outputs); - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_TRUE(model.ToProto().SerializeToString(&model_bytes)); -} - -// Thread-safe first-error recorder that also acts as a stop flag for the worker loops. -class ErrorSink { - public: - void Record(const std::string& message) { - bool expected = false; - if (failed_.compare_exchange_strong(expected, true)) { - std::lock_guard lock(mutex_); - first_error_ = message; - } - } - - bool Failed() const { return failed_.load(); } - - std::string FirstError() const { - std::lock_guard lock(mutex_); - return first_error_; - } - - private: - std::atomic failed_{false}; - mutable std::mutex mutex_; - std::string first_error_; -}; - -} // namespace - -TEST(WebGpuConcurrentContextTestStandalone, ShaderDumpWritesRemainComplete) { - TemporaryDirectory temp_dir{ORT_TSTR("webgpu_shader_dump_test")}; - const auto dump_path = std::filesystem::path(temp_dir.Path()) / ORT_TSTR("shaders.txt"); - constexpr int kThreads = 8; - constexpr int kRecordsPerThread = 100; - constexpr size_t kPayloadSize = 4096; - - { - auto dump_shader = webgpu::detail::CreateShaderDumpFunction(ToUTF8String(dump_path.native())); - std::barrier start{kThreads}; - std::vector threads; - threads.reserve(kThreads); - for (int thread_id = 0; thread_id < kThreads; ++thread_id) { - threads.emplace_back([&, thread_id] { - start.arrive_and_wait(); - for (int record_id = 0; record_id < kRecordsPerThread; ++record_id) { - const std::string marker = "record-" + std::to_string(thread_id) + "-" + std::to_string(record_id); - dump_shader(marker + "-start\n" + std::string(kPayloadSize, static_cast('a' + thread_id)) + - "\n" + marker + "-end"); - } - }); - } - for (auto& thread : threads) { - thread.join(); - } - } - - std::ifstream dump_file(dump_path); - ASSERT_TRUE(dump_file.is_open()); - const std::string contents{std::istreambuf_iterator{dump_file}, std::istreambuf_iterator{}}; - for (int thread_id = 0; thread_id < kThreads; ++thread_id) { - for (int record_id = 0; record_id < kRecordsPerThread; ++record_id) { - const std::string marker = "record-" + std::to_string(thread_id) + "-" + std::to_string(record_id); - const std::string record = marker + "-start\n" + - std::string(kPayloadSize, static_cast('a' + thread_id)) + - "\n" + marker + "-end\n"; - const size_t position = contents.find(record); - ASSERT_NE(position, std::string::npos) << "Missing or corrupted shader dump record " << marker; - EXPECT_EQ(contents.find(record, position + 1), std::string::npos) - << "Duplicate shader dump record " << marker; - } - } -} - -// Fixture: builds the shared model once and provides session/feed factories. A long-lived -// keepalive session pins the shared WebGPU context (ref-count > 0) for the whole test so that -// churn/destroy in one thread never tears the context down under the others. -class WebGpuConcurrentContextTest : public ::testing::Test { - protected: - static constexpr int64_t kNumElements = 256 * 1024; // 1 MB per initializer - static constexpr int kChainLen = 8; // 8 uploads + 8 dispatches per session - static constexpr float kExpected = 1.0f + 0.5f * kChainLen; - - void SetUp() override { - if (DefaultWebGpuExecutionProvider() == nullptr) { - GTEST_SKIP() << "WebGPU execution provider is not available."; - } - ASSERT_NO_FATAL_FAILURE(BuildAddChainModel(kChainLen, kNumElements, model_bytes_)); - keepalive_ = MakeSession(); - } - - std::unique_ptr MakeSession() { - SessionOptions so; - so.session_logid = "webgpu_concurrent_ctx"; - auto session = std::make_unique(so, GetEnvironment()); - ORT_THROW_IF_ERROR(session->RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); - ORT_THROW_IF_ERROR(session->Load(model_bytes_.data(), static_cast(model_bytes_.size()))); - ORT_THROW_IF_ERROR(session->Initialize()); - return session; - } - - NameMLValMap MakeFeeds() const { - std::vector x_values(static_cast(kNumElements), 1.0f); - OrtValue x_value; - CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], - std::vector{kNumElements}, x_values, &x_value); - return NameMLValMap{{"X", x_value}}; - } - - // Runs one inference and validates the numerical result. Records into `sink` on failure. - void RunOnce(InferenceSession& session, ErrorSink& sink, const std::string& tag) { - std::vector output_names{"Y"}; - std::vector fetches; - Status s = session.Run(RunOptions{}, MakeFeeds(), output_names, &fetches); - if (!s.IsOK()) { - sink.Record(tag + " Run failed: " + s.ErrorMessage()); - return; - } - const Tensor& out = fetches[0].Get(); - const float* data = out.Data(); - const int64_t n = out.Shape().Size(); - for (int64_t i = 0; i < n; i += (n / 8) + 1) { - if (std::abs(data[i] - kExpected) > 1e-3f) { - sink.Record(tag + " wrong output: " + std::to_string(data[i])); - return; - } - } - } - - // Repeatedly runs an existing session until `iters` reached or a failure is recorded. - void RunLoop(InferenceSession& session, int iters, ErrorSink& sink, const std::string& tag) { - for (int i = 0; i < iters && !sink.Failed(); ++i) { - try { - RunOnce(session, sink, tag); - } catch (const std::exception& e) { - sink.Record(tag + " threw: " + e.what()); - return; - } - } - } - - static void JoinAll(std::vector& threads) { - for (auto& t : threads) { - t.join(); - } - } - - std::string model_bytes_; - std::unique_ptr keepalive_; -}; - -// Case A: one session, run() concurrently from many threads. InferenceSession serializes these -// via session_mutex_, but this still must never crash or deadlock on the shared context. -TEST_F(WebGpuConcurrentContextTest, SingleSessionMultiThreadRun) { - constexpr int kThreads = 4; - constexpr int kIters = 30; - - auto session = MakeSession(); - ErrorSink sink; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { RunLoop(*session, kIters, sink, "A.run" + std::to_string(t)); }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case B: many threads, each with its own pre-created session. All sessions share the default -// context, so their Run paths execute concurrently against it. Before the fix they also shared a -// single command encoder, which is what this case exposed. -TEST_F(WebGpuConcurrentContextTest, PerThreadSessionRun) { - constexpr int kThreads = 4; - constexpr int kIters = 30; - - std::vector> sessions(kThreads); - for (int t = 0; t < kThreads; ++t) { - sessions[t] = MakeSession(); - } - - ErrorSink sink; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { RunLoop(*sessions[t], kIters, sink, "B.sess" + std::to_string(t)); }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case C: mixed. Runner threads dispatch on pre-initialized sessions while builder threads keep -// creating + Initializing (initializer-upload path) + running fresh sessions. This is the -// closest shape to the original WebNN crash (dispatch racing initializer upload). -TEST_F(WebGpuConcurrentContextTest, MixedCreateAndRun) { - constexpr int kRunners = 3; - constexpr int kBuilders = 3; - constexpr int kIters = 30; - - std::vector> runner_sessions(kRunners); - for (int t = 0; t < kRunners; ++t) { - runner_sessions[t] = MakeSession(); - } - - ErrorSink sink; - std::vector threads; - for (int t = 0; t < kRunners; ++t) { - threads.emplace_back([&, t]() { RunLoop(*runner_sessions[t], kIters, sink, "C.runner" + std::to_string(t)); }); - } - for (int t = 0; t < kBuilders; ++t) { - threads.emplace_back([&, t]() { - const std::string tag = "C.builder" + std::to_string(t); - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - try { - auto session = MakeSession(); // Initialize -> initializer upload - RunOnce(*session, sink, tag); - } catch (const std::exception& e) { - sink.Record(tag + " threw: " + e.what()); - return; - } - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case D: churn. Many threads each repeatedly create + run + destroy their own session, -// exercising concurrent allocation and release against the shared context-level buffer manager. -TEST_F(WebGpuConcurrentContextTest, ChurnCreateRunDestroy) { - constexpr int kThreads = 4; - constexpr int kIters = 15; - - ErrorSink sink; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { - const std::string tag = "D.churn" + std::to_string(t); - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - try { - auto session = MakeSession(); - RunOnce(*session, sink, tag); - session.reset(); // drop session -> release its resources while others run - } catch (const std::exception& e) { - sink.Record(tag + " threw: " + e.what()); - return; - } - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case E: a session warming up pipelines concurrently with a warm session must remain correct. -// -// A builder thread creates a session over a model with many distinct op types and runs it once, -// which compiles one pipeline per op. Meanwhile a runner thread keeps dispatching on an already -// warm session. -// -// Timing is reported to make lock contention visible without imposing a machine-dependent -// performance threshold on this correctness test. -TEST_F(WebGpuConcurrentContextTest, ColdAndWarmSessionsRunConcurrently) { - using Clock = std::chrono::steady_clock; - using Ms = std::chrono::duration; - - std::string cold_model_bytes; - ASSERT_NO_FATAL_FAILURE(BuildUnaryFanOutModel(kNumElements, cold_model_bytes)); - - // Warm session: every pipeline it needs is already compiled and cached on the context. - auto warm_session = MakeSession(); - ErrorSink sink; - RunOnce(*warm_session, sink, "E.warmup"); - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); - - // Baseline latency with no other session active. - constexpr int kBaselineIters = 20; - std::vector baseline_ms; - baseline_ms.reserve(kBaselineIters); - for (int i = 0; i < kBaselineIters; ++i) { - const auto start = Clock::now(); - RunOnce(*warm_session, sink, "E.baseline"); - baseline_ms.push_back(Ms(Clock::now() - start).count()); - } - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); - std::sort(baseline_ms.begin(), baseline_ms.end()); - const double baseline_median = baseline_ms[baseline_ms.size() / 2]; - - std::atomic builder_done{false}; - std::vector contended_ms; - double builder_ms = 0.0; - - std::thread builder([&]() { - const auto start = Clock::now(); - try { - SessionOptions so; - so.session_logid = "webgpu_concurrent_ctx_cold"; - InferenceSession cold_session(so, GetEnvironment()); - ORT_THROW_IF_ERROR(cold_session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); - ORT_THROW_IF_ERROR(cold_session.Load(cold_model_bytes.data(), static_cast(cold_model_bytes.size()))); - ORT_THROW_IF_ERROR(cold_session.Initialize()); - - std::vector output_names; - for (size_t i = 0; i < std::size(kUnaryOps); ++i) { - output_names.push_back("Y" + std::to_string(i)); - } - std::vector fetches; - // This Run is what triggers compilation of one pipeline per op type. - ORT_THROW_IF_ERROR(cold_session.Run(RunOptions{}, MakeFeeds(), output_names, &fetches)); - } catch (const std::exception& e) { - sink.Record(std::string("E.builder threw: ") + e.what()); - } - builder_ms = Ms(Clock::now() - start).count(); - builder_done.store(true); - }); - - // Always take at least one measurement so the loop cannot produce an empty sample set if the - // builder happens to finish first. - const auto contended_start = Clock::now(); - do { - const auto start = Clock::now(); - RunOnce(*warm_session, sink, "E.contended"); - contended_ms.push_back(Ms(Clock::now() - start).count()); - } while (!builder_done.load() && !sink.Failed()); - const double contended_window_ms = Ms(Clock::now() - contended_start).count(); - builder.join(); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); - - const double contended_max = *std::max_element(contended_ms.begin(), contended_ms.end()); - std::sort(contended_ms.begin(), contended_ms.end()); - const double contended_median = contended_ms[contended_ms.size() / 2]; - - // Fraction of the contended window during which the warm session kept doing useful work, using - // its uncontended latency as the reference. 1.0 means the cold session cost it nothing. - const double throughput_efficiency = - (static_cast(contended_ms.size()) * baseline_median) / contended_window_ms; - - std::cout << "[ WebGPU ] cold-session warmup: " << builder_ms << " ms for " - << std::size(kUnaryOps) << " pipelines\n" - << "[ WebGPU ] warm-session latency: baseline median " << baseline_median - << " ms, contended median " << contended_median - << " ms, contended max " << contended_max - << " ms over " << contended_ms.size() << " runs\n" - << "[ WebGPU ] warm-session throughput efficiency while cold session warmed up: " - << throughput_efficiency << std::endl; -} - -// Case F: the public session allocator wraps this same internal allocator. Allocations made -// concurrently with Run must not race the session's command recording state. -TEST_F(WebGpuConcurrentContextTest, SessionAllocatorAndRunConcurrently) { - constexpr int kIters = 40; - auto session = MakeSession(); - auto allocator = session->GetAllocator(OrtMemoryInfo(WEBGPU_BUFFER, - OrtAllocatorType::OrtDeviceAllocator, - webgpu::WebGpuDevice, - OrtMemTypeDefault)); - ASSERT_NE(allocator, nullptr); - - ErrorSink sink; - std::barrier start{2}; - std::thread runner([&]() { - start.arrive_and_wait(); - RunLoop(*session, kIters, sink, "F.run"); - }); - std::thread allocator_user([&]() { - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - Tensor tensor(DataTypeImpl::GetType(), TensorShape{kNumElements}, allocator); - ASSERT_NE(tensor.MutableDataRaw(), nullptr); - } - } catch (const std::exception& e) { - sink.Record(std::string("F.allocator threw: ") + e.what()); - } - }); - runner.join(); - allocator_user.join(); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case G: exercise the plugin Env allocator's shared implementation and getter ownership in a -// native build. Its private command state must remain thread-safe across callers. -TEST_F(WebGpuConcurrentContextTest, SharedAllocatorMultiThreadCreateTensor) { - constexpr int kThreads = 4; - constexpr int kIters = 60; - auto& context = webgpu::WebGpuContextFactory::GetContext(0); - auto context_ref = std::shared_ptr(&context, [](webgpu::WebGpuContext*) {}); - auto recording = std::make_shared(); - std::weak_ptr recording_lifetime = recording; - auto allocator = std::make_shared( - [context_ref = std::move(context_ref)]() -> const webgpu::BufferManager& { - return context_ref->BufferManager(); - }, - [recording = std::move(recording)]() -> webgpu::CommandRecordingState& { return *recording; }, - false, - []() { return true; }); - - ErrorSink sink; - std::barrier start{kThreads}; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - Tensor tensor(DataTypeImpl::GetType(), TensorShape{1024 + t * 16}, allocator); - ASSERT_NE(tensor.MutableDataRaw(), nullptr); - } - } catch (const std::exception& e) { - sink.Record("G.thread" + std::to_string(t) + " threw: " + e.what()); - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); - EXPECT_FALSE(recording_lifetime.expired()); - AllocatorStats stats{}; - allocator->GetStats(&stats); - EXPECT_EQ(stats.num_allocs, 0); - allocator.reset(); - EXPECT_TRUE(recording_lifetime.expired()); -} - -// Case H: OrtEnv owns one data-transfer implementation per EP factory. Concurrent CopyTensors -// calls must not encode and flush through the same recording state simultaneously. -TEST_F(WebGpuConcurrentContextTest, SharedDataTransferMultiThreadCopy) { - constexpr int kThreads = 4; - constexpr int kIters = 30; - constexpr size_t kElements = 4096; - auto& context = webgpu::WebGpuContextFactory::GetContext(0); - webgpu::CommandRecordingState recording; - webgpu::DataTransferImpl data_transfer(context.BufferManager(), recording); - - std::array gpu_buffers; - for (auto& buffer : gpu_buffers) { - wgpu::BufferDescriptor desc{}; - desc.size = kElements * sizeof(float); - desc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; - buffer = context.Device().CreateBuffer(&desc); - } - - ErrorSink sink; - std::barrier start{kThreads}; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { - std::vector input(kElements, static_cast(t + 1)); - std::vector output(kElements); - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - ORT_THROW_IF_ERROR(data_transfer.CopyTensor(input.data(), false, gpu_buffers[t].Get(), true, - input.size() * sizeof(float))); - ORT_THROW_IF_ERROR(data_transfer.CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, - output.size() * sizeof(float))); - if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { - sink.Record("H.thread" + std::to_string(t) + " copied incorrect data"); - } - } - } catch (const std::exception& e) { - sink.Record("H.thread" + std::to_string(t) + " threw: " + e.what()); - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case I: separate data-transfer objects share the context-level BufferManager but own distinct -// command recording timelines. Their command encoders and pending buffers must remain isolated. -TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { - constexpr int kThreads = 4; - constexpr int kIters = 30; - constexpr size_t kElements = 4096; - auto& context = webgpu::WebGpuContextFactory::GetContext(0); - - std::array recordings; - std::array, kThreads> data_transfers; - std::array gpu_buffers; - for (int t = 0; t < kThreads; ++t) { - data_transfers[t] = std::make_unique(context.BufferManager(), recordings[t]); - wgpu::BufferDescriptor desc{}; - desc.size = kElements * sizeof(float); - desc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; - gpu_buffers[t] = context.Device().CreateBuffer(&desc); - } - - ErrorSink sink; - std::barrier start{kThreads}; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { - std::vector input(kElements, static_cast(t + 1)); - std::vector output(kElements); - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - ORT_THROW_IF_ERROR(data_transfers[t]->CopyTensor(input.data(), false, gpu_buffers[t].Get(), true, - input.size() * sizeof(float))); - ORT_THROW_IF_ERROR(data_transfers[t]->CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, - output.size() * sizeof(float))); - if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { - sink.Record("I.thread" + std::to_string(t) + " copied incorrect data"); - } - } - } catch (const std::exception& e) { - sink.Record("I.thread" + std::to_string(t) + " threw: " + e.what()); - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// DIAGNOSTIC (disabled by default): keeps N sessions over an identical model alive and runs them -// round-robin on one thread, then holds so an external sampler can read steady-state GPU memory. -// This is the shape that distinguishes a context-wide buffer pool (steady state ~= one session's -// high-water mark) from per-session pools (~= N times that). -// -// onnxruntime_provider_test.exe --gtest_also_run_disabled_tests \ -// --gtest_filter=WebGpuPoolMemory.DISABLED_MultiSessionSameShape -// -// Tunable via environment variables: ORT_DIAG_SESSIONS, ORT_DIAG_MB, ORT_DIAG_ITERS, -// ORT_DIAG_HOLD_MS, ORT_DIAG_VARY (1 = give each session a different tensor size), -// ORT_DIAG_CACHE (storage buffer cache mode; defaults to "bucket", the product default - -// note DefaultWebGpuExecutionProvider() disables the storage cache, which would hide the -// very behavior this test measures). -TEST(WebGpuPoolMemory, DISABLED_MultiSessionSameShape) { - if (DefaultWebGpuExecutionProvider() == nullptr) { - GTEST_SKIP() << "WebGPU execution provider is not available."; - } - - auto env_int = [](const char* name, int fallback) { - const std::string value = Env::Default().GetEnvironmentVar(name); - return value.empty() ? fallback : std::stoi(value); - }; - - const int num_sessions = env_int("ORT_DIAG_SESSIONS", 4); - const int mb = env_int("ORT_DIAG_MB", 64); - const int iters = env_int("ORT_DIAG_ITERS", 10); - const int hold_ms = env_int("ORT_DIAG_HOLD_MS", 5000); - const bool vary = env_int("ORT_DIAG_VARY", 0) != 0; - std::string cache_mode = Env::Default().GetEnvironmentVar("ORT_DIAG_CACHE"); - if (cache_mode.empty()) { - cache_mode = webgpu::options::kBufferCacheMode_Bucket; - } - constexpr int kChainLen = 8; - - auto make_ep = [&cache_mode]() { - ConfigOptions config_options; - ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kStorageBufferCacheMode, - cache_mode.c_str())); - return WebGpuExecutionProviderWithOptions(config_options); - }; - std::cout << "DIAG_POOL storage_cache=" << cache_mode << "\n"; - - std::vector> sessions; - std::vector feeds; - std::vector output_names{"Y"}; - - // Scalar weights keep GPU memory dominated by intermediate tensors rather than initializers. - auto build_scalar_weight_chain = [](int chain_len, int64_t num_elements, std::string& bytes) { - const std::unordered_map domain_to_version{{"", 13}}; - Model model("webgpu_pool_memory", false, ModelMetaData(), PathString(), - IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, - std::vector(), - DefaultLoggingManager().DefaultLogger()); - Graph& graph = model.MainGraph(); - - ONNX_NAMESPACE::TypeProto float_1d; - float_1d.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - float_1d.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(num_elements); - - ONNX_NAMESPACE::TypeProto float_scalar; - float_scalar.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - float_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - - NodeArg* prev = &graph.GetOrCreateNodeArg("X", &float_1d); - for (int i = 0; i < chain_len; ++i) { - const std::string w_name = "w" + std::to_string(i); - ONNX_NAMESPACE::TensorProto w_tensor; - w_tensor.set_name(w_name); - w_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - w_tensor.add_dims(1); - const float scalar = 1.0f + 0.01f * i; - w_tensor.set_raw_data(&scalar, sizeof(float)); - graph.AddInitializedTensor(w_tensor); - - NodeArg* w_arg = &graph.GetOrCreateNodeArg(w_name, &float_scalar); - const std::string out_name = (i == chain_len - 1) ? "Y" : ("h" + std::to_string(i)); - NodeArg* out_arg = &graph.GetOrCreateNodeArg(out_name, &float_1d); - std::vector inputs{prev, w_arg}; - std::vector outputs{out_arg}; - graph.AddNode("op" + std::to_string(i), (i % 2 == 0) ? "Mul" : "Add", "", inputs, outputs); - prev = out_arg; - } - - graph.SetOutputs(std::vector{prev}); - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_TRUE(model.ToProto().SerializeToString(&bytes)); - }; - - for (int s = 0; s < num_sessions; ++s) { - const int64_t num_elements = - static_cast(mb) * 1024 * 1024 / 4 + (vary ? s * 4 * 1024 * 1024 : 0); - - std::string model_bytes; - ASSERT_NO_FATAL_FAILURE(build_scalar_weight_chain(kChainLen, num_elements, model_bytes)); - - SessionOptions so; - so.session_logid = "webgpu_pool_memory"; - auto session = std::make_unique(so, GetEnvironment()); - ASSERT_STATUS_OK(session->RegisterExecutionProvider(make_ep())); - ASSERT_STATUS_OK(session->Load(model_bytes.data(), static_cast(model_bytes.size()))); - ASSERT_STATUS_OK(session->Initialize()); - - std::vector x_values(static_cast(num_elements), 1.0f); - OrtValue x_value; - CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], - std::vector{num_elements}, x_values, &x_value); - - sessions.push_back(std::move(session)); - feeds.push_back(NameMLValMap{{"X", x_value}}); - } - - for (int i = 0; i < iters; ++i) { - for (int s = 0; s < num_sessions; ++s) { - std::vector fetches; - ASSERT_STATUS_OK(sessions[s]->Run(RunOptions{}, feeds[s], output_names, &fetches)); - } - } - - std::cout << "DIAG_POOL sessions=" << num_sessions << " mb=" << mb - << " vary=" << (vary ? 1 : 0) << " iters=" << iters << " holding...\n" - << std::flush; - std::this_thread::sleep_for(std::chrono::milliseconds(hold_ms)); -} - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc index 9389919e8b3ef..4ad1734faabc6 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc @@ -109,11 +109,13 @@ void CopyTensorRoundTrip(Allocator& allocator, float value) { auto cpu_input = Ort::Value::CreateTensor( cpu_memory_info, input_data.data(), input_data.size(), kShape.data(), kShape.size()); auto gpu_tensor = Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size()); + auto gpu_copy = Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size()); auto cpu_output = Ort::Value::CreateTensor( cpu_memory_info, output_data.data(), output_data.size(), kShape.data(), kShape.size()); ThrowOnError(ort_env->CopyTensor(cpu_input, gpu_tensor, nullptr)); - ThrowOnError(ort_env->CopyTensor(gpu_tensor, cpu_output, nullptr)); + ThrowOnError(ort_env->CopyTensor(gpu_tensor, gpu_copy, nullptr)); + ThrowOnError(ort_env->CopyTensor(gpu_copy, cpu_output, nullptr)); if (output_data != input_data) { throw std::runtime_error("CopyTensor round trip returned incorrect data"); } @@ -725,6 +727,17 @@ TEST_F(PluginEpWebGpuConcurrency, SharedAllocatorCreatesAndCopiesConcurrently) { ASSERT_FALSE(error.Failed()) << error.Message(); } +TEST_F(PluginEpWebGpuConcurrency, SameSessionRunsAreSerialized) { + auto session = CreateSession(); + FirstError error; + RunWorkers(error, [&](int thread_id) { + for (int iteration = 0; iteration < kIterations && !error.Failed(); ++iteration) { + RunWithCpuInputAndOutput(*session, static_cast(thread_id * kIterations + iteration + 1)); + } + }); + ASSERT_FALSE(error.Failed()) << error.Message(); +} + TEST_F(PluginEpWebGpuConcurrency, SameSessionAllocatorCreatesAndCopiesConcurrently) { auto allocator_session = CreateSession(); Ort::Allocator allocator(*allocator_session, Device().GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 49a49de72a830..215ad4177019f 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -198,7 +198,7 @@ void TestCopyAfterDeferredDispatch(bool upload) { dispatch.compute_pipeline = pipeline; dispatch.bind_group = context.Device().CreateBindGroup(&bind_group_desc); recording.deferred_dispatches.push_back(std::move(dispatch)); - recording.has_unsubmitted_work = true; + buffer_manager.BeginRecording(recording); if (upload) { // The dispatch must consume the original input before Upload overwrites it. @@ -225,6 +225,71 @@ TEST(WebGpuContextTest, MemCpyFollowsDeferredDispatch) { TestCopyAfterDeferredDispatch(false); } +TEST(WebGpuContextTest, BufferReuseFollowsCachePolicyAndRecordingSubmission) { + ConfigOptions options; + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + constexpr auto usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + + for (auto mode : {webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::LazyRelease, + webgpu::BufferCacheMode::Bucket, webgpu::BufferCacheMode::Simple, + webgpu::BufferCacheMode::Graph, webgpu::BufferCacheMode::GraphSimple}) { + SCOPED_TRACE(static_cast(mode)); + webgpu::BufferManager manager(context, mode, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState first; + webgpu::CommandRecordingState second; + auto buffer = manager.Create(first, 64, usage); + wgpu::Buffer retained{buffer}; + manager.GetCommandEncoder(first).ClearBuffer(buffer, 0, 64); + manager.Release(buffer, &first); + + auto other_buffer = manager.Create(second, 64, usage); + EXPECT_NE(other_buffer, buffer); + manager.GetCommandEncoder(second).ClearBuffer(other_buffer, 0, 64); + ASSERT_STATUS_OK(context.Flush(manager, second)); + // A submission on another recording cannot make the first buffer reusable. + auto after_other_submission = manager.Create(second, 64, usage); + EXPECT_NE(after_other_submission, buffer); + + ASSERT_STATUS_OK(context.Flush(manager, first)); + auto reused = manager.Create(second, 64, usage); + if (mode == webgpu::BufferCacheMode::Disabled || mode == webgpu::BufferCacheMode::LazyRelease) { + EXPECT_NE(reused, buffer); + } else { + EXPECT_EQ(reused, buffer); + } + manager.Release(reused); + manager.Release(other_buffer); + manager.Release(after_other_submission); + } +} + +TEST(WebGpuContextTest, AbandonedRecordingDoesNotReturnBuffersToPool) { + ConfigOptions options; + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + webgpu::BufferManager manager(context, webgpu::BufferCacheMode::Bucket, + webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState recording; + constexpr auto usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + auto buffer = manager.Create(recording, 64, usage); + // Keep the handle alive so a new device allocation cannot reuse its address. + wgpu::Buffer retained{buffer}; + manager.GetCommandEncoder(recording).ClearBuffer(buffer, 0, 64); + manager.Release(buffer, &recording); + recording.command_encoder = nullptr; + manager.DiscardPendingBuffers(recording); + + // Reuse the recording's address after abandoning its previous commands. + manager.GetCommandEncoder(recording); + ASSERT_STATUS_OK(context.Flush(manager, recording)); + auto replacement = manager.Create(recording, 64, usage); + EXPECT_NE(replacement, buffer); + manager.Release(replacement); +} + TEST(WebGpuContextTest, SessionAllocatorSubmitsReusedBufferClearOutsideRun) { ConfigOptions options; auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); diff --git a/onnxruntime/test/providers/webgpu/webgpu_stream_regression_test.cc b/onnxruntime/test/providers/webgpu/webgpu_stream_regression_test.cc index 188cfe41db600..4dabd8aa00c3e 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_stream_regression_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_stream_regression_test.cc @@ -18,16 +18,10 @@ #include "test/test_environment.h" #include "test/util/include/asserts.h" -#if !defined(BUILD_DAWN_SHARED_LIBRARY) -#include "dawn/dawn_proc.h" -#endif -#include "dawn/native/DawnNative.h" - namespace onnxruntime { namespace test { namespace { -#if !defined(BUILD_DAWN_SHARED_LIBRARY) class UnusedKernel final : public OpKernel { public: explicit UnusedKernel(const OpKernelInfo& info) : OpKernel(info) {} @@ -39,41 +33,14 @@ class UnusedKernelContext final : public OpKernelContext { UnusedKernelContext() : OpKernelContext(nullptr, DefaultLoggingManager().DefaultLogger(), nullptr) {} }; -thread_local webgpu::CommandRecordingState* observed_recording = nullptr; -thread_local bool* observed_exclusion = nullptr; -thread_local size_t* observed_clears = nullptr; - -void ObserveClearBuffer(WGPUCommandEncoder encoder, WGPUBuffer buffer, uint64_t offset, uint64_t size) { - if (observed_recording != nullptr) { - // Probe from another thread while the real FillZero is inside ClearBuffer. Joining the - // nonblocking probe makes this independent of sleeps or which thread is scheduled first. - auto* recording = observed_recording; - bool acquired = false; - std::thread contender([&] { - acquired = recording->mutex.try_lock(); - if (acquired) { - recording->mutex.unlock(); - } - }); - contender.join(); - *observed_exclusion = !acquired; - ++*observed_clears; - } - dawn::native::GetProcs().commandEncoderClearBuffer(encoder, buffer, offset, size); -} -#endif - -TEST(WebGpuContextTest, FillZeroExcludesConcurrentSessionAllocator) { -#if defined(BUILD_DAWN_SHARED_LIBRARY) - GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required by this mutex probe."; -#else +TEST(WebGpuContextTest, FillZeroClearsSessionBuffer) { ConfigOptions options; auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); auto& webgpu_ep = static_cast(*ep); auto& context = webgpu::WebGpuContextFactory::GetContext(0); auto& recording = webgpu_ep.Recording(); - Model model("fill_zero_lock", false, DefaultLoggingManager().DefaultLogger()); + Model model("fill_zero", false, DefaultLoggingManager().DefaultLogger()); auto& node = model.MainGraph().AddNode("unused", "Identity", "", {}, {}); auto kernel_def = KernelDefBuilder().SetName("Identity").Provider(kWebGpuExecutionProvider).SinceVersion(1).Build(); const std::unordered_map initializers; @@ -92,28 +59,14 @@ TEST(WebGpuContextTest, FillZeroExcludesConcurrentSessionAllocator) { Tensor tensor(DataTypeImpl::GetType(), TensorShape{16}, buffer.Get(), OrtMemoryInfo(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice, OrtMemTypeDefault)); - auto procs = dawn::native::GetProcs(); - procs.commandEncoderClearBuffer = ObserveClearBuffer; - bool excluded = false; - size_t clears = 0; - observed_recording = &recording; - observed_exclusion = &excluded; - observed_clears = &clears; - struct RestoreProcs { - ~RestoreProcs() { - observed_recording = nullptr; - observed_exclusion = nullptr; - observed_clears = nullptr; - dawnProcSetProcs(&dawn::native::GetProcs()); - } - } restore; - dawnProcSetProcs(&procs); + std::array data; + data.fill(42); + webgpu_ep.BufferManager().Upload(recording, data.data(), buffer.Get(), sizeof(data)); compute_context.FillZero(tensor); - - EXPECT_EQ(clears, 1u); - EXPECT_TRUE(excluded) << "FillZero encoded a clear without holding the Session recording mutex"; - ASSERT_STATUS_OK(context.Flush(webgpu_ep.BufferManager(), recording)); -#endif + webgpu_ep.BufferManager().Download(recording, buffer.Get(), data.data(), sizeof(data)); + for (auto value : data) { + EXPECT_EQ(value, 0u); + } } TEST(WebGpuContextTest, SharedDeviceErrorScopesRemainThreadLocal) { From 1df175445b248e26d763bec218183c3d211aa77b Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 29 Sep 2026 16:14:38 +0800 Subject: [PATCH 2/5] Track WebGPU pending buffers on release --- onnxruntime/core/providers/webgpu/allocator.h | 6 +- .../core/providers/webgpu/buffer_manager.cc | 38 ++++------- .../core/providers/webgpu/buffer_manager.h | 9 +-- .../core/providers/webgpu/compute_context.h | 2 +- .../core/providers/webgpu/ep/README.md | 6 +- .../core/providers/webgpu/webgpu_context.cc | 8 ++- .../core/providers/webgpu/webgpu_context.h | 19 ++++-- .../providers/webgpu/webgpu_context_test.cc | 67 +++++++++++++++++-- 8 files changed, 105 insertions(+), 50 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/allocator.h b/onnxruntime/core/providers/webgpu/allocator.h index 44824f5624e3f..c7b411d872cb6 100644 --- a/onnxruntime/core/providers/webgpu/allocator.h +++ b/onnxruntime/core/providers/webgpu/allocator.h @@ -30,8 +30,10 @@ class GpuBufferAllocator : public IAllocator { // BufferManager. This allows the EP to route allocations to different // buffer managers (e.g., per-graph) without explicit refresh calls. // Read-only initializers skip cached-buffer clears and can be mapped at creation on UMA. - // should_submit_zero_initialize controls built-in Alloc. Plugin Alloc submits independent clears; - // a matching plugin AllocOnStream instead defers clears on the supplied Session stream. + // should_submit_zero_initialize is used only by built-in WebGPU; plugin builds ignore it. + // TODO: Remove this callback once built-in WebGPU can distinguish external allocators used outside Run + // from internal allocators used during Run. + // Plugin Alloc submits independent clears; a matching AllocOnStream defers them on the Session stream. // Only the plugin Env allocator omits recording_getter, as it never uses a Session stream. GpuBufferAllocator(std::function buffer_manager_getter, std::function recording_getter, diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 58c562c6a320b..475fad95e257a 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -496,21 +496,7 @@ BufferManager::BufferManager(WebGpuContext& context, BufferCacheMode storage_buf storage_cache_{CreateBufferCacheManager(storage_buffer_cache_mode)}, uniform_cache_{CreateBufferCacheManager(uniform_buffer_cache_mode)}, query_resolve_cache_{CreateBufferCacheManager(query_resolve_buffer_cache_mode)}, - default_cache_{CreateBufferCacheManager(default_buffer_cache_mode)}, - supports_buffer_reuse_{storage_cache_->SupportsBufferReuse() || uniform_cache_->SupportsBufferReuse() || - query_resolve_cache_->SupportsBufferReuse() || default_cache_->SupportsBufferReuse()} { -} - -void BufferManager::BeginRecording(const CommandRecordingState& recording) const { - if (supports_buffer_reuse_) { - std::lock_guard lock{mutex_}; - pending_buffers_.try_emplace(&recording); - } -} - -const wgpu::CommandEncoder& BufferManager::GetCommandEncoder(CommandRecordingState& recording) const { - BeginRecording(recording); - return context_.GetCommandEncoder(recording); + default_cache_{CreateBufferCacheManager(default_buffer_cache_mode)} { } void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const { @@ -541,7 +527,7 @@ void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuff staging_buffer.Unmap(); ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); - auto& command_encoder = GetCommandEncoder(recording); + auto& command_encoder = context_.GetCommandEncoder(recording); context_.EndComputePass(recording); command_encoder.CopyBufferToBuffer(staging_buffer, 0, dst, 0, copy_size); ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); @@ -560,7 +546,7 @@ void BufferManager::MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGP src_size, ", dst_size=", dst_size, ", copy_size=", copy_size, "."); ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); - auto& command_encoder = GetCommandEncoder(recording); + auto& command_encoder = context_.GetCommandEncoder(recording); context_.EndComputePass(recording); command_encoder.CopyBufferToBuffer(src, 0, dst, 0, copy_size); } @@ -584,7 +570,7 @@ WGPUBuffer BufferManager::Create(CommandRecordingState& recording, size_t size, auto buffer_guard = wgpu::Buffer::Acquire(buffer); ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); context_.EndComputePass(recording); - GetCommandEncoder(recording).ClearBuffer(buffer, 0, buffer_size); + context_.GetCommandEncoder(recording).ClearBuffer(buffer, 0, buffer_size); if (submit_zero_initialize) { ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); } @@ -624,11 +610,10 @@ void BufferManager::Release(WGPUBuffer buffer, const CommandRecordingState* reco EnforceBufferUnmapped(context_, buffer); std::lock_guard lock{mutex_}; auto& cache = GetCacheManager(buffer); - if (cache.SupportsBufferReuse()) { - if (auto it = pending_buffers_.find(recording); it != pending_buffers_.end()) { - it->second.emplace_back(wgpu::Buffer::Acquire(buffer)); - return; - } + if (cache.SupportsBufferReuse() && recording != nullptr && + recording->has_unsubmitted_work.load(std::memory_order_relaxed)) { + pending_buffers_[recording].emplace_back(wgpu::Buffer::Acquire(buffer)); + return; } cache.ReleaseBuffer(buffer); @@ -647,7 +632,7 @@ void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, v desc.usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead; auto staging_buffer = context_.Device().CreateBuffer(&desc); - auto& command_encoder = GetCommandEncoder(recording); + auto& command_encoder = context_.GetCommandEncoder(recording); context_.EndComputePass(recording); command_encoder.CopyBufferToBuffer(src, 0, staging_buffer, 0, buffer_size); ORT_THROW_IF_ERROR(context_.Flush(*this, recording)); @@ -684,6 +669,8 @@ void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, v void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording) const { std::lock_guard lock{mutex_}; + // Serialize the transition to idle with Free, so a concurrent release cannot miss this refresh. + recording.has_unsubmitted_work.store(false, std::memory_order_relaxed); if (auto it = pending_buffers_.find(&recording); it != pending_buffers_.end()) { for (auto& buffer : it->second) { GetCacheManager(buffer.Get()).ReleaseBuffer(buffer.MoveToCHandle()); @@ -701,8 +688,9 @@ void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording) cons default_cache_->OnRefresh(recording.graph_capture_state); } -void BufferManager::DiscardPendingBuffers(const CommandRecordingState& recording) const { +void BufferManager::DiscardPendingBuffers(CommandRecordingState& recording) const { std::lock_guard lock{mutex_}; + recording.has_unsubmitted_work.store(false, std::memory_order_relaxed); pending_buffers_.erase(&recording); } diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index 8eb644481208c..1a9abd9f11288 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/onnxruntime/core/providers/webgpu/buffer_manager.h @@ -94,9 +94,6 @@ class IBufferCacheManager { class BufferManager { public: BufferManager(WebGpuContext& context, BufferCacheMode storage_buffer_cache_mode, BufferCacheMode uniform_buffer_cache_mode, BufferCacheMode query_resolve_buffer_cache_mode, BufferCacheMode default_buffer_cache_mode); - // Register deferred work before releasing any buffers it uses. Submission or abandonment ends it. - void BeginRecording(const CommandRecordingState& recording) const; - const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) const; void Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const; void MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const; WGPUBuffer Create(CommandRecordingState& recording, size_t size, wgpu::BufferUsage usage, @@ -108,7 +105,7 @@ class BufferManager { void Download(CommandRecordingState& recording, WGPUBuffer src, void* dst, size_t size) const; void RefreshPendingBuffers(CommandRecordingState& recording) const; // Drop retained references when a recording is abandoned instead of submitted. - void DiscardPendingBuffers(const CommandRecordingState& recording) const; + void DiscardPendingBuffers(CommandRecordingState& recording) const; std::vector> ExtractCachedBuffers(wgpu::BufferUsage usage); void AbsorbCachedBuffers(wgpu::BufferUsage usage, @@ -118,13 +115,13 @@ class BufferManager { IBufferCacheManager& GetCacheManager(wgpu::BufferUsage usage) const; IBufferCacheManager& GetCacheManager(WGPUBuffer buffer) const; WebGpuContext& context_; + // Protects both the caches and pending_buffers_. mutable std::mutex mutex_; std::unique_ptr storage_cache_; std::unique_ptr uniform_cache_; std::unique_ptr query_resolve_cache_; std::unique_ptr default_cache_; - const bool supports_buffer_reuse_; - // An entry exists only while that recording has unsubmitted work, even if no buffers were freed. + // Created on release, only for reusable buffers waiting for this recording's submission. mutable InlinedHashMap> pending_buffers_; }; diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index 44547d2efdacb..0d8b63e89b3b3 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/onnxruntime/core/providers/webgpu/compute_context.h @@ -236,7 +236,7 @@ class ComputeContext final : public ComputeContextBase { auto& recording = ep_.Recording(); ORT_THROW_IF_ERROR(webgpu_context_.EncodeDeferredDispatches(recording)); webgpu_context_.EndComputePass(recording); - auto& command_encoder = ep_.BufferManager().GetCommandEncoder(recording); + auto& command_encoder = webgpu_context_.GetCommandEncoder(recording); WGPUBuffer buffer = reinterpret_cast(dst.MutableDataRaw()); command_encoder.ClearBuffer(buffer, 0, dst.SizeInBytes()); } diff --git a/onnxruntime/core/providers/webgpu/ep/README.md b/onnxruntime/core/providers/webgpu/ep/README.md index 695ba2a6253af..3fad1a94628e1 100644 --- a/onnxruntime/core/providers/webgpu/ep/README.md +++ b/onnxruntime/core/providers/webgpu/ep/README.md @@ -63,8 +63,10 @@ Session runs. Keep a Session allocator and its owning Session alive until their and synchronize writes before another Session consumes a shared tensor. BufferManager owns pending releases for reusable buffers, grouped by recording and protected by its cache lock. Submission returns only that recording's buffers to the pool; caches without reuse need no such tracking. -Recording starts, submission, and abandonment update this tracking under the same cache lock. -CommandRecordingState contains only command recording data. Context-shared program caches retain their locks. +Pending lists are created only when buffers are released, with no buffer-manager registration per dispatch. +CommandRecordingState exposes an atomic unsubmitted-work flag so concurrent frees after submission can +return buffers directly to the cache. Submission and abandonment clear that flag under the cache lock, +preventing concurrent releases from missing cleanup. Context-shared program caches retain their locks. Environment transfers with no stream use command state local to each CopyTensors call. GPU-to-GPU copies without a stream submit before returning, but do not wait for GPU completion. A subsequent diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 79c88af4d7e7b..42ff70c05d76d 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -740,7 +740,7 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra command.bind_group = CreateBindGroup(bind_buffers, bind_buffers_segments, *bind_group_layout, program.Name()); command.pending_build = std::move(pending_build); - buffer_mgr.BeginRecording(recording); + recording.has_unsubmitted_work.store(true, std::memory_order_relaxed); if (uniform_buffer) { // The bind group owns a reference now, so return the allocator's reference immediately. buffer_mgr.Release(uniform_buffer, &recording); @@ -1087,6 +1087,11 @@ Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording) { Status status = EncodeDeferredDispatches(recording); if (!recording.command_encoder) { + if (status.IsOK()) { + buffer_mgr.RefreshPendingBuffers(recording); + } else { + buffer_mgr.DiscardPendingBuffers(recording); + } return status; } @@ -1237,7 +1242,6 @@ void WebGpuContext::Replay(const std::vector& captu recording.pending_kernels.emplace_back(*command.pending_kernel_info); } - buffer_manager.BeginRecording(recording); DispatchCommand(command, recording); if (recording.num_pending_dispatches >= max_num_pending_dispatches_) { ORT_THROW_IF_ERROR(Flush(buffer_manager, recording)); diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index df20d6619d56a..226d102cad5b4 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include #include @@ -118,6 +119,9 @@ struct CommandRecordingState { wgpu::CommandEncoder command_encoder; wgpu::ComputePassEncoder compute_pass_encoder; uint32_t num_pending_dispatches = 0; + // Concurrent allocator frees inspect only this flag, never the encoders or deferred dispatches. + // BufferManager clears it under its cache lock after submission or abandonment. + std::atomic has_unsubmitted_work{false}; std::vector deferred_dispatches; std::vector pending_kernels; GraphCaptureState graph_capture_state{GraphCaptureState::Default}; @@ -256,6 +260,14 @@ class WebGpuContext final { bool DeviceHasFeature(wgpu::FeatureName feature) const { return device_features_.contains(feature); } const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return subgroup_matrix_configs_; } + const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) { + if (!recording.command_encoder) { + recording.command_encoder = device_.CreateCommandEncoder(); + recording.has_unsubmitted_work.store(true, std::memory_order_relaxed); + } + return recording.command_encoder; + } + const wgpu::ComputePassEncoder& GetComputePassEncoder(CommandRecordingState& recording) { if (!recording.compute_pass_encoder) { auto& command_encoder = GetCommandEncoder(recording); @@ -388,13 +400,6 @@ class WebGpuContext final { void Initialize(const WebGpuContextConfig& config); - const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) { - if (!recording.command_encoder) { - recording.command_encoder = device_.CreateCommandEncoder(); - } - return recording.command_encoder; - } - wgpu::BindGroup CreateBindGroup(const std::vector& bind_buffers, const std::vector& bind_buffers_segments, const wgpu::BindGroupLayout& bind_group_layout, diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 215ad4177019f..e2db86d49c760 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -198,7 +198,7 @@ void TestCopyAfterDeferredDispatch(bool upload) { dispatch.compute_pipeline = pipeline; dispatch.bind_group = context.Device().CreateBindGroup(&bind_group_desc); recording.deferred_dispatches.push_back(std::move(dispatch)); - buffer_manager.BeginRecording(recording); + recording.has_unsubmitted_work.store(true, std::memory_order_relaxed); if (upload) { // The dispatch must consume the original input before Upload overwrites it. @@ -241,12 +241,14 @@ TEST(WebGpuContextTest, BufferReuseFollowsCachePolicyAndRecordingSubmission) { webgpu::CommandRecordingState second; auto buffer = manager.Create(first, 64, usage); wgpu::Buffer retained{buffer}; - manager.GetCommandEncoder(first).ClearBuffer(buffer, 0, 64); + context.GetCommandEncoder(first).ClearBuffer(buffer, 0, 64); manager.Release(buffer, &first); + auto within_batch = manager.Create(first, 64, usage); + EXPECT_NE(within_batch, buffer); auto other_buffer = manager.Create(second, 64, usage); EXPECT_NE(other_buffer, buffer); - manager.GetCommandEncoder(second).ClearBuffer(other_buffer, 0, 64); + context.GetCommandEncoder(second).ClearBuffer(other_buffer, 0, 64); ASSERT_STATUS_OK(context.Flush(manager, second)); // A submission on another recording cannot make the first buffer reusable. auto after_other_submission = manager.Create(second, 64, usage); @@ -260,11 +262,66 @@ TEST(WebGpuContextTest, BufferReuseFollowsCachePolicyAndRecordingSubmission) { EXPECT_EQ(reused, buffer); } manager.Release(reused); + manager.Release(within_batch); manager.Release(other_buffer); manager.Release(after_other_submission); } } +TEST(WebGpuContextTest, BufferReleasedAfterSubmissionIsReusableWithoutAnotherFlush) { + ConfigOptions options; + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + webgpu::BufferManager manager(context, webgpu::BufferCacheMode::Bucket, + webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState recording; + webgpu::CommandRecordingState other; + constexpr auto usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + auto buffer = manager.Create(recording, 64, usage); + wgpu::Buffer retained{buffer}; + context.GetCommandEncoder(recording).ClearBuffer(buffer, 0, 64); + ASSERT_STATUS_OK(context.Flush(manager, recording)); + + // An allocator can outlive its last Run. Free must not wait for another submission to this recording. + manager.Release(buffer, &recording); + auto reused = manager.Create(other, 64, usage); + EXPECT_EQ(reused, buffer); + manager.Release(reused); +} + +TEST(WebGpuContextTest, EmptyFlushRefreshesIdleCacheWithoutReleasingAnotherRecording) { + ConfigOptions options; + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + webgpu::BufferManager manager(context, webgpu::BufferCacheMode::Simple, + webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState idle; + webgpu::CommandRecordingState active; + constexpr auto usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + auto idle_buffer = manager.Create(idle, 64, usage); + auto active_buffer = manager.Create(active, 64, usage); + wgpu::Buffer retained_idle{idle_buffer}; + wgpu::Buffer retained_active{active_buffer}; + context.GetCommandEncoder(active).ClearBuffer(active_buffer, 0, 64); + manager.Release(active_buffer, &active); + manager.Release(idle_buffer, &idle); + + ASSERT_STATUS_OK(context.Flush(manager, idle)); + auto reused = manager.Create(idle, 64, usage); + EXPECT_EQ(reused, idle_buffer); + auto still_pending = manager.Create(idle, 64, usage); + EXPECT_NE(still_pending, active_buffer); + + ASSERT_STATUS_OK(context.Flush(manager, active)); + auto after_submission = manager.Create(idle, 64, usage); + EXPECT_EQ(after_submission, active_buffer); + manager.Release(reused); + manager.Release(still_pending); + manager.Release(after_submission); +} + TEST(WebGpuContextTest, AbandonedRecordingDoesNotReturnBuffersToPool) { ConfigOptions options; auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); @@ -277,13 +334,13 @@ TEST(WebGpuContextTest, AbandonedRecordingDoesNotReturnBuffersToPool) { auto buffer = manager.Create(recording, 64, usage); // Keep the handle alive so a new device allocation cannot reuse its address. wgpu::Buffer retained{buffer}; - manager.GetCommandEncoder(recording).ClearBuffer(buffer, 0, 64); + context.GetCommandEncoder(recording).ClearBuffer(buffer, 0, 64); manager.Release(buffer, &recording); recording.command_encoder = nullptr; manager.DiscardPendingBuffers(recording); // Reuse the recording's address after abandoning its previous commands. - manager.GetCommandEncoder(recording); + context.GetCommandEncoder(recording); ASSERT_STATUS_OK(context.Flush(manager, recording)); auto replacement = manager.Create(recording, 64, usage); EXPECT_NE(replacement, buffer); From 0fd162601549988b71333ded8452ccbf9f952c8c Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 30 Sep 2026 13:18:20 +0800 Subject: [PATCH 3/5] Avoid session capture state reads in WebGPU streamless allocation --- .../core/providers/webgpu/allocator.cc | 5 +- .../core/providers/webgpu/buffer_manager.cc | 4 ++ .../webgpu/plugin/test_webgpu_concurrency.cc | 56 ++++++++++++++++++- .../providers/webgpu/webgpu_context_test.cc | 47 ++++++++++++++++ 4 files changed, 106 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index 483a4bcabcb90..bab768ab90234 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -35,11 +35,8 @@ GpuBufferAllocator::GpuBufferAllocator( // Built-in callers can supply a different submission policy. void* GpuBufferAllocator::Alloc(size_t size) { #if defined(ORT_USE_EP_API_ADAPTERS) + // Streamless clears are independent of Run and must not read its capture state. CommandRecordingState recording; - if (recording_getter_) { - // Preserve the per-graph cache's release policy without borrowing the Session's commands. - recording.graph_capture_state = recording_getter_().graph_capture_state; - } return Allocate(size, recording, true); #else return Allocate(size, recording_getter_(), should_submit_zero_initialize_ && should_submit_zero_initialize_()); diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 475fad95e257a..ab2be2995c658 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -610,6 +610,10 @@ void BufferManager::Release(WGPUBuffer buffer, const CommandRecordingState* reco EnforceBufferUnmapped(context_, buffer); std::lock_guard lock{mutex_}; auto& cache = GetCacheManager(buffer); + // Relaxed ordering suffices: recording a buffer's use (and setting the flag) must happen-before + // its release, on the same thread or through a synchronized ownership handoff. Only independent + // buffer releases may race with recording. The flag publishes no command data, and mutex_ orders + // releases against clearing it after submission or abandonment. if (cache.SupportsBufferReuse() && recording != nullptr && recording->has_unsubmitted_work.load(std::memory_order_relaxed)) { pending_buffers_[recording].emplace_back(wgpu::Buffer::Acquire(buffer)); diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc index 4ad1734faabc6..6cd6e188eb495 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc @@ -101,7 +101,7 @@ void RunWorkers(FirstError& error, Work work) { } template -void CopyTensorRoundTrip(Allocator& allocator, float value) { +void CopyTensorRoundTrip(Allocator& allocator, float value, bool verify_zero_initialization = false) { const auto cpu_memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); std::array input_data{}; input_data.fill(value); @@ -113,6 +113,14 @@ void CopyTensorRoundTrip(Allocator& allocator, float value) { auto cpu_output = Ort::Value::CreateTensor( cpu_memory_info, output_data.data(), output_data.size(), kShape.data(), kShape.size()); + if (verify_zero_initialization) { + output_data.fill(-1.0f); + ThrowOnError(ort_env->CopyTensor(gpu_tensor, cpu_output, nullptr)); + if (output_data != std::array{}) { + throw std::runtime_error("Streamless allocation was not zero initialized"); + } + } + ThrowOnError(ort_env->CopyTensor(cpu_input, gpu_tensor, nullptr)); ThrowOnError(ort_env->CopyTensor(gpu_tensor, gpu_copy, nullptr)); ThrowOnError(ort_env->CopyTensor(gpu_copy, cpu_output, nullptr)); @@ -142,10 +150,13 @@ class PluginEpWebGpuConcurrency : public ::testing::Test { Device(), OrtDeviceMemoryType_DEFAULT, OrtDeviceAllocator, nullptr); } - std::unique_ptr CreateSession() const { + std::unique_ptr CreateSession(bool enable_graph_capture = false) const { Ort::SessionOptions session_options; session_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); std::unordered_map ep_options; + if (enable_graph_capture) { + ep_options["enableGraphCapture"] = "1"; + } session_options.AppendExecutionProvider_V2(*ort_env, {Device()}, ep_options); return std::make_unique( *ort_env, ORT_TSTR("testdata/mul_1.onnx"), session_options); @@ -750,6 +761,47 @@ TEST_F(PluginEpWebGpuConcurrency, SameSessionAllocatorCreatesAndCopiesConcurrent ASSERT_FALSE(error.Failed()) << error.Message(); } +TEST_F(PluginEpWebGpuConcurrency, SameSessionAllocatorCreatesAndCopiesDuringGraphReplay) { + auto session = CreateSession(true); + Ort::Allocator allocator(*session, Device().GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); + auto gpu_input = Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size()); + auto gpu_output = Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size()); + const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); + std::array input_data{}; + std::array output_data{}; + auto cpu_input = Ort::Value::CreateTensor( + cpu_memory, input_data.data(), input_data.size(), kShape.data(), kShape.size()); + auto cpu_output = Ort::Value::CreateTensor( + cpu_memory, output_data.data(), output_data.size(), kShape.data(), kShape.size()); + Ort::IoBinding binding(*session); + binding.BindInput("X", gpu_input); + binding.BindOutput("Y", gpu_output); + const auto run_and_verify = [&](float value) { + input_data.fill(value); + ThrowOnError(ort_env->CopyTensor(cpu_input, gpu_input, nullptr)); + session->Run(Ort::RunOptions{nullptr}, binding); + ThrowOnError(ort_env->CopyTensor(gpu_output, cpu_output, nullptr)); + VerifyOutput(output_data, value); + }; + + // Warm up and capture before starting allocator workers. Only replay is concurrent; + // bound tensors and the binding itself stay on the Run thread. + run_and_verify(1.0f); + CopyTensorRoundTrip(allocator, 1.0f, true); + FirstError error; + RunWorkers(error, [&](int thread_id) { + for (int iteration = 0; iteration < kIterations && !error.Failed(); ++iteration) { + const float value = static_cast(thread_id * kIterations + iteration + 2); + if (thread_id == 0) { + run_and_verify(value); + } else { + CopyTensorRoundTrip(allocator, value, true); + } + } + }); + ASSERT_FALSE(error.Failed()) << error.Message(); +} + TEST_F(PluginEpWebGpuConcurrency, DedicatedSessionAllocatorFeedsConcurrentSessions) { auto allocator_session = CreateSession(); Ort::Allocator allocator(*allocator_session, Device().GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index e2db86d49c760..cab90b712ac49 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -322,6 +322,53 @@ TEST(WebGpuContextTest, EmptyFlushRefreshesIdleCacheWithoutReleasingAnotherRecor manager.Release(after_submission); } +TEST(WebGpuContextTest, IndependentClearPreservesPendingAndCapturedUniformBuffers) { + ConfigOptions options; + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + webgpu::BufferManager manager(context, webgpu::BufferCacheMode::Graph, + webgpu::BufferCacheMode::GraphSimple, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState session_recording; + webgpu::CommandRecordingState independent_recording; + constexpr auto uniform_usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst; + constexpr auto storage_usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | + wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; + + std::vector captured_commands; + context.CaptureBegin(&captured_commands, manager, session_recording); + auto uniform_buffer = manager.Create(session_recording, 64, uniform_usage); + wgpu::Buffer retained_uniform{uniform_buffer}; + context.GetCommandEncoder(session_recording).ClearBuffer(uniform_buffer, 0, 64); + manager.Release(uniform_buffer, &session_recording); + + const auto submit_independent_clear = [&]() { + // Exercise the cached clear and immediate submission used by plugin streamless allocations. + std::array dirty_data; + dirty_data.fill(0xffffffffu); + auto dirty_buffer = manager.Create(independent_recording, sizeof(dirty_data), storage_usage); + manager.Upload(independent_recording, dirty_data.data(), dirty_buffer, sizeof(dirty_data)); + manager.Release(dirty_buffer, &independent_recording); + auto cleared_buffer = wgpu::Buffer::Acquire( + manager.Create(independent_recording, sizeof(dirty_data), storage_usage, true, true)); + EXPECT_EQ(cleared_buffer.Get(), dirty_buffer); + const std::array zeros{}; + EXPECT_EQ(ReadBufferWithExternalCommandEncoder(context, cleared_buffer.Get()), zeros); + manager.Release(cleared_buffer.MoveToCHandle(), &independent_recording); + }; + + submit_independent_clear(); + auto before_submission = wgpu::Buffer::Acquire(manager.Create(independent_recording, 64, uniform_usage)); + EXPECT_NE(before_submission.Get(), uniform_buffer); + EXPECT_TRUE(session_recording.has_unsubmitted_work.load(std::memory_order_relaxed)); + + ASSERT_STATUS_OK(context.Flush(manager, session_recording)); + context.CaptureEnd(session_recording); + submit_independent_clear(); + auto after_submission = wgpu::Buffer::Acquire(manager.Create(independent_recording, 64, uniform_usage)); + EXPECT_NE(after_submission.Get(), uniform_buffer); +} + TEST(WebGpuContextTest, AbandonedRecordingDoesNotReturnBuffersToPool) { ConfigOptions options; auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); From 5cfd909918e86bf13b0f2b7b384c49ecd4cdc8e1 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 30 Sep 2026 19:34:06 +0800 Subject: [PATCH 4/5] Fix WebGPU pending buffer retirement across managers 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. --- .../core/providers/webgpu/buffer_manager.cc | 14 ++- .../core/providers/webgpu/buffer_manager.h | 3 +- .../core/providers/webgpu/ep/README.md | 7 +- .../core/providers/webgpu/webgpu_context.cc | 16 ++- .../webgpu/plugin/test_webgpu_concurrency.cc | 112 +++++++++++++++++ .../providers/webgpu/webgpu_context_test.cc | 113 ++++++++++++++++++ 6 files changed, 253 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index ab2be2995c658..228a118f1ce96 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -610,6 +610,8 @@ void BufferManager::Release(WGPUBuffer buffer, const CommandRecordingState* reco EnforceBufferUnmapped(context_, buffer); std::lock_guard lock{mutex_}; auto& cache = GetCacheManager(buffer); + // Keep reusable buffers out of the cache while this recording has unsubmitted work. + // Graph caches also need this protection from streamless allocations. // Relaxed ordering suffices: recording a buffer's use (and setting the flag) must happen-before // its release, on the same thread or through a synchronized ownership handoff. Only independent // buffer releases may race with recording. The flag publishes no command data, and mutex_ orders @@ -671,7 +673,7 @@ void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, v staging_buffer.Unmap(); } -void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording) const { +void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording, GraphCaptureState graph_capture_state) const { std::lock_guard lock{mutex_}; // Serialize the transition to idle with Free, so a concurrent release cannot miss this refresh. recording.has_unsubmitted_work.store(false, std::memory_order_relaxed); @@ -682,14 +684,14 @@ void BufferManager::RefreshPendingBuffers(CommandRecordingState& recording) cons pending_buffers_.erase(it); } - if (recording.graph_capture_state == GraphCaptureState::Replaying) { + if (graph_capture_state == GraphCaptureState::Replaying) { return; } - storage_cache_->OnRefresh(recording.graph_capture_state); - uniform_cache_->OnRefresh(recording.graph_capture_state); - query_resolve_cache_->OnRefresh(recording.graph_capture_state); - default_cache_->OnRefresh(recording.graph_capture_state); + storage_cache_->OnRefresh(graph_capture_state); + uniform_cache_->OnRefresh(graph_capture_state); + query_resolve_cache_->OnRefresh(graph_capture_state); + default_cache_->OnRefresh(graph_capture_state); } void BufferManager::DiscardPendingBuffers(CommandRecordingState& recording) const { diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index 1a9abd9f11288..f7e1fb033c816 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/onnxruntime/core/providers/webgpu/buffer_manager.h @@ -51,7 +51,6 @@ class IBufferCacheManager { public: virtual ~IBufferCacheManager() = default; - // Reusable buffers must stay out of the cache until their recording submits. virtual bool SupportsBufferReuse() const { return true; } // calculate actual buffer size to allocate based on the requested size. @@ -103,7 +102,7 @@ class BufferManager { // A null recording releases directly to the cache (e.g., plugin Env allocations). void Release(WGPUBuffer buffer, const CommandRecordingState* recording = nullptr) const; void Download(CommandRecordingState& recording, WGPUBuffer src, void* dst, size_t size) const; - void RefreshPendingBuffers(CommandRecordingState& recording) const; + void RefreshPendingBuffers(CommandRecordingState& recording, GraphCaptureState graph_capture_state) const; // Drop retained references when a recording is abandoned instead of submitted. void DiscardPendingBuffers(CommandRecordingState& recording) const; diff --git a/onnxruntime/core/providers/webgpu/ep/README.md b/onnxruntime/core/providers/webgpu/ep/README.md index 3fad1a94628e1..51d585f3fe3c8 100644 --- a/onnxruntime/core/providers/webgpu/ep/README.md +++ b/onnxruntime/core/providers/webgpu/ep/README.md @@ -86,8 +86,11 @@ copies, mixed feed copies, CPU outputs bound to GPU, concurrent independent Sess same-Session Runs, same-Session and dedicated-Session allocator concurrency, and shared Env allocation/transfers are covered by the tests in `onnxruntime/test/providers/webgpu/plugin`, built into `onnxruntime_provider_test`. -Concurrent graph capture and concurrent profiling are unsupported. Cross-device transfer and -arbitrary foreign stream overrides are not established by these tests. Performance must be measured separately. +The tests also cover initial graph capture and replay across independent Sessions, and concurrent +Run calls for different graph IDs within one Session, serialized by ORT. Initial-capture tests use +preallocated bindings; they do not cover external Session allocator calls overlapping capture. +Concurrent profiling, cross-device transfer, and arbitrary foreign stream overrides are not +established by these tests. Performance must be measured separately. ### Missing parts diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 42ff70c05d76d..7120fbd285370 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -1085,12 +1085,24 @@ Status WebGpuContext::PopErrorScope() { Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording) { + // Graph runs and prepacking can release allocator/uniform buffers into the shared manager. + // Retire this recording in both managers, keeping graph cache policy local to the active one. + const auto refresh_pending_buffers = [&]() { + buffer_mgr.RefreshPendingBuffers(recording, recording.graph_capture_state); + if (&buffer_mgr != buffer_mgr_.get()) { + buffer_mgr_->RefreshPendingBuffers(recording, GraphCaptureState::Default); + } + }; + Status status = EncodeDeferredDispatches(recording); if (!recording.command_encoder) { if (status.IsOK()) { - buffer_mgr.RefreshPendingBuffers(recording); + refresh_pending_buffers(); } else { buffer_mgr.DiscardPendingBuffers(recording); + if (&buffer_mgr != buffer_mgr_.get()) { + buffer_mgr_->DiscardPendingBuffers(recording); + } } return status; } @@ -1134,7 +1146,7 @@ Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, } auto command_buffer = recording.command_encoder.Finish(); device_queue_.Submit(1, &command_buffer); - buffer_mgr.RefreshPendingBuffers(recording); + refresh_pending_buffers(); recording.command_encoder = nullptr; recording.num_pending_dispatches = 0; return status; diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc index 6cd6e188eb495..136879b5a6208 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_concurrency.cc @@ -18,6 +18,7 @@ #include #include +#include "core/common/inlined_containers.h" #include "core/graph/constants.h" #include "core/graph/onnx_protobuf.h" #include "core/session/onnxruntime_cxx_api.h" @@ -491,6 +492,58 @@ class PluginEpWebGpuConcurrency : public ::testing::Test { EXPECT_EQ(copy_kernels, 1u); } + void RunGraphsConcurrently(bool separate_sessions, bool capture_before_workers) const { + InlinedVector> sessions; + const int session_count = separate_sessions ? kThreads : 1; + for (int i = 0; i < session_count; ++i) { + sessions.emplace_back(CreateSession(true)); + } + const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); + InlinedVector allocators; + for (const auto& session : sessions) { + allocators.emplace_back(*session, Device().GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); + } + InlinedVector inputs; + InlinedVector outputs; + InlinedVector bindings; + InlinedVector options; + for (int thread_id = 0; thread_id < kThreads; ++thread_id) { + auto& session = *sessions[separate_sessions ? thread_id : 0]; + auto& allocator = allocators[separate_sessions ? thread_id : 0]; + inputs.emplace_back(Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size())); + outputs.emplace_back(Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size())); + std::array input_data{}; + input_data.fill(static_cast(thread_id + 1)); + auto cpu_input = Ort::Value::CreateTensor( + cpu_memory, input_data.data(), input_data.size(), kShape.data(), kShape.size()); + ThrowOnError(ort_env->CopyTensor(cpu_input, inputs.back(), nullptr)); + bindings.emplace_back(session); + bindings.back().BindInput("X", inputs.back()); + bindings.back().BindOutput("Y", outputs.back()); + options.emplace_back(); + options.back().AddConfigEntry("gpu_graph_id", std::to_string(thread_id).c_str()); + if (capture_before_workers) { + session.Run(options.back(), bindings.back()); + } + } + + FirstError error; + RunWorkers(error, [&](int thread_id) { + auto& session = *sessions[separate_sessions ? thread_id : 0]; + for (int iteration = 0; iteration < kIterations && !error.Failed(); ++iteration) { + // Without pre-capture, the first call includes warm-up and capture retries. + // Later calls replay using the same bound buffers and graph annotation ID. + session.Run(options[thread_id], bindings[thread_id]); + std::array output_data{}; + auto cpu_output = Ort::Value::CreateTensor( + cpu_memory, output_data.data(), output_data.size(), kShape.data(), kShape.size()); + ThrowOnError(ort_env->CopyTensor(outputs[thread_id], cpu_output, nullptr)); + VerifyOutput(output_data, static_cast(thread_id + 1)); + } + }); + ASSERT_FALSE(error.Failed()) << error.Message(); + } + private: std::optional registration_; Ort::ConstEpDevice webgpu_ep_device_{nullptr}; @@ -802,6 +855,65 @@ TEST_F(PluginEpWebGpuConcurrency, SameSessionAllocatorCreatesAndCopiesDuringGrap ASSERT_FALSE(error.Failed()) << error.Message(); } +TEST_F(PluginEpWebGpuConcurrency, SessionAllocatorReusesBuffersFreedDuringGraphReplay) { + auto session = CreateSession(true); + Ort::Allocator allocator(*session, Device().GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); + auto gpu_input = Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size()); + auto gpu_output = Ort::Value::CreateTensor(allocator, kShape.data(), kShape.size()); + Ort::IoBinding binding(*session); + binding.BindInput("X", gpu_input); + binding.BindOutput("Y", gpu_output); + session->Run(Ort::RunOptions{nullptr}, binding); + + // Hold a full bucket so all available buffers of this size belong to this test. + // The model's tensors use the separate 64-byte bucket. + constexpr size_t kBufferSize = 4096; + constexpr size_t kBucketCapacity = 200; + InlinedVector> buffers; + InlinedHashSet handles; + for (size_t i = 0; i < kBucketCapacity; ++i) { + buffers.emplace_back(allocator.GetAllocation(kBufferSize)); + ASSERT_NE(buffers.back()->get(), nullptr); + handles.insert(buffers.back()->get()); + } + + FirstError error; + RunWorkers(error, [&](int thread_id) { + if (thread_id == 0) { + for (int iteration = 0; iteration < kIterations && !error.Failed(); ++iteration) { + session->Run(Ort::RunOptions{nullptr}, binding); + } + } else { + for (size_t i = static_cast(thread_id - 1); i < buffers.size(); i += kThreads - 1) { + buffers[i].reset(); + std::this_thread::yield(); + } + } + }); + ASSERT_FALSE(error.Failed()) << error.Message(); + // Submit once after every Free has completed, including frees at the end of the replay loop. + session->Run(Ort::RunOptions{nullptr}, binding); + + InlinedVector reacquired; + for (size_t i = 0; i < kBucketCapacity; ++i) { + reacquired.emplace_back(allocator.GetAllocation(kBufferSize)); + EXPECT_EQ(handles.erase(reacquired.back().get()), 1u); + } + EXPECT_TRUE(handles.empty()); +} + +TEST_F(PluginEpWebGpuConcurrency, SameSessionCapturedGraphsReplayConcurrently) { + RunGraphsConcurrently(/*separate_sessions=*/false, /*capture_before_workers=*/true); +} + +TEST_F(PluginEpWebGpuConcurrency, SameSessionGraphsCaptureAndReplayConcurrently) { + RunGraphsConcurrently(/*separate_sessions=*/false, /*capture_before_workers=*/false); +} + +TEST_F(PluginEpWebGpuConcurrency, DifferentSessionsGraphsCaptureAndReplayConcurrently) { + RunGraphsConcurrently(/*separate_sessions=*/true, /*capture_before_workers=*/false); +} + TEST_F(PluginEpWebGpuConcurrency, DedicatedSessionAllocatorFeedsConcurrentSessions) { auto allocator_session = CreateSession(); Ort::Allocator allocator(*allocator_session, Device().GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index cab90b712ac49..bf3fae88ac9e5 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -369,6 +369,119 @@ TEST(WebGpuContextTest, IndependentClearPreservesPendingAndCapturedUniformBuffer EXPECT_NE(after_submission.Get(), uniform_buffer); } +TEST(WebGpuContextTest, GraphFlushRetiresSharedBuffersForSubmittedRecording) { + RunWithFreshDefaultContext([]() { + ConfigOptions options; + ASSERT_STATUS_OK(options.AddConfigEntry(kStorageBufferCacheMode, "bucket")); + ASSERT_STATUS_OK(options.AddConfigEntry(kUniformBufferCacheMode, "simple")); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + auto& shared_manager = context.BufferManager(); + webgpu::BufferManager graph_manager(context, webgpu::BufferCacheMode::Graph, + webgpu::BufferCacheMode::GraphSimple, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + + // Exercise both the Bucket and Simple shared caches during capture and replay. + for (auto state : {webgpu::GraphCaptureState::Capturing, webgpu::GraphCaptureState::Replaying}) { + for (auto usage : {wgpu::BufferUsage::Storage, wgpu::BufferUsage::Uniform}) { + SCOPED_TRACE(static_cast(state)); + SCOPED_TRACE(static_cast(usage)); + usage |= wgpu::BufferUsage::CopyDst; + webgpu::CommandRecordingState session; + webgpu::CommandRecordingState other; + webgpu::CommandRecordingState idle; + session.graph_capture_state = state; + auto released = shared_manager.Create(idle, 64, usage); + auto other_released = shared_manager.Create(idle, 64, usage); + wgpu::Buffer retained{released}; + wgpu::Buffer other_retained{other_released}; + context.GetCommandEncoder(session).ClearBuffer(released, 0, 64); + context.GetCommandEncoder(other).ClearBuffer(other_released, 0, 64); + shared_manager.Release(released, &session); + shared_manager.Release(other_released, &other); + + ASSERT_STATUS_OK(context.Flush(graph_manager, idle)); + auto before_submission = wgpu::Buffer::Acquire(shared_manager.Create(idle, 64, usage)); + EXPECT_NE(before_submission.Get(), released); + EXPECT_NE(before_submission.Get(), other_released); + + ASSERT_STATUS_OK(context.Flush(graph_manager, session)); + EXPECT_EQ(session.graph_capture_state, state); + auto reused = wgpu::Buffer::Acquire(shared_manager.Create(idle, 64, usage)); + EXPECT_EQ(reused.Get(), released); + auto still_pending = wgpu::Buffer::Acquire(shared_manager.Create(idle, 64, usage)); + EXPECT_NE(still_pending.Get(), other_released); + + ASSERT_STATUS_OK(context.Flush(shared_manager, other)); + auto other_reused = wgpu::Buffer::Acquire(shared_manager.Create(idle, 64, usage)); + EXPECT_EQ(other_reused.Get(), other_released); + } + } + }); +} + +TEST(WebGpuContextTest, InitializerFlushRetiresSharedUniformBuffers) { + RunWithFreshDefaultContext([]() { + ConfigOptions options; + ASSERT_STATUS_OK(options.AddConfigEntry(kUniformBufferCacheMode, "simple")); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + auto& shared_manager = context.BufferManager(); + constexpr auto usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst; + + // Prepacking flushes the initializer manager, while uniforms use the shared manager. + // An empty flush must also promote idle buffers from SimpleCacheManager's pending list. + for (bool encode_commands : {true, false}) { + SCOPED_TRACE(encode_commands); + webgpu::CommandRecordingState recording; + auto released = shared_manager.Create(recording, 64, usage); + wgpu::Buffer retained{released}; + if (encode_commands) { + context.GetCommandEncoder(recording).ClearBuffer(released, 0, 64); + } + shared_manager.Release(released, &recording); + ASSERT_STATUS_OK(context.Flush(context.InitializerBufferManager(), recording)); + auto reused = wgpu::Buffer::Acquire(shared_manager.Create(recording, 64, usage)); + EXPECT_EQ(reused.Get(), released); + } + }); +} + +TEST(WebGpuContextTest, FailedFlushDiscardsPendingBuffersFromBothManagers) { + RunWithFreshDefaultContext([]() { + ConfigOptions options; + ASSERT_STATUS_OK(options.AddConfigEntry(kStorageBufferCacheMode, "bucket")); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + auto& shared_manager = context.BufferManager(); + webgpu::BufferManager manager(context, webgpu::BufferCacheMode::Bucket, + webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState recording; + constexpr auto usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopyDst; + auto released = manager.Create(recording, 64, usage); + auto shared_released = shared_manager.Create(recording, 64, usage); + wgpu::Buffer retained{released}; + wgpu::Buffer shared_retained{shared_released}; + webgpu::CapturedCommandInfo dispatch; + dispatch.program_key = "missing_pipeline_for_failed_flush_test"; + recording.deferred_dispatches.push_back(std::move(dispatch)); + recording.has_unsubmitted_work.store(true, std::memory_order_relaxed); + manager.Release(released, &recording); + shared_manager.Release(shared_released, &recording); + + EXPECT_FALSE(context.Flush(manager, recording).IsOK()); + EXPECT_EQ(recording.command_encoder, nullptr); + EXPECT_TRUE(recording.deferred_dispatches.empty()); + // Reusing the recording after failure must not revive discarded buffers in either manager. + ASSERT_STATUS_OK(context.Flush(manager, recording)); + auto replacement = wgpu::Buffer::Acquire(manager.Create(recording, 64, usage)); + auto shared_replacement = wgpu::Buffer::Acquire(shared_manager.Create(recording, 64, usage)); + EXPECT_NE(replacement.Get(), released); + EXPECT_NE(shared_replacement.Get(), shared_released); + }); +} + TEST(WebGpuContextTest, AbandonedRecordingDoesNotReturnBuffersToPool) { ConfigOptions options; auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); From 0493b97826a3ad65ffe8e10a2390fddff5ca2239 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 30 Sep 2026 20:02:56 +0800 Subject: [PATCH 5/5] Restore WebGPU shared-helper tests and clarify allocator policy --- .../core/providers/webgpu/ep/allocator.h | 4 +- .../webgpu/concurrent_context_test.cc | 204 ++++++++++++++++++ 2 files changed, 206 insertions(+), 2 deletions(-) create mode 100644 onnxruntime/test/providers/webgpu/concurrent_context_test.cc diff --git a/onnxruntime/core/providers/webgpu/ep/allocator.h b/onnxruntime/core/providers/webgpu/ep/allocator.h index d2ae375d66ca5..87bef112e0f8e 100644 --- a/onnxruntime/core/providers/webgpu/ep/allocator.h +++ b/onnxruntime/core/providers/webgpu/ep/allocator.h @@ -37,8 +37,8 @@ namespace ep { // Alloc vs AllocOnStream is a stream-based distinction, not an external-vs-internal API distinction: // BindInput can allocate on a Session stream before Run; streamless allocation during Run uses Alloc. // Read-only initializers and writable prepacked weights use separate GpuBufferAllocator instances -// with InitializerBufferManager(). Read-only initializers skip clears; prepack and native callers -// can supply different plain-Alloc submission policies. +// with InitializerBufferManager(). Read-only initializers skip clears. Plugin plain Alloc always +// submits cached clears independently; only built-in WebGPU uses the submission-policy callback. OrtAllocator* CreateWebGpuSessionAllocator(AllocatorPtr allocator); bool TryReleaseWebGpuSessionAllocator(OrtAllocator* allocator); diff --git a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc new file mode 100644 index 0000000000000..21363f82f2752 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc @@ -0,0 +1,204 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/graph/model.h" +#include "core/graph/onnx_protobuf.h" +#include "core/platform/env.h" +#include "core/providers/webgpu/program_manager.h" +#include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/inference_session.h" + +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/temp_dir.h" + +namespace onnxruntime { +namespace test { + +TEST(WebGpuConcurrentContextTestStandalone, ShaderDumpWritesRemainComplete) { + TemporaryDirectory temp_dir{ORT_TSTR("webgpu_shader_dump_test")}; + const auto dump_path = std::filesystem::path(temp_dir.Path()) / ORT_TSTR("shaders.txt"); + constexpr int kThreads = 8; + constexpr int kRecordsPerThread = 100; + constexpr size_t kPayloadSize = 4096; + + { + auto dump_shader = webgpu::detail::CreateShaderDumpFunction(ToUTF8String(dump_path.native())); + std::barrier start{kThreads}; + std::vector threads; + threads.reserve(kThreads); + for (int thread_id = 0; thread_id < kThreads; ++thread_id) { + threads.emplace_back([&, thread_id] { + start.arrive_and_wait(); + for (int record_id = 0; record_id < kRecordsPerThread; ++record_id) { + const std::string marker = "record-" + std::to_string(thread_id) + "-" + std::to_string(record_id); + dump_shader(marker + "-start\n" + std::string(kPayloadSize, static_cast('a' + thread_id)) + + "\n" + marker + "-end"); + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + } + + std::ifstream dump_file(dump_path); + ASSERT_TRUE(dump_file.is_open()); + const std::string contents{std::istreambuf_iterator{dump_file}, std::istreambuf_iterator{}}; + for (int thread_id = 0; thread_id < kThreads; ++thread_id) { + for (int record_id = 0; record_id < kRecordsPerThread; ++record_id) { + const std::string marker = "record-" + std::to_string(thread_id) + "-" + std::to_string(record_id); + const std::string record = marker + "-start\n" + + std::string(kPayloadSize, static_cast('a' + thread_id)) + + "\n" + marker + "-end\n"; + const size_t position = contents.find(record); + ASSERT_NE(position, std::string::npos) << "Missing or corrupted shader dump record " << marker; + EXPECT_EQ(contents.find(record, position + 1), std::string::npos) + << "Duplicate shader dump record " << marker; + } + } +} + +// DIAGNOSTIC (disabled by default): keeps N sessions over an identical model alive and runs them +// round-robin on one thread, then holds so an external sampler can read steady-state GPU memory. +// This is the shape that distinguishes a context-wide buffer pool (steady state ~= one session's +// high-water mark) from per-session pools (~= N times that). +// +// onnxruntime_provider_test.exe --gtest_also_run_disabled_tests \ +// --gtest_filter=WebGpuPoolMemory.DISABLED_MultiSessionSameShape +// +// Tunable via environment variables: ORT_DIAG_SESSIONS, ORT_DIAG_MB, ORT_DIAG_ITERS, +// ORT_DIAG_HOLD_MS, ORT_DIAG_VARY (1 = give each session a different tensor size), +// ORT_DIAG_CACHE (storage buffer cache mode; defaults to "bucket", the product default - +// note DefaultWebGpuExecutionProvider() disables the storage cache, which would hide the +// very behavior this test measures). +TEST(WebGpuPoolMemory, DISABLED_MultiSessionSameShape) { + if (DefaultWebGpuExecutionProvider() == nullptr) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + auto env_int = [](const char* name, int fallback) { + const std::string value = Env::Default().GetEnvironmentVar(name); + return value.empty() ? fallback : std::stoi(value); + }; + + const int num_sessions = env_int("ORT_DIAG_SESSIONS", 4); + const int mb = env_int("ORT_DIAG_MB", 64); + const int iters = env_int("ORT_DIAG_ITERS", 10); + const int hold_ms = env_int("ORT_DIAG_HOLD_MS", 5000); + const bool vary = env_int("ORT_DIAG_VARY", 0) != 0; + std::string cache_mode = Env::Default().GetEnvironmentVar("ORT_DIAG_CACHE"); + if (cache_mode.empty()) { + cache_mode = webgpu::options::kBufferCacheMode_Bucket; + } + constexpr int kChainLen = 8; + + auto make_ep = [&cache_mode]() { + ConfigOptions config_options; + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kStorageBufferCacheMode, + cache_mode.c_str())); + return WebGpuExecutionProviderWithOptions(config_options); + }; + std::cout << "DIAG_POOL storage_cache=" << cache_mode << "\n"; + + std::vector> sessions; + std::vector feeds; + std::vector output_names{"Y"}; + + // Scalar weights keep GPU memory dominated by intermediate tensors rather than initializers. + auto build_scalar_weight_chain = [](int chain_len, int64_t num_elements, std::string& bytes) { + const std::unordered_map domain_to_version{{"", 13}}; + Model model("webgpu_pool_memory", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector(), + DefaultLoggingManager().DefaultLogger()); + Graph& graph = model.MainGraph(); + + ONNX_NAMESPACE::TypeProto float_1d; + float_1d.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + float_1d.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(num_elements); + + ONNX_NAMESPACE::TypeProto float_scalar; + float_scalar.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + float_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + NodeArg* prev = &graph.GetOrCreateNodeArg("X", &float_1d); + for (int i = 0; i < chain_len; ++i) { + const std::string w_name = "w" + std::to_string(i); + ONNX_NAMESPACE::TensorProto w_tensor; + w_tensor.set_name(w_name); + w_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + w_tensor.add_dims(1); + const float scalar = 1.0f + 0.01f * i; + w_tensor.set_raw_data(&scalar, sizeof(float)); + graph.AddInitializedTensor(w_tensor); + + NodeArg* w_arg = &graph.GetOrCreateNodeArg(w_name, &float_scalar); + const std::string out_name = (i == chain_len - 1) ? "Y" : ("h" + std::to_string(i)); + NodeArg* out_arg = &graph.GetOrCreateNodeArg(out_name, &float_1d); + std::vector inputs{prev, w_arg}; + std::vector outputs{out_arg}; + graph.AddNode("op" + std::to_string(i), (i % 2 == 0) ? "Mul" : "Add", "", inputs, outputs); + prev = out_arg; + } + + graph.SetOutputs(std::vector{prev}); + ASSERT_STATUS_OK(graph.Resolve()); + ASSERT_TRUE(model.ToProto().SerializeToString(&bytes)); + }; + + for (int s = 0; s < num_sessions; ++s) { + const int64_t num_elements = + static_cast(mb) * 1024 * 1024 / 4 + (vary ? s * 4 * 1024 * 1024 : 0); + + std::string model_bytes; + ASSERT_NO_FATAL_FAILURE(build_scalar_weight_chain(kChainLen, num_elements, model_bytes)); + + SessionOptions so; + so.session_logid = "webgpu_pool_memory"; + auto session = std::make_unique(so, GetEnvironment()); + ASSERT_STATUS_OK(session->RegisterExecutionProvider(make_ep())); + ASSERT_STATUS_OK(session->Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session->Initialize()); + + std::vector x_values(static_cast(num_elements), 1.0f); + OrtValue x_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], + std::vector{num_elements}, x_values, &x_value); + + sessions.push_back(std::move(session)); + feeds.push_back(NameMLValMap{{"X", x_value}}); + } + + for (int i = 0; i < iters; ++i) { + for (int s = 0; s < num_sessions; ++s) { + std::vector fetches; + ASSERT_STATUS_OK(sessions[s]->Run(RunOptions{}, feeds[s], output_names, &fetches)); + } + } + + std::cout << "DIAG_POOL sessions=" << num_sessions << " mb=" << mb + << " vary=" << (vary ? 1 : 0) << " iters=" << iters << " holding...\n" + << std::flush; + std::this_thread::sleep_for(std::chrono::milliseconds(hold_ms)); +} + +} // namespace test +} // namespace onnxruntime