From ac22780299165228ef0030524ffac3a0ec2ee88b Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:34:12 -0700 Subject: [PATCH] [ExecuTorch][WebGPU] Fold remaining hardcoded 65535 dispatch-limit checks into queried_max_workgroups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20916 The backend queries the device's `maxComputeWorkgroupsPerDimension` through `utils::queried_max_workgroups` (keeping `65535` only as the documented spec-default fallback when the query fails), but a few landed dispatch-limit guards still hardcode the literal `65535u` or re-inline the same query. This folds them onto the shared helper, matching the convention already used by `Linear.cpp`/`Bmm.cpp` and the backend's "query the device limit, don't hardcode" rule. Key changes: - `RmsNorm.cpp` — replace both `> 65535u` guards (build path and resize hook) with `utils::queried_max_workgroups(...)`; drop the now-stale `(65535)` from the two error strings. - `QuantizedLinear.cpp` — delete the file-local `max_workgroups_per_dim` (a near-duplicate of the util) and call `utils::queried_max_workgroups` at both sites; `steel_supported()` (a `maxComputeInvocationsPerWorkgroup` gate) is unchanged. - `WebGPUUtils.h` — `clamp_workgroup_count` now delegates to `queried_max_workgroups` instead of re-inlining the same query + fallback. No behavior change on compliant devices: `queried_max_workgroups` returns the reported `maxComputeWorkgroupsPerDimension` (the WebGPU spec guarantees >= 65535) or falls back to 65535 if the query fails, so each swapped `> ...` guard can only relax versus the old `> 65535u` literal — no currently-passing shape newly throws. Co-authored-with: Claude Code. ghstack-source-id: 402967330 @exported-using-ghexport Differential Revision: [D111999110](https://our.internmc.facebook.com/intern/diff/D111999110/) --- backends/webgpu/runtime/WebGPUUtils.h | 8 +------- .../ops/quantized_linear/QuantizedLinear.cpp | 14 ++------------ backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp | 8 ++++---- 3 files changed, 7 insertions(+), 23 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 4326fde8a19..90ead722ace 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -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 diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index b0728764310..3dd83b7ab3e 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -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. @@ -85,7 +75,7 @@ uint32_t tiled_wg_count( const char* kind) { const int64_t total_wgs = utils::div_up(m, tile_m) * utils::div_up(n, tile_n); - if (total_wgs > static_cast(max_workgroups_per_dim(device))) { + if (total_wgs > static_cast(utils::queried_max_workgroups(device))) { throw std::runtime_error( std::string("WebGPU ") + op_name + ": " + kind + " tile count exceeds the 1D dispatch limit"); @@ -109,7 +99,7 @@ steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { const uint64_t total = static_cast((m + kQ4gswSteelTile - 1u) / kQ4gswSteelTile) * static_cast((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(total); } diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp index c07e58e8ec8..39aaae3a6bf 100644 --- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp +++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp @@ -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; @@ -104,9 +104,9 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& 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