[ExecuTorch][WebGPU] Use 2D-folded dispatch for cat and slice (pytorch#23900)

cat and slice dispatched a 1D grid with one invocation per element, so
any
input (cat) or output (slice) larger than 65535 workgroups failed to
load
even though the buffers fit comfortably within device limits.

Switch both to compute_2d_workgroup_count and the folded flat index the
other elementwise ops already use, including in the resize hooks.

This change was authored with Claude Code.

Test Plan:
- New heavy grid_limit op tests for cat and slice cross the 1D limit;
they
  fail to load before this change and pass after, alongside the existing
  cat/slice cases (native Dawn on Metal).
- gen_wgsl_headers.py --check, test_wgsl_codegen.py, lintrunner.

Author: digantdesai

backends/webgpu/runtime/ops/cat/Cat.cpp

@@ -71,7 +71,7 @@ void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
   // Validate + cache input meta/wgc BEFORE any GPU alloc (no leak on throw).
   std::vector<TensorMeta> in_metas(ids.size());
-  std::vector<uint32_t> wg_counts(ids.size());
+  std::vector<utils::WgCount> wg_counts(ids.size());
   int64_t concat_sum = 0;
   for (size_t k = 0; k < ids.size(); k++) {
     const int id = ids[k];
@@ -92,7 +92,7 @@ void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         static_cast<size_t>(in_metas[k].numel) * sizeof(float)) {
       throw std::runtime_error("cat: non-fp32 input (nbytes != numel * 4)");
     }
-    wg_counts[k] = utils::compute_1d_workgroup_count(
+    wg_counts[k] = utils::compute_2d_workgroup_count(
         device, in_metas[k].numel, wg_size, "cat");
     concat_sum += in_tensor.dims[dim];
   }
@@ -158,8 +158,12 @@ void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
     in_meta_bufs[k] = in_meta_buf;
     params_bufs[k] = params_buf;
-    dispatch_idxs[k] =
-        graph.add_dispatch({bundle.pipeline, bundle.bind_group, wg_counts[k]});
+    dispatch_idxs[k] = graph.add_dispatch(
+        {bundle.pipeline,
+         bundle.bind_group,
+         wg_counts[k].x,
+         "cat",
+         wg_counts[k].y});
     if (!shared_resources.has_value()) {
       shared_resources.emplace(std::move(bundle));
     }
@@ -211,9 +215,10 @@ void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       params.off_k = off;
       wgpuQueueWriteBuffer(
           g.queue(), params_bufs[k], 0, &params, sizeof(params));
-      g.dispatch_at(dispatch_idxs[k]).workgroup_count_x =
-          utils::compute_1d_workgroup_count(
-              g.device(), in_meta.numel, wg_size, "cat(resize)");
+      const utils::WgCount wg = utils::compute_2d_workgroup_count(
+          g.device(), in_meta.numel, wg_size, "cat(resize)");
+      g.dispatch_at(dispatch_idxs[k]).workgroup_count_x = wg.x;
+      g.dispatch_at(dispatch_idxs[k]).workgroup_count_y = wg.y;
       off += static_cast<uint32_t>(in_dims[cdim]);
     }
   };

backends/webgpu/runtime/ops/cat/cat.wgsl

@@ -19,8 +19,11 @@ struct Params {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let in_bufi = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large tensors).
+    let in_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (in_bufi >= in_meta.numel) {
         return;
     }

backends/webgpu/runtime/ops/cat/cat_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from cat.wgsl - DO NOT EDIT.
-// wgsl-sha256: 1a5f66607e2959c12d757989a3c1ae6fc6f2ede139bc9afb6e14dd3961a6fcb7
+// wgsl-sha256: 5bef4284287c6d3b7f7aefcbaf642f96347575f5271dd190975a6d59d093af36
 inline constexpr const char* kCatWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<f32>;
@@ -36,8 +36,11 @@ struct Params {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let in_bufi = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large tensors).
+    let in_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (in_bufi >= in_meta.numel) {
         return;
     }

backends/webgpu/runtime/ops/slice/Slice.cpp

@@ -152,7 +152,7 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   params.step = static_cast<uint32_t>(step);
 
   uint32_t wg_size = utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX);
-  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+  const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
       device, out_meta.numel, wg_size, "slice");
 
   WGPUConstantEntry wg_size_constant = {};
@@ -186,7 +186,12 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       &wg_size_constant,
       1);
 
-  graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
+  graph.add_dispatch(
+      {bundle.pipeline,
+       bundle.bind_group,
+       workgroup_count.x,
+       "slice",
+       workgroup_count.y});
   const size_t dispatch_idx = graph.num_dispatches() - 1;
 
   // Dynamic shapes: live start/end -> out[dim] len + meta/params/dispatch.
@@ -228,9 +233,10 @@ void slice_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     p.start = static_cast<uint32_t>(start);
     p.step = static_cast<uint32_t>(step);
     wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
-    g.dispatch_at(dispatch_idx).workgroup_count_x =
-        utils::compute_1d_workgroup_count(
-            g.device(), om.numel, wg_size, "slice(resize)");
+    const utils::WgCount wg = utils::compute_2d_workgroup_count(
+        g.device(), om.numel, wg_size, "slice(resize)");
+    g.dispatch_at(dispatch_idx).workgroup_count_x = wg.x;
+    g.dispatch_at(dispatch_idx).workgroup_count_y = wg.y;
   };
   if (is_symint(graph, start_id)) {
     graph.add_resize_hook(start_id, recompute);

backends/webgpu/runtime/ops/slice/slice.wgsl

@@ -20,8 +20,11 @@ struct Params {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_bufi = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large tensors).
+    let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (out_bufi >= out_meta.numel) {
         return;
     }

backends/webgpu/runtime/ops/slice/slice_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from slice.wgsl - DO NOT EDIT.
-// wgsl-sha256: 6a895a8c321cd3ddaffc468d7843c9dea3eaeb9d3de0088a1a09419b3ccfd10b
+// wgsl-sha256: eb05f6d85e338fa3bb3dcf485e75a7fde71171fdb98098b08b235ed2bf009113
 inline constexpr const char* kSliceWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<f32>;
@@ -37,8 +37,11 @@ struct Params {
 override wg_size: u32 = 64u;
 
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let out_bufi = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large tensors).
+    let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size);
     if (out_bufi >= out_meta.numel) {
         return;
     }

backends/webgpu/test/op_tests/cases.py

@@ -1052,7 +1052,17 @@ def _unsqueeze_suite() -> WebGPUTestSuite:
 
 @register_op_test("slice")
 def _slice_suite() -> WebGPUTestSuite:
-    return _fn_config_suite(SliceModule, _SLICE_CONFIGS)
+    suite = _fn_config_suite(SliceModule, _SLICE_CONFIGS)
+    # Output past the 65535 x 64 1D dispatch limit -> 2D-folded grid.
+    suite.cases.append(
+        Case(
+            name="grid_limit",
+            construct={"fn": lambda x: x[:, :1]},
+            inputs=((1, 2, 2100, 2100),),
+            heavy=True,
+        )
+    )
+    return suite
 
 
 @register_op_test("permute")
@@ -1076,6 +1086,15 @@ def _cat_suite() -> WebGPUTestSuite:
         cases=[
             Case(name=n, construct={"dim": dim}, inputs=tuple(shapes))
             for n, (shapes, dim) in _CAT_CONFIGS.items()
+        ]
+        + [
+            # Per-input dispatch past the 65535 x 64 1D limit -> 2D-folded grid.
+            Case(
+                name="grid_limit",
+                construct={"dim": 1},
+                inputs=((1, 1, 2100, 2100), (1, 1, 2100, 2100)),
+                heavy=True,
+            )
         ],
         golden_dtype="float32",  # concatenation copies values; fp64 bit-identical
     )

backends/webgpu/test/test_wgsl_codegen.py

@@ -220,14 +220,14 @@ def test_generated_output_manifest_digest(self) -> None:
             digest.update(b"\0")
             digest.update(output.read_bytes())
             digest.update(b"\0")
-        self.assertEqual(len(outputs), 136)
+        self.assertEqual(len(outputs), 138)
         self.assertEqual(
             digest.hexdigest(),
-            "0512f8d258952e446ffaedcb653b6a3a720eccf8a6b5327d95fd454a912214a3",
+            "fc01291b581fcbbb1112ed5b0c7ce5fcab2e1476d2ea92018ae300c9495cf132",
         )
         self.assertEqual(
             hashlib.sha256(g.registry_path().read_bytes()).hexdigest(),
-            "28aaa7a8d3e916df43e407120e91d487d0d51cbc5ca93c56bd822d25d109890e",
+            "3f5a3acd2645ea422a14091daa8fc2b7e5163d2c93aee06aed3bc0d2aa1d43d4",
         )
 
     def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: