Skip to content

Implement WebGPU subgroup-size-control infrastructure #32056

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jiajia Qin (qjia7) merged 2 commits into from
Aug 17, 2026
Merged

Implement WebGPU subgroup-size-control infrastructure #32056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2858f03
Implement WebGPU subgroup-size-control infrastructure
jchen10 Aug 13, 2026
586f73c
Fix use default WorkgroupSize
jchen10 Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
8 changes: 7 additions & 1 deletion onnxruntime/core/providers/webgpu/program.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand Down
6 changes: 6 additions & 0 deletions onnxruntime/core/providers/webgpu/program.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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<ProgramUniformVariableValue>& UniformVariables() const { return variables_; }
inline const std::vector<ProgramOverridableConstantValue>& OverridableConstants() const { return overridable_constants_; }

Expand Down Expand Up @@ -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<ProgramUniformVariableValue> variables_;
std::vector<ProgramOverridableConstantValue> overridable_constants_;
Expand Down
10 changes: 8 additions & 2 deletions onnxruntime/core/providers/webgpu/program_cache_key.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,11 @@ std::string CalculateProgramCacheKey(const ProgramBase& program,
SS(ss, kStringInitialSizeCacheKey);

// final key format:
// <KEY>=<PROGRAM_NAME>[<CUSTOM_CACHE_HINT>]:<WORKGROUP_SIZE>:<DISPATCH_FLAG>:<UNIFORMS>:<INPUTS_INFO>
// <KEY>=<PROGRAM_NAME>[<CUSTOM_CACHE_HINT>]:<WORKGROUP_SIZE>:<SUBGROUP_SIZE>:<DISPATCH_FLAG>:<UNIFORMS>:<INPUTS_INFO>
//
// <CUSTOM_CACHE_HINT> = <HINT_0>|<HINT_1>|...
// <WORKGROUP_SIZE> = <X_IF_OVERRIDED>,<Y_IF_OVERRIDED>,<Z_IF_OVERRIDED>
// <WORKGROUP_SIZE> = <X_IF_OVERRIDDEN>,<Y_IF_OVERRIDDEN>,<Z_IF_OVERRIDDEN>
// <SUBGROUP_SIZE> = <SUBGROUP_SIZE_IF_OVERRIDDEN>
// <DISPATCH_FLAG> = <!IS_1D_DISPATCH>
// <UNIFORMS> = <UNIFORMS_INFO_0>|<UNIFORMS_INFO_1>|...
// <UNIFORMS_INFO_i> = <UNIFORM_LENGTH>
Expand Down Expand Up @@ -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;
Expand Down
35 changes: 31 additions & 4 deletions onnxruntime/core/providers/webgpu/shader_helper.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,9 +41,11 @@
// 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();

Check warning on line 46 in onnxruntime/core/providers/webgpu/shader_helper.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Using deprecated casting style.  Use static_cast<uint32_t>(...) instead  [readability/casting] [4]

Raw Output:
onnxruntime/core/providers/webgpu/shader_helper.cc:46:  Using deprecated casting style.  Use static_cast<uint32_t>(...) instead  [readability/casting] [4]
auto workgroup_size_y = program_.WorkgroupSizeY() == 0 ? uint32_t(1) : program_.WorkgroupSizeY();

Check warning on line 47 in onnxruntime/core/providers/webgpu/shader_helper.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Using deprecated casting style.  Use static_cast<uint32_t>(...) instead  [readability/casting] [4]

Raw Output:
onnxruntime/core/providers/webgpu/shader_helper.cc:47:  Using deprecated casting style.  Use static_cast<uint32_t>(...) instead  [readability/casting] [4]
auto workgroup_size_z = program_.WorkgroupSizeZ() == 0 ? uint32_t(1) : program_.WorkgroupSizeZ();

Check warning on line 48 in onnxruntime/core/providers/webgpu/shader_helper.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Using deprecated casting style.  Use static_cast<uint32_t>(...) instead  [readability/casting] [4]

Raw Output:
onnxruntime/core/providers/webgpu/shader_helper.cc:48:  Using deprecated casting style.  Use static_cast<uint32_t>(...) instead  [readability/casting] [4]

ORT_RETURN_IF_NOT(workgroup_size_x <= limits_.maxComputeWorkgroupSizeX &&
workgroup_size_y <= limits_.maxComputeWorkgroupSizeY &&
Expand All @@ -56,12 +58,30 @@
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.");
Comment thread
jchen10 marked this conversation as resolved.
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<u32>,\n"
" @builtin(workgroup_id) workgroup_id : vec3<u32>,\n"
" @builtin(local_invocation_index) local_idx : u32,\n"
Expand Down Expand Up @@ -404,6 +424,9 @@
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";
Expand All @@ -422,6 +445,10 @@
<< ";\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);
Expand Down
1 change: 1 addition & 0 deletions onnxruntime/core/providers/webgpu/webgpu_context.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -741,6 +741,7 @@ std::vector<wgpu::FeatureName> WebGpuContext::GetAvailableRequiredFeatures(const
wgpu::FeatureName::TimestampQuery,
wgpu::FeatureName::ShaderF16,
wgpu::FeatureName::Subgroups,
wgpu::FeatureName::SubgroupSizeControl,
#if !defined(__wasm__)
wgpu::FeatureName::BufferMapExtendedUsages,
#endif
Expand Down
Loading