[ExecuTorch][WebGPU] Make shader workgroup sizes runtime-configurable (pytorch#21006)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#20954 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/86/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/JCNTH/86/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/JCNTH/86/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

@@ -48,6 +48,16 @@ inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
   return desired;
 }
 
+// Clamp to device limit, then floor to pow2 (reduction kernels halve stride).
+inline uint32_t clamp_workgroup_size_pow2(WGPUDevice device, uint32_t desired) {
+  uint32_t v = clamp_workgroup_size(device, desired);
+  uint32_t p = 1u;
+  while (p <= (v >> 1u)) {
+    p <<= 1u;
+  }
+  return p;
+}
+
 struct WgCount {
   uint32_t x;
   uint32_t y;

backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp

@@ -244,6 +244,11 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const uint32_t wg_size =
       utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX);
   const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u);
+  // GEMV (bicol) is a pow2 tree reduction; compute its size only when used.
+  const uint32_t gemv_wg_size = use_gemv
+      ? utils::clamp_workgroup_size_pow2(
+            device, kQ4gswLinearCoop4BicolWorkgroupSizeX)
+      : 0u;
   // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible.
   const bool use_steel = !use_gemv && steel_supported(device) &&
       steel_workgroup_count(device, M, N, K) > 0u;
@@ -362,16 +367,17 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   WGPUPipelineLayout pipeline_layout =
       wgpuDeviceCreatePipelineLayout(device, &pl_desc);
 
+  // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed.
+  const bool fixed_wg = use_steel || use_shmem_gemm;
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
-  wg_size_constant.value = static_cast<double>(wg_size);
+  wg_size_constant.value =
+      static_cast<double>(use_gemv ? gemv_wg_size : wg_size);
 
   WGPUComputePipelineDescriptor pipeline_desc = {};
   pipeline_desc.layout = pipeline_layout;
   pipeline_desc.compute.module = shader;
   pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
-  // Only tiled GEMM overrides wg_size; GEMV/shmem (64) + steel (256) are fixed.
-  const bool fixed_wg = use_gemv || use_steel || use_shmem_gemm;
   pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u;
   pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant;
   WGPUComputePipeline pipeline =

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl

@@ -17,10 +17,11 @@ struct Params {
 @group(0) @binding(5) var<uniform> params: Params;
 
 // Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
-const WG: u32 = 64u;
-var<workgroup> partial: array<vec2<f32>, WG>;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
+var<workgroup> partial: array<vec2<f32>, wg_size>;
 
-@compute @workgroup_size(WG, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(num_workgroups) ngrp: vec3<u32>,
@@ -69,12 +70,12 @@ fn main(
         acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
         acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
       }
-      w = w + WG;
+      w = w + wg_size;
     }
 
     partial[lid.x] = vec2<f32>(acc0, acc1);
     workgroupBarrier();
-    var s: u32 = WG >> 1u;
+    var s: u32 = wg_size >> 1u;
     loop {
       if (s == 0u) {
         break;

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_coop4_bicol.wgsl - DO NOT EDIT.
-// wgsl-sha256: 44c53ca7879f7a3382a45703a67cbc6fc221ecbcada58f8c413df50abbc6e544
+// wgsl-sha256: ef677c2771c3d4a73704f4725bb490174149b971bc62e696595b8d13941c979e
 inline constexpr const char* kQ4gswLinearCoop4BicolWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
@@ -34,10 +34,11 @@ struct Params {
 @group(0) @binding(5) var<uniform> params: Params;
 
 // Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
-const WG: u32 = 64u;
-var<workgroup> partial: array<vec2<f32>, WG>;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
+var<workgroup> partial: array<vec2<f32>, wg_size>;
 
-@compute @workgroup_size(WG, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(num_workgroups) ngrp: vec3<u32>,
@@ -86,12 +87,12 @@ fn main(
         acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
         acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
       }
-      w = w + WG;
+      w = w + wg_size;
     }
 
     partial[lid.x] = vec2<f32>(acc0, acc1);
     workgroupBarrier();
-    var s: u32 = WG >> 1u;
+    var s: u32 = wg_size >> 1u;
     loop {
       if (s == 0u) {
         break;

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h

@@ -16,7 +16,8 @@ struct Params {
 }
 @group(0) @binding(5) var<uniform> params: Params;
 
-override wg_size: u32 = 64u;
+// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob).
+const wg_size: u32 = 64u;
 
 // Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
 const WG_M: u32 = 32u;   // output rows per workgroup

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_shmem.wgsl - DO NOT EDIT.
-// wgsl-sha256: 0a1219d3b6781315a21066089fca3f92235587e8af8eb734185f35ea4bfc8a52
+// wgsl-sha256: f0180cad41701807b34912b57d9715c02588c9447783bbf691f93836168fe981
 inline constexpr const char* kQ4gswLinearGemmShmemWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
@@ -33,7 +33,8 @@ struct Params {
 }
 @group(0) @binding(5) var<uniform> params: Params;
 
-override wg_size: u32 = 64u;
+// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob).
+const wg_size: u32 = 64u;
 
 // Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
 const WG_M: u32 = 32u;   // output rows per workgroup

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl

@@ -34,6 +34,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<${buffer_scalar_type(DTYPE)}, 1024>;   // BM*BK
 var<workgroup> Bs: array<${buffer_scalar_type(DTYPE)}, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce
+// wgsl-sha256: fe8b71afc5634e7db5607deb58f30c6517f93e1d66370ea017bcd9071201c6ca
 inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"(
 enable f16;
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@@ -50,6 +50,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<f16, 1024>;   // BM*BK
 var<workgroup> Bs: array<f16, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {