[webgpu] Enable graph capture (microsoft#24900)

This PR enables graph capture capabilities in the WebGPU provider, which
is similar with jsep one microsoft#18989.

All limitations are similar with JS/CUDA EP:
1. Models with control-flow ops (i.e. If, Loop and Scan ops) are not
supported.
2. Usage of graph capture is limited to models where-in all ops in the
model can be partitioned to the WebGPU EP or CPU EP and no memory copy
between them.
3. Shapes of inputs/outputs cannot change across inference calls.
4. IOBinding is required. And all inputs/outputs are pre-allocated gpu
buffers.

When users use graph capture feature, we suppose they will do some
pre-process and post-process for the inference's inputs and outputs in
order to keep the whole pipeline on GPU to avoid some unnecessary cpu to
gpu or gpu to cpu copying. The usage will be like below:
```
// Initialize Dawn
{
  // 1. Create Dawn instance
  ...
  instance = wgpu::CreateInstance(&instanceDescriptor); 
  // 2.  Create the adapter
  ...
  instance.RequestAdapter
  // 3. Create device from adapter
  ...
  adapter.RequestDevice
}

// Create session options
webgpu_options_ = std::make_unique<Ort::SessionOptions>();
std::unordered_map<std::string, std::string> provider_options;
provider_options["dawnProcTable"] = std::to_string(reinterpret_cast<size_t>(&dawn::native::GetProcs()));
provider_options["webgpuInstance"] = std::to_string(reinterpret_cast<size_t>(instance_.Get()));
provider_options["webgpuDevice"] = std::to_string(reinterpret_cast<size_t>(device_.Get()));
provider_options["deviceId"] = "1";
provider_options["enableGraphCapture"] = "1";
// add WebGPU provider
webgpu_options_->AppendExecutionProvider("WebGPU", provider_options);
...
// create webgpu session
webgpu_session_ = std::make_unique<Ort::Session>(*env_, model_path_.c_str(), *webgpu_options_);
...
Ort::MemoryInfo memory_info_gpu("WebGPU_Buffer", OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
Ort::Allocator allocator(*webgpu_session_, memory_info_gpu);
auto input_buffer = allocator.GetAllocation(input_tensor_size_ * sizeof(float));
auto output_buffer = allocator.GetAllocation(output_tensor_size_ * sizeof(float));
// Create IoBinding objects
Ort::IoBinding webgpu_binding(*webgpu_session_);
// Upload cpu data to input_buffer or copy gpu buffer to input_buffer
...
// Create an OrtValue tensor backed by data on gpu memory
Ort::Value bound_x = Ort::Value::CreateTensor(memory_info_gpu, reinterpret_cast<float*>(input_buffer.get()), input_tensor_size_,
                input_dims_.data(), input_dims_.size());

Ort::Value bound_y = Ort::Value::CreateTensor(memory_info_gpu, reinterpret_cast<float*>(output_buffer.get()), output_tensor_size_,
                output_dims_.data(), output_dims_.size());
webgpu_binding.BindInput("input", bound_x);
webgpu_binding.BindOutput("output", bound_y);

// Run inference
webgpu_session_->Run(Ort::RunOptions{nullptr}, webgpu_binding); // normal run + capturing
...
// post process output_buffer's content
...
// Update input_buffer's content
...

// Run again
webgpu_session_->Run(Ort::RunOptions{nullptr}, webgpu_binding); // replay()
...
// post process output_buffer's content
...
```

Author: qjia7

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 @@ class DisabledCacheManager : public IBufferCacheManager {
     wgpuBufferRelease(buffer);
   }
 
-  void OnRefresh() override {
+  void OnRefresh(const SessionState& /*session_status*/) override {
     // no-op
   }
 };
@@ -59,7 +59,7 @@ class LazyReleaseCacheManager : public IBufferCacheManager {
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh() override {
+  void OnRefresh(const SessionState& /*session_status*/) override {
     Release();
     pending_buffers_.clear();
   }
@@ -103,7 +103,7 @@ class SimpleCacheManager : public IBufferCacheManager {
     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 @@ class BucketCacheManager : public IBufferCacheManager {
     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 @@ class BucketCacheManager : public IBufferCacheManager {
   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 @@ std::unique_ptr<IBufferCacheManager> CreateBufferCacheManager(BufferCacheMode ca
       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 @@ std::ostream& operator<<(std::ostream& os, BufferCacheMode mode) {
     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 @@ BufferManager::BufferManager(WebGpuContext& context, BufferCacheMode storage_buf
       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 @@ void BufferManager::Upload(void* src, WGPUBuffer dst, size_t size) {
   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 @@ void BufferManager::MemCpy(WGPUBuffer src, WGPUBuffer dst, size_t size) {
   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 @@ WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage) {
   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 @@ WGPUBuffer BufferManager::CreateUMA(size_t size, wgpu::BufferUsage usage) {
   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 @@ void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) {
   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 @@ void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) {
 
   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 {