[webgpu] Enable graph capture

onnxruntime/core/providers/webgpu/allocator.cc

@@ -3,7 +3,7 @@
 
 #include "core/framework/session_state.h"
 #include "core/providers/webgpu/allocator.h"
-#include "core/providers/webgpu/webgpu_context.h"
+#include "core/providers/webgpu/buffer_manager.h"
 
 namespace onnxruntime {
 namespace webgpu {
@@ -15,18 +15,17 @@ void* GpuBufferAllocator::Alloc(size_t size) {
 
   stats_.num_allocs++;
 
-#if !defined(__wasm__)
-  if (!session_initialized_ && context_.DeviceHasFeature(wgpu::FeatureName::BufferMapExtendedUsages)) {
-    return context_.BufferManager().CreateUMA(size);
+  // Check if the buffer manager supports UMA and we're not yet in an initialized session
+  if (!session_initialized_ && buffer_manager_.SupportsUMA()) {
+    return buffer_manager_.CreateUMA(size);
   }
-#endif  // !defined(__wasm__)
 
-  return context_.BufferManager().Create(size);
+  return buffer_manager_.Create(size);
 }
 
 void GpuBufferAllocator::Free(void* p) {
   if (p != nullptr) {
-    context_.BufferManager().Release(static_cast<WGPUBuffer>(p));
+    buffer_manager_.Release(static_cast<WGPUBuffer>(p));
     stats_.num_allocs--;
   }
 }

onnxruntime/core/providers/webgpu/allocator.h

@@ -9,27 +9,29 @@
 namespace onnxruntime {
 namespace webgpu {
 
-class WebGpuContext;
+class BufferManager;
 
 class GpuBufferAllocator : public IAllocator {
  public:
-  GpuBufferAllocator(const WebGpuContext& context)
+  GpuBufferAllocator(const BufferManager& buffer_manager)
       : IAllocator(
             OrtMemoryInfo(WEBGPU_BUFFER, OrtAllocatorType::OrtDeviceAllocator,
                           OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NONE, 0),
                           OrtMemTypeDefault)),
-        context_{context} {
+        buffer_manager_{buffer_manager} {
   }
 
   virtual void* Alloc(size_t size) override;
   virtual void Free(void* p) override;
   void GetStats(AllocatorStats* stats) override;
-
   void OnSessionInitializationEnd();
 
+  // Return the associated BufferManager
+  const BufferManager& GetBufferManager() const { return buffer_manager_; }
+
  private:
   AllocatorStats stats_;
-  const WebGpuContext& context_;
+  const BufferManager& buffer_manager_;
   bool session_initialized_ = false;
 };
 

onnxruntime/core/providers/webgpu/buffer_manager.cc

@@ -37,7 +37,7 @@
     wgpuBufferRelease(buffer);
   }
 
-  void OnRefresh() override {
+  void OnRefresh(const SessionState& /*session_status*/) override {
     // no-op
   }
 };
@@ -59,7 +59,7 @@
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh() override {
+  void OnRefresh(const SessionState& /*session_status*/) override {
     Release();
     pending_buffers_.clear();
   }
@@ -103,7 +103,7 @@
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh() override {
+  void OnRefresh(const SessionState& /*session_status*/) override {
     for (auto& buffer : pending_buffers_) {
       buffers_[static_cast<size_t>(wgpuBufferGetSize(buffer))].emplace_back(buffer);
     }
@@ -196,12 +196,9 @@
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh() override {
-    // TODO: consider graph capture. currently not supported
-
+  void OnRefresh(const SessionState& /*session_status*/) override {
     for (auto& buffer : pending_buffers_) {
       auto buffer_size = static_cast<size_t>(wgpuBufferGetSize(buffer));
-
       auto it = buckets_.find(buffer_size);
       if (it != buckets_.end() && it->second.size() < buckets_limit_[buffer_size]) {
         it->second.emplace_back(buffer);
@@ -249,6 +246,155 @@
   std::vector<size_t> buckets_keys_;
 };
 
+class GraphCacheManager : public IBufferCacheManager {
+ public:
+  GraphCacheManager() : buckets_limit_{BUCKET_DEFAULT_LIMIT_TABLE} {
+    Initialize();
+  }
+  GraphCacheManager(std::unordered_map<size_t, size_t>&& buckets_limit) : buckets_limit_{buckets_limit} {
+    Initialize();
+  }
+
+  size_t CalculateBufferSize(size_t request_size) override {
+    // binary serch size
+    auto it = std::lower_bound(buckets_keys_.begin(), buckets_keys_.end(), request_size);
+    if (it == buckets_keys_.end()) {
+      return NormalizeBufferSize(request_size);
+    } else {
+      return *it;
+    }
+  }
+
+  WGPUBuffer TryAcquireCachedBuffer(size_t buffer_size) override {
+    auto it = buckets_.find(buffer_size);
+    if (it != buckets_.end() && !it->second.empty()) {
+      auto buffer = it->second.back();
+      it->second.pop_back();
+      return buffer;
+    }
+    return nullptr;
+  }
+
+  void RegisterBuffer(WGPUBuffer /*buffer*/, size_t /*request_size*/) override {
+    // no-op
+  }
+
+  void ReleaseBuffer(WGPUBuffer buffer) override {
+    pending_buffers_.emplace_back(buffer);
+  }
+
+  void OnRefresh(const SessionState& /*session_status*/) override {
+    // Initialize buckets if they don't exist yet
+    if (buckets_.empty()) {
+      for (const auto& pair : buckets_limit_) {
+        buckets_.emplace(pair.first, std::vector<WGPUBuffer>());
+      }
+    }
+
+    for (auto& buffer : pending_buffers_) {
+      auto buffer_size = static_cast<size_t>(wgpuBufferGetSize(buffer));
+      auto it = buckets_.find(buffer_size);
+      if (it != buckets_.end()) {
+        it->second.emplace_back(buffer);
+      } else {
+        // insert a new bucket if it doesn't exist
+        buckets_[buffer_size] = std::vector<WGPUBuffer>{buffer};
+      }
+    }
+
+    pending_buffers_.clear();
+  }
+
+  ~GraphCacheManager() {
+    for (auto& buffer : pending_buffers_) {
+      wgpuBufferRelease(buffer);
+    }
+    for (auto& pair : buckets_) {
+      for (auto& buffer : pair.second) {
+        wgpuBufferRelease(buffer);
+      }
+    }
+  }
+
+ protected:
+  void Initialize() {
+    buckets_keys_.reserve(buckets_limit_.size());
+    for (const auto& pair : buckets_limit_) {
+      buckets_keys_.push_back(pair.first);
+    }
+    std::sort(buckets_keys_.begin(), buckets_keys_.end());
+
+#ifndef NDEBUG  // if debug build
+    ORT_ENFORCE(std::all_of(buckets_keys_.begin(), buckets_keys_.end(), [](size_t size) { return size % 16 == 0; }),
+                "Bucket sizes must be multiples of 16.");
+
+    for (size_t i = 1; i < buckets_keys_.size(); ++i) {
+      ORT_ENFORCE(buckets_keys_[i] > buckets_keys_[i - 1], "Bucket sizes must be in increasing order.");
+    }
+#endif
+  }
+  std::unordered_map<size_t, size_t> buckets_limit_;
+  std::unordered_map<size_t, std::vector<WGPUBuffer>> buckets_;
+  std::vector<WGPUBuffer> pending_buffers_;
+  std::vector<size_t> buckets_keys_;
+};
+
+class GraphSimpleCacheManager : public IBufferCacheManager {
+  size_t CalculateBufferSize(size_t request_size) override {
+    return NormalizeBufferSize(request_size);
+  }
+
+  WGPUBuffer TryAcquireCachedBuffer(size_t buffer_size) override {
+    auto it = buffers_.find(buffer_size);
+    if (it != buffers_.end() && !it->second.empty()) {
+      auto buffer = it->second.back();
+      it->second.pop_back();
+      return buffer;
+    }
+
+    return nullptr;
+  }
+
+  void RegisterBuffer(WGPUBuffer /*buffer*/, size_t /*request_size*/) override {
+    // no-op
+  }
+
+  void ReleaseBuffer(WGPUBuffer buffer) override {
+    pending_buffers_.emplace_back(buffer);
+  }
+
+  void OnRefresh(const SessionState& session_status) override {
+    for (auto& buffer : pending_buffers_) {
+      if (session_status == SessionState::Default) {
+        buffers_[static_cast<size_t>(wgpuBufferGetSize(buffer))].emplace_back(buffer);
+      } else {
+        captured_buffers_.emplace_back(buffer);
+      }
+    }
+    pending_buffers_.clear();
+  }
+
+ public:
+  ~GraphSimpleCacheManager() {
+    for (auto& buffer : pending_buffers_) {
+      wgpuBufferRelease(buffer);
+    }
+    for (auto& pair : buffers_) {
+      for (auto& buffer : pair.second) {
+        wgpuBufferRelease(buffer);
+      }
+    }
+    for (auto& buffer : captured_buffers_) {
+      wgpuBufferRelease(buffer);
+    }
+  }
+
+ protected:
+  std::map<size_t, std::vector<WGPUBuffer>> buffers_;
+  std::vector<WGPUBuffer> pending_buffers_;
+  std::vector<WGPUBuffer> captured_buffers_;
+};
+
 std::unique_ptr<IBufferCacheManager> CreateBufferCacheManager(BufferCacheMode cache_mode) {
   switch (cache_mode) {
     case BufferCacheMode::Disabled:
@@ -259,6 +405,10 @@
       return std::make_unique<SimpleCacheManager>();
     case BufferCacheMode::Bucket:
       return std::make_unique<BucketCacheManager>();
+    case BufferCacheMode::Graph:
+      return std::make_unique<GraphCacheManager>();
+    case BufferCacheMode::GraphSimple:
+      return std::make_unique<GraphSimpleCacheManager>();
     default:
       ORT_NOT_IMPLEMENTED("Unsupported buffer cache mode");
   }
@@ -278,6 +428,12 @@
     case BufferCacheMode::Bucket:
       os << "Bucket";
       break;
+    case BufferCacheMode::Graph:
+      os << "Graph";
+      break;
+    case BufferCacheMode::GraphSimple:
+      os << "GraphSimple";
+      break;
     default:
       os << "Unknown(" << static_cast<int>(mode) << ")";
   }
@@ -292,7 +448,7 @@
       default_cache_{CreateBufferCacheManager(BufferCacheMode::Disabled)} {
 }
 
-void BufferManager::Upload(void* src, WGPUBuffer dst, size_t size) {
+void BufferManager::Upload(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) {
@@ -317,10 +473,10 @@
   auto& command_encoder = context_.GetCommandEncoder();
   context_.EndComputePass();
   command_encoder.CopyBufferToBuffer(staging_buffer, 0, dst, 0, buffer_size);
-  context_.Flush();
+  context_.Flush(*this);
 }
 
-void BufferManager::MemCpy(WGPUBuffer src, WGPUBuffer dst, size_t size) {
+void BufferManager::MemCpy(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);
@@ -337,7 +493,7 @@
   command_encoder.CopyBufferToBuffer(src, 0, dst, 0, buffer_size);
 }
 
-WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage) {
+WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage) const {
   auto& cache = GetCacheManager(usage);
   auto buffer_size = cache.CalculateBufferSize(size);
 
@@ -358,7 +514,7 @@
   return buffer;
 }
 
-WGPUBuffer BufferManager::CreateUMA(size_t size, wgpu::BufferUsage usage) {
+WGPUBuffer BufferManager::CreateUMA(size_t size, wgpu::BufferUsage usage) const {
   ORT_ENFORCE(usage & wgpu::BufferUsage::Storage, "UMA buffer must be a storage buffer.");
   auto& cache = GetCacheManager(usage);
   auto buffer_size = cache.CalculateBufferSize(size);
@@ -378,12 +534,21 @@
   return buffer;
 }
 
-void BufferManager::Release(WGPUBuffer buffer) {
+bool BufferManager::SupportsUMA() const {
+#if !defined(__wasm__)
+  // Check if the device supports the BufferMapExtendedUsages feature
+  return context_.DeviceHasFeature(wgpu::FeatureName::BufferMapExtendedUsages);
+#else
+  return false;
+#endif  // !defined(__wasm__)
+}
+
+void BufferManager::Release(WGPUBuffer buffer) const {
   EnforceBufferUnmapped(context_, buffer);
   GetCacheManager(buffer).ReleaseBuffer(buffer);
 }
 
-void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) {
+void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) const {
   EnforceBufferUnmapped(context_, src);
   auto buffer_size = NormalizeBufferSize(size);
 
@@ -395,7 +560,7 @@
   auto& command_encoder = context_.GetCommandEncoder();
   context_.EndComputePass();
   command_encoder.CopyBufferToBuffer(src, 0, staging_buffer, 0, buffer_size);
-  context_.Flush();
+  context_.Flush(*this);
 
   // TODO: revise wait in whole project
 
@@ -405,13 +570,14 @@
 
   auto mapped_data = staging_buffer.GetConstMappedRange();
   memcpy(dst, mapped_data, size);
+  staging_buffer.Unmap();
 }
 
-void BufferManager::RefreshPendingBuffers() {
-  storage_cache_->OnRefresh();
-  uniform_cache_->OnRefresh();
-  query_resolve_cache_->OnRefresh();
-  default_cache_->OnRefresh();
+void BufferManager::RefreshPendingBuffers(const SessionState& session_status) const {
+  storage_cache_->OnRefresh(session_status);
+  uniform_cache_->OnRefresh(session_status);
+  query_resolve_cache_->OnRefresh(session_status);
+  default_cache_->OnRefresh(session_status);
 }
 
 IBufferCacheManager& BufferManager::GetCacheManager(wgpu::BufferUsage usage) const {