[webgpu] revise implementation of buffer split support (#26429)

### Description

This PR addresses a few concerns:
- revert `const ProgramBase&` -> `ProgramBase&`: this itself is not
doing something wrong but gives much more pressure for who reads the
code to understand whether/where the program object is modified. It also
can introduce further unexpected modifications to the program object
(for example the indirect dispatch code)
- change bool option
`"ep.webgpuexecutionprovider.smallStorageBufferBindingSizeForTesting"`
to `"ep.webgpuexecutionprovider.maxStorageBufferBindingSize"` so now
it's possible to set any value in option. (setting to <128MB will cause
an assert failure)
- segments are optional in cache key (only present when not equals to 1,
which is the common case)
- avoid some unnecessary API calls, which is OK for native but may
affect web perf.
- clean up the code a little bit and add a few comments

---------

Co-authored-by: Copilot <[email protected]>

Author: fs-eire

onnxruntime/core/providers/webgpu/compute_context.h

@@ -123,7 +123,7 @@ class ComputeContext {
   //
   // Run a compute shader program.
   //
-  inline Status RunProgram(ProgramBase& program) {
+  inline Status RunProgram(const ProgramBase& program) {
     return webgpu_context_.Run(*this, program);
   }
 

onnxruntime/core/providers/webgpu/program.cc

@@ -362,6 +362,7 @@ ProgramBase& ProgramBase::SetDispatchGroupSize(uint32_t x, uint32_t y, uint32_t
 
 ProgramBase& ProgramBase::SetIndirectDispatchTensor(const Tensor* indirect_dispatch_tensor) {
   indirect_dispatch_tensor_ = indirect_dispatch_tensor;
+  AddInput({indirect_dispatch_tensor, ProgramTensorMetadataDependency::None});
   return *this;
 }
 

onnxruntime/core/providers/webgpu/program.h

@@ -226,7 +226,6 @@ struct ProgramInput {
   ProgramInput(const Tensor* tensor, ProgramTensorMetadataDependency dependency, const TensorShape& override_shape, int component);
 
   const Tensor* tensor;
-  uint32_t segments = 1;
   ProgramTensorMetadataDependency dependency;
   ProgramVariableDataType var_type;
   bool use_override_shape;
@@ -246,7 +245,6 @@ struct ProgramOutput {
   ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, const TensorShape& override_shape, int component);
 
   Tensor* tensor;
-  uint32_t segments = 1;
   ProgramTensorMetadataDependency dependency;
   ProgramVariableDataType var_type;
   bool is_atomic;
@@ -348,18 +346,6 @@ class ProgramBase {
   inline const ProgramMetadata& Metadata() const { return metadata_; }
   inline const std::string& CacheHint() const { return cache_hint_; }
   inline const std::vector<ProgramInput>& Inputs() const { return inputs_; }
-  inline void setSegmentsForInput(size_t index, uint32_t segments) {
-    if (index >= inputs_.size()) {
-      throw std::out_of_range("input index out of range");
-    }
-    inputs_[index].segments = segments;
-  }
-  inline void setSegmentsForOutput(size_t index, uint32_t segments) {
-    if (index >= outputs_.size()) {
-      throw std::out_of_range("output index out of range");
-    }
-    outputs_[index].segments = segments;
-  }
   inline const std::vector<ProgramOutput>& Outputs() const { return outputs_; }
   inline const std::vector<TensorShape>& Indices() const { return indices_; }
   inline uint32_t DispatchGroupSizeX() const { return dispatch_group_size_x_; }

onnxruntime/core/providers/webgpu/program_cache_key.cc

@@ -17,8 +17,12 @@ namespace webgpu {
 
 namespace {
 // append the info of an input or output to the cachekey
-void AppendTensorInfo(std::ostream& ss, const TensorShape& tensor_shape, ProgramVariableDataType var_type, ProgramTensorMetadataDependency dependency,
-                      bool& first, uint32_t segments = 1) {
+void AppendTensorInfo(std::ostream& ss,
+                      const TensorShape& tensor_shape,
+                      ProgramVariableDataType var_type,
+                      ProgramTensorMetadataDependency dependency,
+                      bool& first,
+                      uint32_t segments) {
   if (first) {
     first = false;
   } else {
@@ -34,7 +38,9 @@ void AppendTensorInfo(std::ostream& ss, const TensorShape& tensor_shape, Program
     ss << ';';
   }
 
-  ss D("Segs=") << segments << ';';
+  if (segments != 1) {
+    ss D("Segs=") << segments << ';';
+  }
 
   if ((dependency & ProgramTensorMetadataDependency::Shape) == ProgramTensorMetadataDependency::Shape) {
     ss D("Dims=") << tensor_shape.ToString();
@@ -44,7 +50,10 @@ void AppendTensorInfo(std::ostream& ss, const TensorShape& tensor_shape, Program
 }
 }  // namespace
 
-std::string CalculateProgramCacheKey(const ProgramBase& program, bool is_1d_dispatch) {
+std::string CalculateProgramCacheKey(const ProgramBase& program,
+                                     std::span<uint32_t> inputs_segments,
+                                     std::span<uint32_t> outputs_segments,
+                                     bool is_1d_dispatch) {
   SS(ss, kStringInitialSizeCacheKey);
 
   // final key format:
@@ -56,7 +65,7 @@ std::string CalculateProgramCacheKey(const ProgramBase& program, bool is_1d_disp
   // <UNIFORMS>          = <UNIFORMS_INFO_0>|<UNIFORMS_INFO_1>|...
   // <UNIFORMS_INFO_i>   = <UNIFORM_LENGTH>
   // <INPUTS_INFO>       = <INPUTS_INFO_0>|<INPUTS_INFO_1>|...
-  // <INPUTS_INFO_i>     = <TENSOR_ELEMENT_TYPE_OR_EMPTY>;<TENSOR_SHAPE_OR_RANK_OR_EMPTY>
+  // <INPUTS_INFO_i>     = <TENSOR_ELEMENT_TYPE_OR_EMPTY>;<TENSOR_SEGMENTS_OR_EMPTY>;<TENSOR_SHAPE_OR_RANK_OR_EMPTY>
   ss << program.Name();
 
   // append custom cache hint if any
@@ -98,19 +107,26 @@ std::string CalculateProgramCacheKey(const ProgramBase& program, bool is_1d_disp
 
   ss << ":" D("Inputs=");
   first = true;
-  for (const auto& input : program.Inputs()) {
-    AppendTensorInfo(ss, input.use_override_shape ? input.override_shape : input.tensor->Shape(), input.var_type, input.dependency, first, input.segments);
+  for (size_t i = 0; i < program.Inputs().size(); i++) {
+    const auto& input = program.Inputs()[i];
+    AppendTensorInfo(ss,
+                     input.use_override_shape ? input.override_shape : input.tensor->Shape(),
+                     input.var_type,
+                     input.dependency,
+                     first,
+                     inputs_segments[i]);
   }
 
   ss << ":" D("Outputs=");
   first = true;
-  for (const auto& output : program.Outputs()) {
+  for (size_t i = 0; i < program.Outputs().size(); i++) {
+    const auto& output = program.Outputs()[i];
     AppendTensorInfo(ss,
                      output.use_override_shape ? output.override_shape : output.tensor->Shape(),
                      output.var_type,
                      output.dependency,
                      first,
-                     output.segments);
+                     outputs_segments[i]);
   }
 
   if (!program.Indices().empty()) {

onnxruntime/core/providers/webgpu/program_cache_key.h

@@ -3,14 +3,18 @@
 
 #pragma once
 
+#include <span>
 #include <string>
 
 #include "core/providers/webgpu/program.h"
 
 namespace onnxruntime {
 namespace webgpu {
 
-std::string CalculateProgramCacheKey(const ProgramBase& program, bool is_1d_dispatch);
+std::string CalculateProgramCacheKey(const ProgramBase& program,
+                                     std::span<uint32_t> inputs_segments,
+                                     std::span<uint32_t> outputs_segments,
+                                     bool is_1d_dispatch);
 
 }  // namespace webgpu
 }  // namespace onnxruntime

onnxruntime/core/providers/webgpu/program_manager.cc

@@ -39,30 +39,35 @@ Status ProgramManager::NormalizeDispatchGroupSize(uint32_t& x, uint32_t& y, uint
   return Status::OK();
 }
 
-Status ProgramManager::CalculateSegmentsForInputsAndOutputs(ProgramBase& program) {
+Status ProgramManager::CalculateSegmentsForInputsAndOutputs(const ProgramBase& program, std::vector<uint32_t>& inputs_segments, std::vector<uint32_t>& outputs_segments) const {
+  inputs_segments.resize(program.Inputs().size(), 1);
+  outputs_segments.resize(program.Outputs().size(), 1);
+
   const uint64_t maxStorageBufferBindingSize = webgpu_context_.DeviceLimits().maxStorageBufferBindingSize;
 
   // Inputs
   for (size_t i = 0; i < program.Inputs().size(); ++i) {
     const auto& input = program.Inputs()[i];
     if (input.tensor && input.tensor->SizeInBytes() > maxStorageBufferBindingSize) {
       uint32_t segments = static_cast<uint32_t>((input.tensor->SizeInBytes() + maxStorageBufferBindingSize - 1) / maxStorageBufferBindingSize);
-      program.setSegmentsForInput(i, segments);
+      inputs_segments[i] = segments;
     }
   }
   // Outputs
   for (size_t i = 0; i < program.Outputs().size(); ++i) {
     const auto& output = program.Outputs()[i];
     if (output.tensor && output.tensor->SizeInBytes() > maxStorageBufferBindingSize) {
       uint32_t segments = static_cast<uint32_t>((output.tensor->SizeInBytes() + maxStorageBufferBindingSize - 1) / maxStorageBufferBindingSize);
-      program.setSegmentsForOutput(i, segments);
+      outputs_segments[i] = segments;
     }
   }
   return Status::OK();
 }
 
 Status ProgramManager::Build(const ProgramBase& program,
                              const ProgramMetadata& program_metadata,
+                             const std::span<uint32_t> inputs_segments,
+                             const std::span<uint32_t> outputs_segments,
 #ifndef NDEBUG  // if debug build
                              const std::string& program_key,
 #endif
@@ -74,6 +79,8 @@ Status ProgramManager::Build(const ProgramBase& program,
   auto& device = webgpu_context_.Device();
   ShaderHelper shader_helper{program,
                              program_metadata,
+                             inputs_segments,
+                             outputs_segments,
                              device,
                              webgpu_context_.DeviceLimits(),
                              normalized_dispatch_x,
@@ -83,8 +90,10 @@ Status ProgramManager::Build(const ProgramBase& program,
 
   ORT_RETURN_IF_ERROR(program.GenerateShaderCode(shader_helper));
 
-  // Finalize inputs after GenerateShaderCode() to ensure indirect buffer is added as the last input
-  shader_helper.FinalizeInputs();
+  // Add indirect buffer as the last shader input when using indirect dispatch.
+  if (program.IndirectDispatchTensor() != nullptr) {
+    shader_helper.AddInput("indirect_buffer", ShaderUsage::None);
+  }
 
   ORT_RETURN_IF_ERROR(shader_helper.ValidateShapeForInputs());
   ORT_RETURN_IF_ERROR(shader_helper.ValidateShapeForOutputs());

onnxruntime/core/providers/webgpu/program_manager.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <span>
 #include <string>
 #include <unordered_map>
 
@@ -38,10 +39,12 @@ class ProgramManager {
   ProgramManager(WebGpuContext& webgpu_context) : webgpu_context_(webgpu_context) {}
 
   Status NormalizeDispatchGroupSize(uint32_t& x, uint32_t& y, uint32_t& z) const;
-  Status CalculateSegmentsForInputsAndOutputs(ProgramBase& program);
+  Status CalculateSegmentsForInputsAndOutputs(const ProgramBase& program, std::vector<uint32_t>& inputs_segments, std::vector<uint32_t>& outputs_segments) const;
 
   Status Build(const ProgramBase& program,
                const ProgramMetadata& metadata,
+               const std::span<uint32_t> inputs_segments,
+               const std::span<uint32_t> outputs_segments,
 #ifndef NDEBUG  // if debug build
                const std::string& program_key,
 #endif

onnxruntime/core/providers/webgpu/shader_helper.cc

@@ -18,13 +18,17 @@ namespace webgpu {
 
 ShaderHelper::ShaderHelper(const ProgramBase& program,
                            const ProgramMetadata& program_metadata,
+                           const std::span<uint32_t> inputs_segments,
+                           const std::span<uint32_t> outputs_segments,
                            const wgpu::Device& device,
                            const wgpu::Limits& limits,
                            uint32_t dispatch_group_size_x,
                            uint32_t dispatch_group_size_y,
                            uint32_t dispatch_group_size_z)
     : device_{device},
       limits_{limits},
+      inputs_segments_{inputs_segments},
+      outputs_segments_{outputs_segments},
       dispatch_group_size_x_{dispatch_group_size_x},
       dispatch_group_size_y_{dispatch_group_size_y},
       dispatch_group_size_z_{dispatch_group_size_z},
@@ -95,21 +99,14 @@ Status ShaderHelper::Init() {
   return Status::OK();
 }
 
-void ShaderHelper::FinalizeInputs() {
-  // Automatically add indirect buffer as the last shader input when using indirect dispatch.
-  if (program_.IndirectDispatchTensor() != nullptr) {
-    AddInput("indirect_buffer", ShaderUsage::None);
-  }
-}
-
 const ShaderVariableHelper& ShaderHelper::AddInput(const std::string& name, ShaderUsage usage) {
   const size_t input_index = input_vars_.size();
   ORT_ENFORCE(input_index < program_.Inputs().size(),
               "Too many inputs in the program (", program_.Inputs().size(), ")");
 
   const auto& dims = program_.Inputs()[input_index].use_override_shape ? program_.Inputs()[input_index].override_shape
                                                                        : program_.Inputs()[input_index].tensor->Shape();
-  return AddVariableImpl(true, name, usage, dims, program_.Inputs()[input_index].segments);
+  return AddVariableImpl(true, name, usage, dims, inputs_segments_[input_index]);
 }
 
 const ShaderVariableHelper& ShaderHelper::AddOutput(const std::string& name, ShaderUsage usage) {
@@ -119,7 +116,7 @@ const ShaderVariableHelper& ShaderHelper::AddOutput(const std::string& name, Sha
 
   const auto& dims = program_.Outputs()[output_index].use_override_shape ? program_.Outputs()[output_index].override_shape
                                                                          : program_.Outputs()[output_index].tensor->Shape();
-  return AddVariableImpl(false, name, usage, dims, program_.Outputs()[output_index].segments);
+  return AddVariableImpl(false, name, usage, dims, outputs_segments_[output_index]);
 }
 
 const ShaderIndicesHelper& ShaderHelper::AddIndices(const std::string& name, ShaderUsage usage) {

onnxruntime/core/providers/webgpu/shader_helper.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <span>
 #include <sstream>
 
 #include "core/providers/webgpu/webgpu_external_header.h"
@@ -67,6 +68,8 @@ class ShaderHelper final {
  public:
   ShaderHelper(const ProgramBase& program,
                const ProgramMetadata& program_metadata,
+               const std::span<uint32_t> inputs_segments,
+               const std::span<uint32_t> outputs_segments,
                const wgpu::Device& device,
                const wgpu::Limits& limits,
                uint32_t dispatch_group_size_x,
@@ -75,11 +78,6 @@ class ShaderHelper final {
 
   Status Init();
 
-  // Finalize inputs by automatically adding the indirect buffer if needed.
-  // This should be called after GenerateShaderCode() to ensure the indirect buffer
-  // is registered as the last input.
-  void FinalizeInputs();
-
   // Add an input variable to the shader.
   //
   // depending on the usage of the variable, additional code may be generated.
@@ -164,6 +162,8 @@ class ShaderHelper final {
 
   const wgpu::Device& device_;
   const wgpu::Limits& limits_;
+  const std::span<uint32_t> inputs_segments_;
+  const std::span<uint32_t> outputs_segments_;
   uint32_t dispatch_group_size_x_;
   uint32_t dispatch_group_size_y_;
   uint32_t dispatch_group_size_z_;