Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 1 addition & 7 deletions backends/webgpu/runtime/WebGPUUtils.h
Original file line number Diff line number Diff line change
Expand Up @@ -142,13 +142,7 @@ make_uniform(WGPUDevice device, const void* data, size_t size) {
// Clamp a 1D workgroup count to the device limit, for grid-stride kernels that
// loop over any excess work (vs compute_1d_workgroup_count, which throws).
inline uint32_t clamp_workgroup_count(WGPUDevice device, uint32_t desired) {
WGPULimits limits = {};
uint32_t max_count =
wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
limits.maxComputeWorkgroupsPerDimension > 0
? limits.maxComputeWorkgroupsPerDimension
: 65535u; // WebGPU spec-default floor
return std::min(desired, max_count);
return std::min(desired, queried_max_workgroups(device));
}

} // namespace executorch::backends::webgpu::utils
14 changes: 2 additions & 12 deletions backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -62,16 +62,6 @@ constexpr uint32_t kQ4gswSteelTile = 64u;
constexpr uint32_t kQ4gswSteelBK = 16u;
constexpr uint32_t kQ4gswSteelInvocations = 256u;

// Max workgroups per 1D dispatch dimension: the device limit, or 65535 when the
// query fails / reports 0.
uint32_t max_workgroups_per_dim(WGPUDevice device) {
WGPULimits limits = {};
return (wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
limits.maxComputeWorkgroupsPerDimension > 0)
? limits.maxComputeWorkgroupsPerDimension
: 65535u;
}

// One workgroup per (tile_m x tile_n) tile, no grid-stride: throw when the tile
// count would exceed the 1D dispatch limit. Shared by the steel + shmem GEMM
// routes; `kind` names the route in the error message.
Expand All @@ -85,7 +75,7 @@ uint32_t tiled_wg_count(
const char* kind) {
const int64_t total_wgs =
utils::div_up<int64_t>(m, tile_m) * utils::div_up<int64_t>(n, tile_n);
if (total_wgs > static_cast<int64_t>(max_workgroups_per_dim(device))) {
if (total_wgs > static_cast<int64_t>(utils::queried_max_workgroups(device))) {
throw std::runtime_error(
std::string("WebGPU ") + op_name + ": " + kind +
" tile count exceeds the 1D dispatch limit");
Expand All @@ -109,7 +99,7 @@ steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) {
const uint64_t total =
static_cast<uint64_t>((m + kQ4gswSteelTile - 1u) / kQ4gswSteelTile) *
static_cast<uint64_t>((n + kQ4gswSteelTile - 1u) / kQ4gswSteelTile);
const uint32_t max_count = max_workgroups_per_dim(device);
const uint32_t max_count = utils::queried_max_workgroups(device);
return (total == 0u || total > max_count) ? 0u : static_cast<uint32_t>(total);
}

Expand Down
8 changes: 4 additions & 4 deletions backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -52,9 +52,9 @@ void resize_rms_norm(
if (rows == 0) {
throw std::runtime_error("WebGPU rms_norm: zero rows");
}
if (rows > 65535u) {
if (rows > utils::queried_max_workgroups(g.device())) {
throw std::runtime_error(
"WebGPU rms_norm: num_rows exceeds the 1D dispatch limit (65535)");
"WebGPU rms_norm: num_rows exceeds the 1D dispatch limit");
}
RmsNormParams p = {};
p.num_rows = rows;
Expand Down Expand Up @@ -104,9 +104,9 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
throw std::runtime_error("WebGPU rms_norm: zero rows");
}
// Validate the 1D dispatch limit before allocating any GPU objects.
if (num_rows > 65535u) {
if (num_rows > utils::queried_max_workgroups(device)) {
throw std::runtime_error(
"WebGPU rms_norm: num_rows exceeds the 1D dispatch limit (65535)");
"WebGPU rms_norm: num_rows exceeds the 1D dispatch limit");
}

// Create uniform buffer for params
Expand Down
Loading