[ExecuTorch][WebGPU] Fold remaining hardcoded 65535 dispatch-limit checks into queried_max_workgroups (#20950)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #20916 by
@JCNTH
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/JCNTH/55/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/JCNTH/55/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/JCNTH/55/orig

@diff-train-skip-merge

Co-authored-by: Julian Ng-Thow-Hing <juliannth@meta.com>
Co-authored-by: Julian Ng-Thow-Hing <107437036+JCNTH@users.noreply.github.com>

Author: pytorchbot

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

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<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");
@@ -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);
 }
 

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<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