[WebGPU] Support concurrent use with explicit Session streams

cmake/onnxruntime_unittests.cmake

@@ -534,6 +534,14 @@ if (onnxruntime_USE_CUDA AND onnxruntime_BUILD_CUDA_EP_AS_PLUGIN AND NOT onnxrun
   list(APPEND onnxruntime_test_providers_src ${onnxruntime_test_providers_cuda_plugin_src})
 endif()
 
+if (onnxruntime_USE_WEBGPU AND onnxruntime_USE_EP_API_ADAPTERS AND NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_REDUCED_OPS_BUILD)
+  file(GLOB onnxruntime_test_providers_webgpu_plugin_src CONFIGURE_DEPENDS
+    "${TEST_SRC_DIR}/providers/webgpu/plugin/*.cc"
+    "${TEST_SRC_DIR}/providers/webgpu/plugin/*.h"
+  )
+  list(APPEND onnxruntime_test_providers_src ${onnxruntime_test_providers_webgpu_plugin_src})
+endif()
+
 if (onnxruntime_USE_CANN)
   file(GLOB_RECURSE onnxruntime_test_providers_cann_src CONFIGURE_DEPENDS
     "${TEST_SRC_DIR}/providers/cann/*"
@@ -661,6 +669,10 @@ if(onnxruntime_USE_CUDA AND onnxruntime_BUILD_CUDA_EP_AS_PLUGIN)
   list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_cuda_plugin)
 endif()
 
+if(onnxruntime_USE_WEBGPU AND onnxruntime_USE_EP_API_ADAPTERS)
+  list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_webgpu)
+endif()
+
 if(onnxruntime_USE_CANN)
   list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_cann)
 endif()
@@ -1483,6 +1495,11 @@ block()
       ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP=1)
   endif()
 
+  if (onnxruntime_USE_WEBGPU AND onnxruntime_USE_EP_API_ADAPTERS)
+    target_compile_definitions(onnxruntime_provider_test PRIVATE
+      ORT_UNIT_TEST_HAS_WEBGPU_PLUGIN_EP=1)
+  endif()
+
   if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_PLUGIN)
     target_link_libraries(onnxruntime_provider_test PRIVATE
       CUDA::cudart
@@ -2725,11 +2742,6 @@ if (onnxruntime_BUILD_SHARED_LIB AND
   file(GLOB onnxruntime_autoep_test_SRC "${ONNXRUNTIME_AUTOEP_TEST_SRC_DIR}/*.h"
                                         "${ONNXRUNTIME_AUTOEP_TEST_SRC_DIR}/*.cc")
 
-  if (NOT onnxruntime_USE_WEBGPU OR NOT onnxruntime_USE_EP_API_ADAPTERS)
-    list(REMOVE_ITEM onnxruntime_autoep_test_SRC
-         "${ONNXRUNTIME_AUTOEP_TEST_SRC_DIR}/test_webgpu_allocators.cc")
-  endif()
-
   set(onnxruntime_autoep_test_LIBS onnxruntime_mocked_allocator ${ONNXRUNTIME_TEST_LIBS} onnxruntime_test_utils
                                    onnx_proto onnx ${onnxruntime_EXTERNAL_LIBRARIES})
 
@@ -2761,12 +2773,6 @@ if (onnxruntime_BUILD_SHARED_LIB AND
           LIBS ${onnxruntime_autoep_test_LIBS}
           DEPENDS ${all_dependencies} example_plugin_ep example_plugin_ep_virt_gpu example_plugin_ep_kernel_registry
   )
-
-  if (onnxruntime_USE_WEBGPU AND onnxruntime_USE_EP_API_ADAPTERS)
-    # The WebGPU plugin is loaded at test-run time, so ensure a focused auto-EP test build produces it and its
-    # co-located runtime dependencies.
-    add_dependencies(onnxruntime_autoep_test onnxruntime_providers_webgpu)
-  endif()
 endif()
 
 if (onnxruntime_BUILD_SHARED_LIB AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" AND NOT onnxruntime_MINIMAL_BUILD)

onnxruntime/core/providers/webgpu/allocator.cc

@@ -2,17 +2,19 @@
 // Licensed under the MIT License.
 
 #include <memory>
+#include <mutex>
 #include <utility>
 
-#include "core/framework/session_state.h"
 #include "core/providers/webgpu/allocator.h"
 #include "core/providers/webgpu/buffer_manager.h"
+#include "core/providers/webgpu/webgpu_context.h"
 
 namespace onnxruntime {
 namespace webgpu {
 
 GpuBufferAllocator::GpuBufferAllocator(
     std::function<const BufferManager&()> buffer_manager_getter,
+    std::function<CommandRecordingState&()> recording_getter,
     bool is_read_only_allocator,
     std::function<bool()> should_submit_zero_initialize)
     : IAllocator(
@@ -22,33 +24,50 @@ GpuBufferAllocator::GpuBufferAllocator(
                         WebGpuDevice,
                         OrtMemTypeDefault)),
       buffer_manager_getter_{std::move(buffer_manager_getter)},
+      recording_getter_{std::move(recording_getter)},
       should_submit_zero_initialize_{std::move(should_submit_zero_initialize)},
       mapped_at_creation_{is_read_only_allocator && buffer_manager_getter_().SupportsUMA()},
       initialize_to_zero_{!is_read_only_allocator} {
 }
 
+// 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.
 void* GpuBufferAllocator::Alloc(size_t size) {
+  auto& recording = recording_getter_();
+  std::lock_guard<std::recursive_mutex> lock{recording.mutex};
+  return Allocate(size, should_submit_zero_initialize_ && should_submit_zero_initialize_());
+}
+
+void* GpuBufferAllocator::Allocate(size_t size, bool submit_zero_initialize) {
   if (size == 0) {
     return nullptr;
   }
 
+  auto& recording = recording_getter_();
+  std::lock_guard<std::recursive_mutex> 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;
 
-  const bool submit_zero_initialize = should_submit_zero_initialize_ && should_submit_zero_initialize_();
-  return buffer_manager_getter_().Create(size, usage, initialize_to_zero_, submit_zero_initialize);
+  return buffer_manager_getter_().Create(recording, size, usage, initialize_to_zero_,
+                                         submit_zero_initialize);
 }
 
 void GpuBufferAllocator::Free(void* p) {
   if (p != nullptr) {
-    buffer_manager_getter_().Release(static_cast<WGPUBuffer>(p));
+    auto& recording = recording_getter_();
+    std::lock_guard<std::recursive_mutex> lock{recording.mutex};
+    buffer_manager_getter_().Release(recording, static_cast<WGPUBuffer>(p));
     stats_.num_allocs--;
   }
 }
 
 void GpuBufferAllocator::GetStats(AllocatorStats* stats) {
+  auto& recording = recording_getter_();
+  std::lock_guard<std::recursive_mutex> lock{recording.mutex};
   *stats = stats_;
 }
 
@@ -70,12 +89,14 @@ void WebGpuNoOpAllocator::Free(void* /*p*/) {
 
 AllocatorPtr CreateWebGpuAllocator(bool device_free,
                                    std::function<const BufferManager&()> buffer_manager_getter,
+                                   std::function<CommandRecordingState&()> recording_getter,
                                    bool is_read_only_allocator,
                                    std::function<bool()> should_submit_zero_initialize) {
   if (device_free) {
     return std::make_shared<WebGpuNoOpAllocator>(is_read_only_allocator);
   }
-  return std::make_shared<GpuBufferAllocator>(std::move(buffer_manager_getter), is_read_only_allocator,
+  return std::make_shared<GpuBufferAllocator>(std::move(buffer_manager_getter), std::move(recording_getter),
+                                              is_read_only_allocator,
                                               std::move(should_submit_zero_initialize));
 }
 

onnxruntime/core/providers/webgpu/allocator.h

@@ -12,28 +12,43 @@ namespace onnxruntime {
 namespace webgpu {
 
 class BufferManager;
+struct CommandRecordingState;
 
 inline constexpr OrtDevice WebGpuDevice{OrtDevice::GPU,
                                         OrtDevice::MemType::DEFAULT,
                                         OrtDevice::VendorIds::NONE,
                                         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.
 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.
   GpuBufferAllocator(std::function<const BufferManager&()> buffer_manager_getter,
+                     std::function<CommandRecordingState&()> recording_getter,
                      bool is_read_only_allocator,
                      std::function<bool()> should_submit_zero_initialize = {});
 
   virtual void* Alloc(size_t size) override;
   virtual void Free(void* p) override;
   void GetStats(AllocatorStats* stats) override;
 
+#if defined(ORT_USE_EP_API_ADAPTERS)
+  bool IsStreamAware() const override { return true; }
+  void* AllocOnStream(size_t size, Stream* stream) override;
+#endif
+
  private:
+  void* Allocate(size_t size, bool submit_zero_initialize);
   AllocatorStats stats_;
   std::function<const BufferManager&()> buffer_manager_getter_;
+  std::function<CommandRecordingState&()> recording_getter_;
   std::function<bool()> should_submit_zero_initialize_;
   bool mapped_at_creation_;
   // Cached writable buffers are cleared explicitly by BufferManager::Create. Fresh buffers rely on Dawn's
@@ -59,6 +74,7 @@ class WebGpuNoOpAllocator : public IAllocator {
 // allocation ever happens.
 AllocatorPtr CreateWebGpuAllocator(bool device_free,
                                    std::function<const BufferManager&()> buffer_manager_getter,
+                                   std::function<CommandRecordingState&()> recording_getter,
                                    bool is_read_only_allocator,
                                    std::function<bool()> should_submit_zero_initialize = {});
 

onnxruntime/core/providers/webgpu/buffer_manager.cc

@@ -495,7 +495,7 @@
       default_cache_{CreateBufferCacheManager(default_buffer_cache_mode)} {
 }
 
-void BufferManager::Upload(void* src, WGPUBuffer dst, size_t size) const {
+void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const {
   // If the buffer is mapped, we can directly write to it.
   void* mapped_data = wgpuBufferGetMappedRange(dst, 0, WGPU_WHOLE_MAP_SIZE);  // ensure the buffer is mapped
   if (mapped_data) {
@@ -522,14 +522,14 @@
   // shader to write the non-aligned remainder.
   staging_buffer.Unmap();
 
-  ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches());
-  auto& command_encoder = context_.GetCommandEncoder();
-  context_.EndComputePass();
+  ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(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));
+  ORT_THROW_IF_ERROR(context_.Flush(*this, recording));
 }
 
-void BufferManager::MemCpy(WGPUBuffer src, WGPUBuffer dst, size_t size) const {
+void BufferManager::MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const {
   ORT_ENFORCE(src != dst, "Source and destination buffers must be different.");
   EnforceBufferUnmapped(context_, src);
   EnforceBufferUnmapped(context_, dst);
@@ -541,29 +541,35 @@
               "Source and destination buffers must have enough space for the copy operation. src_size=",
               src_size, ", dst_size=", dst_size, ", copy_size=", copy_size, ".");
 
-  ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches());
-  auto& command_encoder = context_.GetCommandEncoder();
-  context_.EndComputePass();
+  ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording));
+  auto& command_encoder = context_.GetCommandEncoder(recording);
+  context_.EndComputePass(recording);
   command_encoder.CopyBufferToBuffer(src, 0, dst, 0, copy_size);
 }
 
-WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage, bool initialize_to_zero,
+WGPUBuffer BufferManager::Create(CommandRecordingState& recording, size_t size, wgpu::BufferUsage usage,
+                                 bool initialize_to_zero,
                                  bool submit_zero_initialize) const {
-  auto& cache = GetCacheManager(usage);
-  auto buffer_size = cache.CalculateBufferSize(size);
-
-  auto buffer = cache.TryAcquireCachedBuffer(buffer_size);
+  size_t buffer_size;
+  WGPUBuffer buffer;
+  {
+    std::lock_guard<std::mutex> lock{mutex_};
+    auto& cache = GetCacheManager(usage);
+    buffer_size = cache.CalculateBufferSize(size);
+    buffer = cache.TryAcquireCachedBuffer(buffer_size);
+  }
   if (buffer) {
     if (initialize_to_zero) {
-      // initialize_to_zero controls whether a cached buffer is cleared. submit_zero_initialize separately controls
-      // whether that clear is submitted before Create returns. Session::Run defers submission to preserve dispatch
-      // batching, while allocations made outside Run submit immediately so subsequent queue work observes the clear.
+      // 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.
       auto buffer_guard = wgpu::Buffer::Acquire(buffer);
-      ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches());
-      context_.EndComputePass();
-      context_.GetCommandEncoder().ClearBuffer(buffer, 0, buffer_size);
+      ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording));
+      context_.EndComputePass(recording);
+      context_.GetCommandEncoder(recording).ClearBuffer(buffer, 0, buffer_size);
       if (submit_zero_initialize) {
-        ORT_THROW_IF_ERROR(context_.Flush(*this));
+        ORT_THROW_IF_ERROR(context_.Flush(*this, recording));
       }
       return buffer_guard.MoveToCHandle();
     }
@@ -581,7 +587,10 @@
 
   ORT_ENFORCE(buffer, "Failed to create GPU buffer: size=", buffer_size, ", usage=", uint64_t(usage), ".");
 
-  cache.RegisterBuffer(buffer, size);
+  {
+    std::lock_guard<std::mutex> lock{mutex_};
+    GetCacheManager(usage).RegisterBuffer(buffer, size);
+  }
   return buffer;
 }
 
@@ -594,15 +603,21 @@
 #endif  // !defined(__wasm__)
 }
 
-void BufferManager::Release(WGPUBuffer buffer) const {
+void BufferManager::Release(CommandRecordingState& recording, WGPUBuffer buffer) const {
   EnforceBufferUnmapped(context_, buffer);
+  if (recording.has_unsubmitted_work) {
+    recording.pending_buffers.emplace_back(wgpu::Buffer::Acquire(buffer));
+    return;
+  }
+
+  std::lock_guard<std::mutex> lock{mutex_};
   GetCacheManager(buffer).ReleaseBuffer(buffer);
 }
 
-void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) const {
+void BufferManager::Download(CommandRecordingState& recording, WGPUBuffer src, void* dst, size_t size) const {
   // Encode pending deferred dispatches before recording the readback; the flush below submits both
   // in order.
-  ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches());
+  ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording));
 
   EnforceBufferUnmapped(context_, src);
   auto buffer_size = NormalizeBufferSize(size);
@@ -612,10 +627,10 @@
   desc.usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead;
 
   auto staging_buffer = context_.Device().CreateBuffer(&desc);
-  auto& command_encoder = context_.GetCommandEncoder();
-  context_.EndComputePass();
+  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));
+  ORT_THROW_IF_ERROR(context_.Flush(*this, recording));
 
   // TODO: revise wait in whole project
 
@@ -647,11 +662,29 @@
   staging_buffer.Unmap();
 }
 
-void BufferManager::RefreshPendingBuffers(GraphCaptureState graph_capture_state) const {
-  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::RefreshPendingBuffers(CommandRecordingState& recording) const {
+  std::lock_guard<std::mutex> lock{mutex_};
+  for (auto& buffer : recording.pending_buffers) {
+    GetCacheManager(buffer.Get()).ReleaseBuffer(buffer.MoveToCHandle());
+  }
+  recording.pending_buffers.clear();
+
+  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);
+}
+
+std::vector<std::pair<size_t, WGPUBuffer>> BufferManager::ExtractCachedBuffers(wgpu::BufferUsage usage) {
+  std::lock_guard<std::mutex> lock{mutex_};
+  return GetCacheManager(usage).ExtractCachedBuffers();
+}
+
+void BufferManager::AbsorbCachedBuffers(
+    wgpu::BufferUsage usage,
+    std::vector<std::pair<size_t, WGPUBuffer>>&& buffers) {
+  std::lock_guard<std::mutex> lock{mutex_};
+  GetCacheManager(usage).AbsorbCachedBuffers(std::move(buffers));
 }
 
 IBufferCacheManager& BufferManager::GetCacheManager(wgpu::BufferUsage usage) const {