[webgpu] revise class ComputeContext in WebGPU EP

onnxruntime/core/providers/webgpu/compute_context.cc

@@ -1,37 +1,19 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include "core/framework/op_kernel.h"
-
 #include "core/providers/webgpu/compute_context.h"
-#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, const WebGpuExecutionProvider& ep)
-    : webgpu_context_{WebGpuContextFactory::GetContext(kernel_context.GetDeviceId())},
+ComputeContext::ComputeContext(OpKernelContext& kernel_context, const WebGpuExecutionProvider& ep, WebGpuContext& webgpu_context)
+    : webgpu_context_{webgpu_context},
       kernel_context_{kernel_context},
       ep_{ep} {
 }
 
-void ComputeContext::PushErrorScope() {
-  if (webgpu_context_.ValidationMode() >= ValidationMode::Full) {
-    webgpu_context_.PushErrorScope();
-  }
-}
-
-Status ComputeContext::PopErrorScope() {
-  if (webgpu_context_.ValidationMode() >= ValidationMode::Full) {
-    return webgpu_context_.PopErrorScope();
-  }
-  return Status::OK();
-}
-
-const webgpu::BufferManager& ComputeContext::BufferManager() const {
-  return ep_.BufferManager();
+const webgpu::BufferManager& ComputeContext::BufferManagerAccessor::Get(const ComputeContext& context) {
+  return context.ep_.BufferManager();
 }
 
 }  // namespace webgpu

onnxruntime/core/providers/webgpu/compute_context.h

@@ -23,11 +23,22 @@ namespace webgpu {
 class WebGpuContext;
 class BufferManager;
 
-class ComputeContext {
+class ComputeContext final {
  public:
-  ComputeContext(OpKernelContext& kernel_context, const WebGpuExecutionProvider& ep);
+  // Nested accessor class to provide controlled access to BufferManager
+  class BufferManagerAccessor {
+    // access to BufferManager is limited to class WebGpuContext.
+    // This ensures no access to BufferManager from other classes, avoiding
+    // potential misuse.
+    friend class WebGpuContext;
 
-  virtual ~ComputeContext() = default;
+   private:
+    static const webgpu::BufferManager& Get(const ComputeContext& context);
+  };
+
+  ComputeContext(OpKernelContext& kernel_context, const WebGpuExecutionProvider& ep, WebGpuContext& webgpu_context);
+
+  ~ComputeContext() = default;
 
   //
   // Get various information from the context.
@@ -120,33 +131,15 @@ class ComputeContext {
     ORT_THROW_IF_ERROR(kernel_context_.GetTempSpaceAllocator(&allocator));
     return {data_type, std::forward<TensorShapeType>(shape), allocator};
   }
+
   //
   // Run a compute shader program.
   //
   inline Status RunProgram(const ProgramBase& program) {
     return webgpu_context_.Run(*this, program);
   }
 
-  //
-  // Get the buffer manager from the GPU allocator.
-  //
-  const webgpu::BufferManager& BufferManager() const;
-
-  //
-  // Push error scope.
-  //
-  // This is useful only when "skip_validation" is not set.
-  //
-  void PushErrorScope();
-
-  //
-  // Pop error scope.
-  //
-  // This is useful only when "skip_validation" is not set.
-  //
-  Status PopErrorScope();
-
- protected:
+ private:
   WebGpuContext& webgpu_context_;
   OpKernelContext& kernel_context_;
   const WebGpuExecutionProvider& ep_;

onnxruntime/core/providers/webgpu/webgpu_context.cc

@@ -442,7 +442,7 @@ Status WebGpuContext::Run(ComputeContext& context, const ProgramBase& program) {
   const size_t uniform_buffer_total_size = (current_offset + max_alignment_of_field - 1) / max_alignment_of_field * max_alignment_of_field;
 
   WGPUBuffer uniform_buffer = nullptr;
-  const webgpu::BufferManager& buffer_mgr = context.BufferManager();
+  const webgpu::BufferManager& buffer_mgr = ComputeContext::BufferManagerAccessor::Get(context);
   if (uniform_buffer_total_size > 0) {
     std::vector<uint8_t> uniform_data_buffer(uniform_buffer_total_size);
 

onnxruntime/core/providers/webgpu/webgpu_kernel.cc

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/webgpu_kernel.h"
+#include "core/providers/webgpu/webgpu_context.h"
+#include "core/providers/webgpu/webgpu_execution_provider.h"
+#include "core/providers/webgpu/program.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+WebGpuKernel::WebGpuKernel(const OpKernelInfo& info)
+    : OpKernel(info),
+      ep_(*static_cast<const WebGpuExecutionProvider*>(info.GetExecutionProvider())) {
+}
+
+Status WebGpuKernel::Compute(OpKernelContext* p_op_kernel_context) const {
+  WebGpuContext& webgpu_context = WebGpuContextFactory::GetContext(ep_.GetDeviceId());
+  ComputeContext context{*p_op_kernel_context, ep_, webgpu_context};
+
+  if (webgpu_context.ValidationMode() >= ValidationMode::Full) {
+    webgpu_context.PushErrorScope();
+  }
+
+  Status s = ComputeInternal(context);
+
+  if (webgpu_context.ValidationMode() >= ValidationMode::Full) {
+    ORT_RETURN_IF_ERROR(webgpu_context.PopErrorScope());
+  }
+
+  return s;
+}
+
+}  // namespace webgpu
+}  // namespace onnxruntime

onnxruntime/core/providers/webgpu/webgpu_kernel.h

@@ -3,7 +3,6 @@
 
 #pragma once
 
-#include "core/providers/webgpu/webgpu_execution_provider.h"
 #include "core/providers/webgpu/compute_context.h"
 
 #include "core/framework/op_kernel.h"
@@ -18,20 +17,9 @@ namespace webgpu {
 // -----------------------------------------------------------------------
 class WebGpuKernel : public OpKernel {
  public:
-  explicit WebGpuKernel(const OpKernelInfo& info)
-      : OpKernel(info),
-        ep_(*static_cast<const WebGpuExecutionProvider*>(info.GetExecutionProvider())) {
-  }
+  explicit WebGpuKernel(const OpKernelInfo& info);
 
-  Status Compute(OpKernelContext* p_op_kernel_context) const override {
-    ComputeContext context{*p_op_kernel_context, ep_};
-
-    context.PushErrorScope();
-    Status s = ComputeInternal(context);
-    ORT_RETURN_IF_ERROR(context.PopErrorScope());
-
-    return s;
-  }
+  Status Compute(OpKernelContext* p_op_kernel_context) const override;
 
   virtual Status ComputeInternal(ComputeContext& context) const = 0;