[webgpu][dawn API optimization] reduce number of calls to buffer APIs (microsoft#24315)

### Description

This PR is one of a series of changes for optimization of Dawn API
usage. See microsoft#24281

Reduce the calls to wgpuBufferAddRef and wgpuBufferRelease (part 1).

Author: fs-eire

onnxruntime/core/providers/webgpu/allocator.cc

@@ -13,15 +13,15 @@ void* GpuBufferAllocator::Alloc(size_t size) {
     return nullptr;
   }
 
-  WGPUBuffer buffer;
-  if (!session_initialized_ && context_.SupportsBufferMapExtendedUsages()) {
-    buffer = context_.BufferManager().CreateUMA(size);
-  } else {
-    buffer = context_.BufferManager().Create(size);
+  stats_.num_allocs++;
+
+#if !defined(__wasm__)
+  if (!session_initialized_ && context_.DeviceHasFeature(wgpu::FeatureName::BufferMapExtendedUsages)) {
+    return context_.BufferManager().CreateUMA(size);
   }
+#endif  // !defined(__wasm__)
 
-  stats_.num_allocs++;
-  return buffer;
+  return context_.BufferManager().Create(size);
 }
 
 void GpuBufferAllocator::Free(void* p) {

onnxruntime/core/providers/webgpu/buffer_manager.cc

@@ -56,14 +56,27 @@ class LazyReleaseCacheManager : public IBufferCacheManager {
   }
 
   void ReleaseBuffer(WGPUBuffer buffer) override {
-    pending_buffers_.emplace_back(wgpu::Buffer::Acquire(buffer));
+    pending_buffers_.emplace_back(buffer);
   }
 
   void OnRefresh() override {
+    Release();
     pending_buffers_.clear();
   }
 
-  std::vector<wgpu::Buffer> pending_buffers_;
+ public:
+  ~LazyReleaseCacheManager() {
+    Release();
+  }
+
+ protected:
+  void Release() {
+    for (auto& buffer : pending_buffers_) {
+      wgpuBufferRelease(buffer);
+    }
+  }
+
+  std::vector<WGPUBuffer> pending_buffers_;
 };
 
 class SimpleCacheManager : public IBufferCacheManager {
@@ -74,7 +87,7 @@ class SimpleCacheManager : public IBufferCacheManager {
   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().MoveToCHandle();
+      auto buffer = it->second.back();
       it->second.pop_back();
       return buffer;
     }
@@ -87,18 +100,31 @@ class SimpleCacheManager : public IBufferCacheManager {
   }
 
   void ReleaseBuffer(WGPUBuffer buffer) override {
-    pending_buffers_.emplace_back(wgpu::Buffer::Acquire(buffer));
+    pending_buffers_.emplace_back(buffer);
   }
 
   void OnRefresh() override {
     for (auto& buffer : pending_buffers_) {
-      buffers_[static_cast<size_t>(buffer.GetSize())].emplace_back(std::move(buffer));
+      buffers_[static_cast<size_t>(wgpuBufferGetSize(buffer))].emplace_back(buffer);
     }
     pending_buffers_.clear();
   }
 
-  std::map<size_t, std::vector<wgpu::Buffer>> buffers_;
-  std::vector<wgpu::Buffer> pending_buffers_;
+ public:
+  ~SimpleCacheManager() {
+    for (auto& buffer : pending_buffers_) {
+      wgpuBufferRelease(buffer);
+    }
+    for (auto& pair : buffers_) {
+      for (auto& buffer : pair.second) {
+        wgpuBufferRelease(buffer);
+      }
+    }
+  }
+
+ protected:
+  std::map<size_t, std::vector<WGPUBuffer>> buffers_;
+  std::vector<WGPUBuffer> pending_buffers_;
 };
 
 // TODO: maybe use different bucket size for storage and uniform buffers?
@@ -155,7 +181,7 @@ class BucketCacheManager : public IBufferCacheManager {
   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().MoveToCHandle();
+      auto buffer = it->second.back();
       it->second.pop_back();
       return buffer;
     }
@@ -167,31 +193,44 @@ class BucketCacheManager : public IBufferCacheManager {
   }
 
   void ReleaseBuffer(WGPUBuffer buffer) override {
-    pending_buffers_.emplace_back(wgpu::Buffer::Acquire(buffer));
+    pending_buffers_.emplace_back(buffer);
   }
 
   void OnRefresh() override {
     // TODO: consider graph capture. currently not supported
 
     for (auto& buffer : pending_buffers_) {
-      auto buffer_size = static_cast<size_t>(buffer.GetSize());
+      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(std::move(buffer));
+        it->second.emplace_back(buffer);
+      } else {
+        wgpuBufferRelease(buffer);
       }
     }
 
     pending_buffers_.clear();
   }
 
+  ~BucketCacheManager() {
+    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());
     buckets_.reserve(buckets_limit_.size());
     for (const auto& pair : buckets_limit_) {
       buckets_keys_.push_back(pair.first);
-      buckets_.emplace(pair.first, std::vector<wgpu::Buffer>());
+      buckets_.emplace(pair.first, std::vector<WGPUBuffer>());
     }
     std::sort(buckets_keys_.begin(), buckets_keys_.end());
 
@@ -205,8 +244,8 @@ class BucketCacheManager : public IBufferCacheManager {
 #endif
   }
   std::unordered_map<size_t, size_t> buckets_limit_;
-  std::unordered_map<size_t, std::vector<wgpu::Buffer>> buckets_;
-  std::vector<wgpu::Buffer> pending_buffers_;
+  std::unordered_map<size_t, std::vector<WGPUBuffer>> buckets_;
+  std::vector<WGPUBuffer> pending_buffers_;
   std::vector<size_t> buckets_keys_;
 };
 
@@ -255,11 +294,10 @@ BufferManager::BufferManager(WebGpuContext& context, BufferCacheMode storage_buf
 
 void BufferManager::Upload(void* src, WGPUBuffer dst, size_t size) {
   // If the buffer is mapped, we can directly write to it.
-  wgpu::Buffer dst_buffer = dst;
-  auto mapped_data = dst_buffer.GetMappedRange();
+  void* mapped_data = wgpuBufferGetMappedRange(dst, 0, WGPU_WHOLE_MAP_SIZE);  // ensure the buffer is mapped
   if (mapped_data) {
     memcpy(mapped_data, src, size);
-    dst_buffer.Unmap();
+    wgpuBufferUnmap(dst);
     return;
   }
 
@@ -288,17 +326,19 @@ void BufferManager::MemCpy(WGPUBuffer src, WGPUBuffer dst, size_t size) {
   EnforceBufferUnmapped(context_, dst);
 
   auto buffer_size = NormalizeBufferSize(size);
-  ORT_ENFORCE(buffer_size <= wgpuBufferGetSize(src) && buffer_size <= wgpuBufferGetSize(dst),
+  auto src_size = static_cast<size_t>(wgpuBufferGetSize(src));
+  auto dst_size = static_cast<size_t>(wgpuBufferGetSize(dst));
+  ORT_ENFORCE(buffer_size <= src_size && buffer_size <= dst_size,
               "Source and destination buffers must have enough space for the copy operation. src_size=",
-              wgpuBufferGetSize(src), ", dst_size=", wgpuBufferGetSize(dst), ", copy_size=", buffer_size, ".");
+              src_size, ", dst_size=", dst_size, ", copy_size=", buffer_size, ".");
 
   auto& command_encoder = context_.GetCommandEncoder();
   context_.EndComputePass();
   command_encoder.CopyBufferToBuffer(src, 0, dst, 0, buffer_size);
 }
 
 WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage) {
-  auto& cache = GetCacheManager(static_cast<WGPUBufferUsage>(usage));
+  auto& cache = GetCacheManager(usage);
   auto buffer_size = cache.CalculateBufferSize(size);
 
   auto buffer = cache.TryAcquireCachedBuffer(buffer_size);
@@ -310,7 +350,6 @@ WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage) {
   wgpu::BufferDescriptor desc{};
   desc.size = buffer_size;
   desc.usage = usage;
-  // desc.label = std::to_string(xx++).c_str();
   buffer = context_.Device().CreateBuffer(&desc).MoveToCHandle();
 
   ORT_ENFORCE(buffer, "Failed to create GPU buffer: size=", buffer_size, ", usage=", uint64_t(usage), ".");
@@ -320,14 +359,16 @@ WGPUBuffer BufferManager::Create(size_t size, wgpu::BufferUsage usage) {
 }
 
 WGPUBuffer BufferManager::CreateUMA(size_t size, wgpu::BufferUsage usage) {
-  ORT_ENFORCE(usage & wgpu::BufferUsage::Storage, "UMA buffer must have storage usage.");
-  auto& cache = GetCacheManager(static_cast<WGPUBufferUsage>(usage));
+  ORT_ENFORCE(usage & wgpu::BufferUsage::Storage, "UMA buffer must be a storage buffer.");
+  auto& cache = GetCacheManager(usage);
   auto buffer_size = cache.CalculateBufferSize(size);
 
+  // Ensure the buffer is mapped for writing at creation.
+  usage |= wgpu::BufferUsage::MapWrite;
+
   wgpu::BufferDescriptor desc{};
   desc.size = buffer_size;
-  // Ensure the buffer is mapped for writing at creation.
-  desc.usage = usage | wgpu::BufferUsage::MapWrite;
+  desc.usage = usage;
   desc.mappedAtCreation = true;
   auto buffer = context_.Device().CreateBuffer(&desc).MoveToCHandle();
 
@@ -373,20 +414,21 @@ void BufferManager::RefreshPendingBuffers() {
   default_cache_->OnRefresh();
 }
 
-IBufferCacheManager& BufferManager::GetCacheManager(WGPUBufferUsage usage) const {
-  if (usage & WGPUBufferUsage_Storage) {
+IBufferCacheManager& BufferManager::GetCacheManager(wgpu::BufferUsage usage) const {
+  if (usage & wgpu::BufferUsage::Storage) {
     return *storage_cache_;
-  } else if (usage & WGPUBufferUsage_Uniform) {
+  } else if (usage & wgpu::BufferUsage::Uniform) {
     return *uniform_cache_;
-  } else if (usage & WGPUBufferUsage_QueryResolve) {
+  } else if (usage & wgpu::BufferUsage::QueryResolve) {
     return *query_resolve_cache_;
   } else {
     return *default_cache_;
   }
 }
 
 IBufferCacheManager& BufferManager::GetCacheManager(WGPUBuffer buffer) const {
-  return GetCacheManager(wgpuBufferGetUsage(buffer));
+  auto usage = static_cast<wgpu::BufferUsage>(wgpuBufferGetUsage(buffer));
+  return GetCacheManager(usage);
 }
 
 std::unique_ptr<BufferManager> BufferManagerFactory::Create(WebGpuContext& context, BufferCacheMode storage_buffer_cache_mode, BufferCacheMode uniform_buffer_cache_mode, BufferCacheMode query_resolve_buffer_cache_mode) {

onnxruntime/core/providers/webgpu/buffer_manager.h

@@ -70,7 +70,7 @@ class BufferManager {
   void RefreshPendingBuffers();
 
  private:
-  IBufferCacheManager& GetCacheManager(WGPUBufferUsage usage) const;
+  IBufferCacheManager& GetCacheManager(wgpu::BufferUsage usage) const;
   IBufferCacheManager& GetCacheManager(WGPUBuffer buffer) const;
 
   WebGpuContext& context_;

onnxruntime/core/providers/webgpu/webgpu_context.cc

@@ -149,10 +149,6 @@ void WebGpuContext::Initialize(const WebGpuBufferCacheConfig& buffer_cache_confi
       device_features_.insert(supported_features.features[i]);
     }
 
-#if !defined(__wasm__)
-    supports_buffer_map_extended_usages_ = device_.HasFeature(wgpu::FeatureName::BufferMapExtendedUsages);
-#endif
-
     // create buffer manager
     buffer_mgr_ = BufferManagerFactory::Create(*this,
                                                buffer_cache_config.storage.mode,

onnxruntime/core/providers/webgpu/webgpu_context.h

@@ -145,8 +145,6 @@ class WebGpuContext final {
   Status Run(ComputeContext& context, const ProgramBase& program);
   void OnRunEnd();
 
-  bool SupportsBufferMapExtendedUsages() const { return supports_buffer_map_extended_usages_; }
-
  private:
   enum class TimestampQueryType {
     None = 0,
@@ -238,7 +236,6 @@ class WebGpuContext final {
 #if defined(ENABLE_PIX_FOR_WEBGPU_EP)
   std::unique_ptr<WebGpuPIXFrameGenerator> pix_frame_generator_ = nullptr;
 #endif  // ENABLE_PIX_FOR_WEBGPU_EP
-  bool supports_buffer_map_extended_usages_ = false;
 };
 
 }  // namespace webgpu