Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/runtime/ops/reduce/Reduce.cpp

@@ -111,8 +111,9 @@ void reduce_impl(
 
   const uint32_t wg_size =
       utils::clamp_workgroup_size(device, kReduceWorkgroupSizeX);
-  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
-      device, static_cast<uint32_t>(outputs), wg_size, op_name);
+  // Cooperative reduction: one workgroup per output element (2D-folded grid).
+  const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
+      device, static_cast<uint32_t>(outputs), 1u, op_name);
 
   ReduceParams params = {};
   params.outer = outer;
@@ -190,8 +191,8 @@ void reduce_impl(
   bg_desc.entries = bg_entries;
   WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
 
-  const size_t dispatch_idx =
-      graph.add_dispatch({pipeline, bind_group, workgroup_count, op_name});
+  const size_t dispatch_idx = graph.add_dispatch(
+      {pipeline, bind_group, workgroup_count.x, op_name, workgroup_count.y});
 
   // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch.
   WGPUBuffer params_buf = uniform_buffer;
@@ -205,7 +206,6 @@ void reduce_impl(
        keepdim,
        is_mean_u,
        build_outputs,
-       wg_size,
        dispatch_idx,
        params_buf](WebGPUGraph& g) {
         const auto& d = g.cur_dims(in_id);
@@ -225,12 +225,13 @@ void reduce_impl(
         p.inner = n;
         p.is_mean = is_mean_u;
         wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
-        g.dispatch_at(dispatch_idx).workgroup_count_x =
-            utils::compute_1d_workgroup_count(
-                g.device(),
-                static_cast<uint32_t>(live_outputs),
-                wg_size,
-                "reduce(resize)");
+        const utils::WgCount wgc = utils::compute_2d_workgroup_count(
+            g.device(),
+            static_cast<uint32_t>(live_outputs),
+            1u,
+            "reduce(resize)");
+        g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
+        g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
         // Propagate reduced output dims for downstream resize hooks.
         int64_t nd = static_cast<int64_t>(d.size());
         int64_t rd = dim < 0 ? dim + nd : dim;

backends/webgpu/runtime/ops/reduce/reduce.wgsl

@@ -11,22 +11,46 @@ struct Params {
 
 override wg_size: u32 = 256;
 
+// Cooperative shared-memory reduction, one workgroup per output element: each
+// thread sums a strided slice of the reduced dim into a shared partial, then
+// thread 0 folds the partials. Same one-workgroup-per-row shared-memory shape as
+// Vulkan's reduce_per_row_buffer.glsl. Fixed 256 upper bound >= any clamped
+// wg_size; only [0, wg_size) is used.
+var<workgroup> partials: array<f32, 256>;
+
 @compute @workgroup_size(wg_size)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let t = gid.x;
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  // One workgroup per output; 2D-fold lifts the 65535 grid cap. `t` is uniform
+  // across the workgroup, so the early return keeps the barrier in uniform flow.
+  let t = wid.x + wid.y * num_workgroups.x;
   let outs = params.outer_ * params.inner_;
   if (t >= outs) {
     return;
   }
   let oo = t / params.inner_;
   let ii = t % params.inner_;
   let base = oo * params.r_ * params.inner_ + ii;
+
   var acc: f32 = 0.0;
-  for (var r: u32 = 0u; r < params.r_; r = r + 1u) {
-    acc = acc + inp[base + r * params.inner_];
+  var k: u32 = lid.x;
+  while (k < params.r_) {
+    acc = acc + inp[base + k * params.inner_];
+    k = k + wg_size;
   }
-  if (params.is_mean == 1u) {
-    acc = acc / f32(params.r_);
+  partials[lid.x] = acc;
+  workgroupBarrier();
+
+  if (lid.x == 0u) {
+    var s: f32 = partials[0];
+    for (var i: u32 = 1u; i < wg_size; i = i + 1u) {
+      s = s + partials[i];
+    }
+    if (params.is_mean == 1u) {
+      s = s / f32(params.r_);
+    }
+    out[t] = s;
   }
-  out[t] = acc;
 }

backends/webgpu/runtime/ops/reduce/reduce_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from reduce.wgsl - DO NOT EDIT.
-// wgsl-sha256: e657d5252093de9e416188ef96acaeea739d188673b4d1e3cce3d6b3d91afd68
+// wgsl-sha256: b4a40b67af55986ea5136f06c260cf3625e4db7809975fc616aeace16b479b2d
 inline constexpr const char* kReduceWGSL = R"(
 struct Params {
   outer_: u32,
@@ -28,24 +28,48 @@ struct Params {
 
 override wg_size: u32 = 256;
 
+// Cooperative shared-memory reduction, one workgroup per output element: each
+// thread sums a strided slice of the reduced dim into a shared partial, then
+// thread 0 folds the partials. Same one-workgroup-per-row shared-memory shape as
+// Vulkan's reduce_per_row_buffer.glsl. Fixed 256 upper bound >= any clamped
+// wg_size; only [0, wg_size) is used.
+var<workgroup> partials: array<f32, 256>;
+
 @compute @workgroup_size(wg_size)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let t = gid.x;
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  // One workgroup per output; 2D-fold lifts the 65535 grid cap. `t` is uniform
+  // across the workgroup, so the early return keeps the barrier in uniform flow.
+  let t = wid.x + wid.y * num_workgroups.x;
   let outs = params.outer_ * params.inner_;
   if (t >= outs) {
     return;
   }
   let oo = t / params.inner_;
   let ii = t % params.inner_;
   let base = oo * params.r_ * params.inner_ + ii;
+
   var acc: f32 = 0.0;
-  for (var r: u32 = 0u; r < params.r_; r = r + 1u) {
-    acc = acc + inp[base + r * params.inner_];
+  var k: u32 = lid.x;
+  while (k < params.r_) {
+    acc = acc + inp[base + k * params.inner_];
+    k = k + wg_size;
   }
-  if (params.is_mean == 1u) {
-    acc = acc / f32(params.r_);
+  partials[lid.x] = acc;
+  workgroupBarrier();
+
+  if (lid.x == 0u) {
+    var s: f32 = partials[0];
+    for (var i: u32 = 1u; i < wg_size; i = i + 1u) {
+      s = s + partials[i];
+    }
+    if (params.is_mean == 1u) {
+      s = s / f32(params.r_);
+    }
+    out[t] = s;
   }
-  out[t] = acc;
 }
 )";