diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index 64c267f970c24..98d718db1472c 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -214,6 +214,13 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; SubgroupMatrixMatMulNBitsProgram mul_program{nbits, config_index, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect}; mul_program.SetWorkgroupSize(work_group_size); + + // On Intel, use a fixed subgroup size of 32 for better performance. + if (context.AdapterInfo().vendor == std::string_view{"intel"} && + context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + mul_program.SetSubgroupSize(32); + } + uint32_t dispatch_x = (N + tile_size_b - 1) / tile_size_b; uint32_t num_m_tiles = (M + tile_size_a - 1) / tile_size_a; uint32_t dispatch_y = num_m_tiles; diff --git a/onnxruntime/core/providers/webgpu/program.cc b/onnxruntime/core/providers/webgpu/program.cc index df23f5d1dc73c..7197e3b08ff8b 100644 --- a/onnxruntime/core/providers/webgpu/program.cc +++ b/onnxruntime/core/providers/webgpu/program.cc @@ -330,7 +330,8 @@ ProgramBase::ProgramBase(std::string_view name, ProgramMetadata&& metadata) indirect_dispatch_tensor_{nullptr}, workgroup_size_x_{0}, workgroup_size_y_{0}, - workgroup_size_z_{0} { + workgroup_size_z_{0}, + subgroup_size_{0} { } ProgramBase& ProgramBase::AddInput(ProgramInput&& input) { @@ -389,6 +390,11 @@ ProgramBase& ProgramBase::SetWorkgroupSize(uint32_t x, uint32_t y, uint32_t z) { return *this; } +ProgramBase& ProgramBase::SetSubgroupSize(uint32_t size) { + subgroup_size_ = size; + return *this; +} + ProgramBase& ProgramBase::AddUniformVariable(ProgramUniformVariableValue&& variable) { variables_.emplace_back(variable); return *this; diff --git a/onnxruntime/core/providers/webgpu/program.h b/onnxruntime/core/providers/webgpu/program.h index 736111d2a18bc..e1b3d04fd35fd 100644 --- a/onnxruntime/core/providers/webgpu/program.h +++ b/onnxruntime/core/providers/webgpu/program.h @@ -324,6 +324,10 @@ class ProgramBase { // set the size of a workgroup grid. ProgramBase& SetWorkgroupSize(uint32_t x, uint32_t y, uint32_t z); + // request a fixed subgroup size via the `subgroup-size-control` feature (0 = driver-chosen, the default). + // requires the device to support wgpu::FeatureName::SubgroupSizeControl; see ShaderHelper::GenerateSourceCode. + ProgramBase& SetSubgroupSize(uint32_t size); + // add a uniform variable. // // the specified uniform variable should match the uniform definition in the class, @@ -364,6 +368,7 @@ class ProgramBase { inline uint32_t WorkgroupSizeX() const { return workgroup_size_x_; } inline uint32_t WorkgroupSizeY() const { return workgroup_size_y_; } inline uint32_t WorkgroupSizeZ() const { return workgroup_size_z_; } + inline uint32_t SubgroupSize() const { return subgroup_size_; } inline const std::vector& UniformVariables() const { return variables_; } inline const std::vector& OverridableConstants() const { return overridable_constants_; } @@ -392,6 +397,7 @@ class ProgramBase { uint32_t workgroup_size_x_; uint32_t workgroup_size_y_; uint32_t workgroup_size_z_; + uint32_t subgroup_size_; std::vector variables_; std::vector overridable_constants_; diff --git a/onnxruntime/core/providers/webgpu/program_cache_key.cc b/onnxruntime/core/providers/webgpu/program_cache_key.cc index d4a6b2bf1d812..ac40e8f4de07c 100644 --- a/onnxruntime/core/providers/webgpu/program_cache_key.cc +++ b/onnxruntime/core/providers/webgpu/program_cache_key.cc @@ -57,10 +57,11 @@ std::string CalculateProgramCacheKey(const ProgramBase& program, SS(ss, kStringInitialSizeCacheKey); // final key format: - // =[]:::: + // =[]::::: // // = ||... - // = ,, + // = ,, + // = // = // = ||... // = @@ -91,6 +92,11 @@ std::string CalculateProgramCacheKey(const ProgramBase& program, } } + // append the requested subgroup size (subgroup-size-control) if any + if (auto subgroup_size = program.SubgroupSize(); subgroup_size != 0) { + ss << ":" D("SubgroupSize=") << subgroup_size; + } + ss << ":" D("DispatchDim=") << (is_1d_dispatch ? "1" : "3"); ss << ":" D("UniformSizes="); bool first = true; diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index 657fc3a6f45ad..74e4d7cc9cd80 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -41,9 +41,11 @@ Status ShaderHelper::Init() { // dispatch group size is normalized so no need to validate it here // validate workgroup size - auto workgroup_size_x = program_.WorkgroupSizeX(); - auto workgroup_size_y = program_.WorkgroupSizeY(); - auto workgroup_size_z = program_.WorkgroupSizeZ(); + // normalize 0 (meaning "use default") the same way GenerateSourceCode() does, so validation + // below reflects the effective workgroup size that will actually be used in the shader. + auto workgroup_size_x = program_.WorkgroupSizeX() == 0 ? uint32_t(WORKGROUP_SIZE) : program_.WorkgroupSizeX(); + auto workgroup_size_y = program_.WorkgroupSizeY() == 0 ? uint32_t(1) : program_.WorkgroupSizeY(); + auto workgroup_size_z = program_.WorkgroupSizeZ() == 0 ? uint32_t(1) : program_.WorkgroupSizeZ(); ORT_RETURN_IF_NOT(workgroup_size_x <= limits_.maxComputeWorkgroupSizeX && workgroup_size_y <= limits_.maxComputeWorkgroupSizeY && @@ -56,12 +58,30 @@ Status ShaderHelper::Init() { ORT_RETURN_IF_NOT(workgroup_size_x * workgroup_size_y * workgroup_size_z <= limits_.maxComputeInvocationsPerWorkgroup, "Workgroup size exceeds the maximum allowed invocations ", limits_.maxComputeInvocationsPerWorkgroup); + // validate the requested subgroup size (subgroup-size-control), if any + if (auto subgroup_size = program_.SubgroupSize(); subgroup_size != 0) { + const auto& adapter_info = webgpu_context_.AdapterInfo(); + ORT_RETURN_IF_NOT(webgpu_context_.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + "Program ", program_.Name(), " requires subgroup-size-control but the device does not support it."); + ORT_RETURN_IF_NOT((subgroup_size & (subgroup_size - 1)) == 0, + "Requested subgroup size ", subgroup_size, " is not a power of two."); + ORT_RETURN_IF_NOT(subgroup_size >= adapter_info.subgroupMinSize && subgroup_size <= adapter_info.subgroupMaxSize, + "Requested subgroup size ", subgroup_size, " is outside the adapter's supported range [", + adapter_info.subgroupMinSize, ", ", adapter_info.subgroupMaxSize, "]"); + ORT_RETURN_IF_NOT(workgroup_size_x % subgroup_size == 0, + "Workgroup size x (", workgroup_size_x, ") must be a multiple of the requested subgroup size (", subgroup_size, ")"); + } + // init body string stream bool is_1d_dispatch = dispatch_group_size_y_ == 1 && dispatch_group_size_z_ == 1; bool use_indirect_dispatch = program_.IndirectDispatchTensor() != nullptr; // append header for main function so it is ready for user to append main function body - body_ss_ << "@compute @workgroup_size(workgroup_size_x, workgroup_size_y, workgroup_size_z)\n" + body_ss_ << "@compute @workgroup_size(workgroup_size_x, workgroup_size_y, workgroup_size_z)"; + if (program_.SubgroupSize() != 0) { + body_ss_ << " @subgroup_size(subgroup_size)"; + } + body_ss_ << "\n" "fn main(@builtin(global_invocation_id) global_id : vec3,\n" " @builtin(workgroup_id) workgroup_id : vec3,\n" " @builtin(local_invocation_index) local_idx : u32,\n" @@ -404,6 +424,9 @@ Status ShaderHelper::GenerateSourceCode(std::string& code, std::vector& sha if (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::Subgroups)) { ss << "enable subgroups;\n"; } + if (program_.SubgroupSize() != 0) { + ss << "enable subgroup_size_control;\n"; + } #if !defined(__wasm__) if (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { ss << "enable chromium_experimental_subgroup_matrix;\n"; @@ -422,6 +445,10 @@ Status ShaderHelper::GenerateSourceCode(std::string& code, std::vector& sha << ";\nconst workgroup_size_z: u32 = " << (program_.WorkgroupSizeZ() == 0 ? uint32_t(1) : program_.WorkgroupSizeZ()) << ";\n"; + if (auto subgroup_size = program_.SubgroupSize(); subgroup_size != 0) { + ss << "const subgroup_size: u32 = " << subgroup_size << ";\n"; + } + for (const auto& constant : program_metadata_.constants) { ss << "const " << constant.name << ": " << constant.type << " = "; WriteConstantValue(ss, constant); diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index a76d9b8d4001d..7395ef3205585 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -741,6 +741,7 @@ std::vector WebGpuContext::GetAvailableRequiredFeatures(const wgpu::FeatureName::TimestampQuery, wgpu::FeatureName::ShaderF16, wgpu::FeatureName::Subgroups, + wgpu::FeatureName::SubgroupSizeControl, #if !defined(__wasm__) wgpu::FeatureName::BufferMapExtendedUsages, #endif