[webgpu] a few optimizations to graph capture implementation (microsoft#25305)

### Description

1. rename `SessionState` to `GraphCaptureState`, since there is already
one SessionState type in ORT.
2. optimize implementation of `ComputeContext::BufferManager()`

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: fs-eire

onnxruntime/core/providers/webgpu/allocator.h

@@ -26,9 +26,6 @@ class GpuBufferAllocator : public IAllocator {
   void GetStats(AllocatorStats* stats) override;
   void OnSessionInitializationEnd();
 
-  // Return the associated BufferManager
-  const BufferManager& GetBufferManager() const { return buffer_manager_; }
-
  private:
   AllocatorStats stats_;
   const BufferManager& buffer_manager_;

onnxruntime/core/providers/webgpu/buffer_manager.cc

@@ -37,7 +37,7 @@ class DisabledCacheManager : public IBufferCacheManager {
     wgpuBufferRelease(buffer);
   }
 
-  void OnRefresh(const SessionState& /*session_status*/) override {
+  void OnRefresh(GraphCaptureState /*graph_capture_state*/) override {
     // no-op
   }
 };
@@ -59,7 +59,7 @@ class LazyReleaseCacheManager : public IBufferCacheManager {
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh(const SessionState& /*session_status*/) override {
+  void OnRefresh(GraphCaptureState /*graph_capture_state*/) override {
     Release();
     pending_buffers_.clear();
   }
@@ -103,7 +103,7 @@ class SimpleCacheManager : public IBufferCacheManager {
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh(const SessionState& /*session_status*/) override {
+  void OnRefresh(GraphCaptureState /*graph_capture_state*/) override {
     for (auto& buffer : pending_buffers_) {
       buffers_[static_cast<size_t>(wgpuBufferGetSize(buffer))].emplace_back(buffer);
     }
@@ -196,7 +196,7 @@ class BucketCacheManager : public IBufferCacheManager {
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh(const SessionState& /*session_status*/) override {
+  void OnRefresh(GraphCaptureState /*graph_capture_state*/) override {
     for (auto& buffer : pending_buffers_) {
       auto buffer_size = static_cast<size_t>(wgpuBufferGetSize(buffer));
       auto it = buckets_.find(buffer_size);
@@ -283,7 +283,7 @@ class GraphCacheManager : public IBufferCacheManager {
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh(const SessionState& /*session_status*/) override {
+  void OnRefresh(GraphCaptureState /*graph_capture_state*/) override {
     // Initialize buckets if they don't exist yet
     if (buckets_.empty()) {
       for (const auto& pair : buckets_limit_) {
@@ -363,9 +363,9 @@ class GraphSimpleCacheManager : public IBufferCacheManager {
     pending_buffers_.emplace_back(buffer);
   }
 
-  void OnRefresh(const SessionState& session_status) override {
+  void OnRefresh(GraphCaptureState graph_capture_state) override {
     for (auto& buffer : pending_buffers_) {
-      if (session_status == SessionState::Default) {
+      if (graph_capture_state == GraphCaptureState::Default) {
         buffers_[static_cast<size_t>(wgpuBufferGetSize(buffer))].emplace_back(buffer);
       } else {
         captured_buffers_.emplace_back(buffer);
@@ -573,11 +573,11 @@ void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) const {
   staging_buffer.Unmap();
 }
 
-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);
+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);
 }
 
 IBufferCacheManager& BufferManager::GetCacheManager(wgpu::BufferUsage usage) const {

onnxruntime/core/providers/webgpu/buffer_manager.h

@@ -15,7 +15,7 @@ namespace webgpu {
 class WebGpuContext;
 
 // For command capture and replay
-enum class SessionState {
+enum class GraphCaptureState {
   Default,
   Capturing,
   Replaying
@@ -59,7 +59,7 @@ class IBufferCacheManager {
   virtual void ReleaseBuffer(WGPUBuffer buffer) = 0;
 
   // when a stream refresh is requested
-  virtual void OnRefresh(const SessionState& session_status) = 0;
+  virtual void OnRefresh(GraphCaptureState graph_capture_state) = 0;
 };
 
 //
@@ -77,7 +77,7 @@ class BufferManager {
   bool SupportsUMA() const;
   void Release(WGPUBuffer buffer) const;
   void Download(WGPUBuffer src, void* dst, size_t size) const;
-  void RefreshPendingBuffers(const SessionState& session_status) const;
+  void RefreshPendingBuffers(GraphCaptureState graph_capture_state) const;
 
  private:
   IBufferCacheManager& GetCacheManager(wgpu::BufferUsage usage) const;

onnxruntime/core/providers/webgpu/compute_context.cc

@@ -7,12 +7,14 @@
 #include "core/providers/webgpu/webgpu_context.h"
 #include "core/providers/webgpu/allocator.h"
 #include "core/providers/webgpu/buffer_manager.h"
+#include "core/providers/webgpu/webgpu_execution_provider.h"
 
 namespace onnxruntime {
 namespace webgpu {
-ComputeContext::ComputeContext(OpKernelContext& kernel_context)
+ComputeContext::ComputeContext(OpKernelContext& kernel_context, const WebGpuExecutionProvider& ep)
     : webgpu_context_{WebGpuContextFactory::GetContext(kernel_context.GetDeviceId())},
-      kernel_context_{kernel_context} {
+      kernel_context_{kernel_context},
+      ep_{ep} {
 }
 
 void ComputeContext::PushErrorScope() {
@@ -29,10 +31,7 @@ Status ComputeContext::PopErrorScope() {
 }
 
 const webgpu::BufferManager& ComputeContext::BufferManager() const {
-  OrtDevice gpu_device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NONE, 0);
-  AllocatorPtr allocator = kernel_context_.GetAllocator(gpu_device);
-  const GpuBufferAllocator* gpu_allocator = static_cast<const GpuBufferAllocator*>(allocator.get());
-  return gpu_allocator->GetBufferManager();
+  return ep_.BufferManager();
 }
 
 }  // namespace webgpu

onnxruntime/core/providers/webgpu/compute_context.h

@@ -16,6 +16,7 @@
 namespace onnxruntime {
 
 class Tensor;
+class WebGpuExecutionProvider;
 
 namespace webgpu {
 
@@ -24,7 +25,7 @@ class BufferManager;
 
 class ComputeContext {
  public:
-  ComputeContext(OpKernelContext& kernel_context);
+  ComputeContext(OpKernelContext& kernel_context, const WebGpuExecutionProvider& ep);
 
   virtual ~ComputeContext() = default;
 
@@ -145,6 +146,7 @@ class ComputeContext {
  protected:
   WebGpuContext& webgpu_context_;
   OpKernelContext& kernel_context_;
+  const WebGpuExecutionProvider& ep_;
 };
 
 }  // namespace webgpu

onnxruntime/core/providers/webgpu/webgpu_context.cc

@@ -691,8 +691,8 @@ void WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr) {
   }
   auto command_buffer = current_command_encoder_.Finish();
   device_queue_.Submit(1, &command_buffer);
-  if (session_status_ != SessionState::Replaying) {
-    buffer_mgr.RefreshPendingBuffers(session_status_);
+  if (graph_capture_state_ != GraphCaptureState::Replaying) {
+    buffer_mgr.RefreshPendingBuffers(graph_capture_state_);
   }
   current_command_encoder_ = nullptr;
   num_pending_dispatches_ = 0;
@@ -724,7 +724,7 @@ void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& comput
   bind_group_desc.label = {program_artifact.name.data(), program_artifact.name.length()};
 
   auto bind_group = wgpuDeviceCreateBindGroup(Device().Get(), &bind_group_desc);
-  if (session_status_ == SessionState::Capturing) {
+  if (graph_capture_state_ == GraphCaptureState::Capturing) {
     external_captured_commands_->push_back({program_artifact.compute_pipeline,
                                             bind_group,
                                             bind_group_layout,
@@ -754,12 +754,12 @@ void WebGpuContext::CaptureBegin(std::vector<webgpu::CapturedCommandInfo>* captu
   // TODO: support profiling with graph capture.
   ORT_ENFORCE(!is_profiling_, "profiling is not supported yet under graph capture mode");
 
-  session_status_ = SessionState::Capturing;
+  graph_capture_state_ = GraphCaptureState::Capturing;
 }
 
 void WebGpuContext::Replay(const std::vector<webgpu::CapturedCommandInfo>& captured_commands, const webgpu::BufferManager& buffer_manager) {
   LOGS_DEFAULT(VERBOSE) << "Replay with external storage";
-  session_status_ = SessionState::Replaying;
+  graph_capture_state_ = GraphCaptureState::Replaying;
   // Replay all captured commands from the provided vector
   const size_t command_count = captured_commands.size();
   for (size_t i = 0; i < command_count; ++i) {
@@ -784,13 +784,13 @@ void WebGpuContext::Replay(const std::vector<webgpu::CapturedCommandInfo>& captu
   // Flush any remaining commands
   Flush(buffer_manager);
 
-  session_status_ = SessionState::Default;
+  graph_capture_state_ = GraphCaptureState::Default;
 }
 
 void WebGpuContext::CaptureEnd() {
   LOGS_DEFAULT(VERBOSE) << "CaptureEnd";
 
-  session_status_ = SessionState::Default;
+  graph_capture_state_ = GraphCaptureState::Default;
   external_captured_commands_ = nullptr;
 }
 

onnxruntime/core/providers/webgpu/webgpu_context.h

@@ -254,7 +254,7 @@ class WebGpuContext final {
   uint64_t gpu_timestamp_offset_ = 0;
   bool is_profiling_ = false;
   bool preserve_device_;
-  SessionState session_status_{SessionState::Default};
+  GraphCaptureState graph_capture_state_{GraphCaptureState::Default};
 
   // External vector to store captured commands, owned by EP
   std::vector<webgpu::CapturedCommandInfo>* external_captured_commands_ = nullptr;

onnxruntime/core/providers/webgpu/webgpu_kernel.h

@@ -9,6 +9,8 @@
 #include "core/framework/op_kernel.h"
 
 namespace onnxruntime {
+
+class WebGpuExecutionProvider;
 namespace webgpu {
 
 // -----------------------------------------------------------------------
@@ -17,11 +19,12 @@ namespace webgpu {
 class WebGpuKernel : public OpKernel {
  public:
   explicit WebGpuKernel(const OpKernelInfo& info)
-      : OpKernel(info) {
+      : OpKernel(info),
+        ep_(*static_cast<const WebGpuExecutionProvider*>(info.GetExecutionProvider())) {
   }
 
   Status Compute(OpKernelContext* p_op_kernel_context) const override {
-    ComputeContext context{*p_op_kernel_context};
+    ComputeContext context{*p_op_kernel_context, ep_};
 
     context.PushErrorScope();
     Status s = ComputeInternal(context);
@@ -31,6 +34,9 @@ class WebGpuKernel : public OpKernel {
   }
 
   virtual Status ComputeInternal(ComputeContext& context) const = 0;
+
+ private:
+  const WebGpuExecutionProvider& ep_;
 };
 
 }  // namespace webgpu