diff --git a/onnxruntime/core/providers/webgpu/compute_context.cc b/onnxruntime/core/providers/webgpu/compute_context.cc index 25caa9b954fc0..d9badde2e5d35 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.cc +++ b/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 diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index 6bf7df74ea861..bf32f60bf5667 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/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,6 +131,7 @@ class ComputeContext { ORT_THROW_IF_ERROR(kernel_context_.GetTempSpaceAllocator(&allocator)); return {data_type, std::forward(shape), allocator}; } + // // Run a compute shader program. // @@ -127,26 +139,7 @@ class ComputeContext { 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_; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index a91e34c334687..9af9cd455b5a4 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/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 uniform_data_buffer(uniform_buffer_total_size); diff --git a/onnxruntime/core/providers/webgpu/webgpu_kernel.cc b/onnxruntime/core/providers/webgpu/webgpu_kernel.cc new file mode 100644 index 0000000000000..5e46779f35833 --- /dev/null +++ b/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(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 diff --git a/onnxruntime/core/providers/webgpu/webgpu_kernel.h b/onnxruntime/core/providers/webgpu/webgpu_kernel.h index e37be2944a22b..3c750e305421c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_kernel.h +++ b/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(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;