[ExecuTorch][WebGPU] Add reduce op (sum/mean) to the WebGPU backend

Pull Request resolved: #20923

Add an `aten.sum.dim_IntList` / `aten.mean.dim` reduction handler for the training tail.

Key changes:
- `runtime/ops/reduce/` — per-row reduction WGSL kernel + handler
- `CMakeLists.txt` `WEBGPU_SRCS` — wire the source

Reuses the shared Vulkan partitioner (`aten.sum.dim_IntList`/`mean.dim` already registered); WebGPU kernel only.

Co-authored-with: Claude Code.
ghstack-source-id: 405026057
@exported-using-ghexport

Differential Revision: [D111755133](https://our.internmc.facebook.com/intern/diff/D111755133/)

Author: JCNTH

backends/webgpu/CMakeLists.txt

@@ -56,6 +56,7 @@ set(WEBGPU_SRCS
     runtime/ops/log_softmax/LogSoftmax.cpp
     runtime/ops/softmax/Softmax.cpp
     runtime/ops/bmm/Bmm.cpp
+    runtime/ops/reduce/Reduce.cpp
     runtime/ops/div/BinaryOp.cpp
     runtime/ops/sub/BinaryOp.cpp
     runtime/ops/linear/Linear.cpp

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

@@ -0,0 +1,273 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Uniform layout matching the WGSL Params struct (16-byte aligned).
+struct ReduceParams {
+  uint32_t outer;
+  uint32_t r;
+  uint32_t inner;
+  uint32_t is_mean;
+};
+static_assert(sizeof(ReduceParams) == 16, "ReduceParams must be 16 bytes");
+
+void decompose(
+    const std::vector<int64_t>& dims,
+    int64_t dim,
+    uint32_t& outer,
+    uint32_t& r,
+    uint32_t& inner) {
+  const int64_t ndim = static_cast<int64_t>(dims.size());
+  if (dim < 0) {
+    dim += ndim;
+  }
+  if (ndim == 0 || dim < 0 || dim >= ndim) {
+    throw std::runtime_error("WebGPU reduce: dim out of range");
+  }
+  uint64_t o = 1, in = 1;
+  for (int64_t d = 0; d < dim; ++d) {
+    o *= static_cast<uint64_t>(dims[d]);
+  }
+  for (int64_t d = dim + 1; d < ndim; ++d) {
+    in *= static_cast<uint64_t>(dims[d]);
+  }
+  outer = static_cast<uint32_t>(o);
+  r = static_cast<uint32_t>(dims[dim]);
+  inner = static_cast<uint32_t>(in);
+}
+
+void reduce_impl(
+    WebGPUGraph& graph,
+    const std::vector<int>& args,
+    bool is_mean,
+    const char* op_name) {
+  const int in_id = args.at(0);
+  const int dim_id = args.at(1);
+  const int keepdim_id = args.at(2);
+  const int out_id = args.at(args.size() - 1);
+
+  WGPUDevice device = graph.device();
+  const auto& in = graph.get_tensor(in_id);
+  const auto& out = graph.get_tensor(out_id);
+
+  bool keepdim = false;
+  if (graph.get_value_type(keepdim_id) == WebGPUGraph::ValueType::Int) {
+    keepdim = graph.get_int(keepdim_id) != 0;
+  }
+
+  if (in.dims.empty()) {
+    throw std::runtime_error("WebGPU reduce: scalar input unsupported");
+  }
+  if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::IntList) {
+    throw std::runtime_error("WebGPU reduce: dim arg is not an IntList");
+  }
+  const std::vector<int64_t>& reduce_dims = graph.get_int_list(dim_id);
+  // Single-dim reduction only for now; multi-dim is a tracked extension.
+  if (reduce_dims.size() != 1) {
+    throw std::runtime_error(
+        "WebGPU reduce: only single-dim reduction is supported");
+  }
+  const int64_t dim = reduce_dims[0];
+
+  uint32_t outer = 0, r = 0, inner = 0;
+  decompose(in.dims, dim, outer, r, inner);
+  if (outer == 0 || r == 0 || inner == 0) {
+    throw std::runtime_error("WebGPU reduce: zero-sized reduction");
+  }
+
+  uint64_t in_numel = 1;
+  for (int64_t d : in.dims) {
+    in_numel *= static_cast<uint64_t>(d);
+  }
+  const uint64_t outputs = static_cast<uint64_t>(outer) * inner;
+  if (in.nbytes != in_numel * sizeof(float) ||
+      out.nbytes != outputs * sizeof(float)) {
+    throw std::runtime_error("WebGPU reduce: fp32-only (byte-size mismatch)");
+  }
+  if (outputs > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU reduce: output count exceeds dispatch limit");
+  }
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kReduceWorkgroupSizeX);
+  // 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;
+  params.r = r;
+  params.inner = inner;
+  params.is_mean = is_mean ? 1u : 0u;
+
+  WGPUBufferDescriptor uniform_desc = {};
+  uniform_desc.size = sizeof(ReduceParams);
+  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  uniform_desc.mappedAtCreation = true;
+  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
+  void* mapped =
+      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(ReduceParams));
+  std::memcpy(mapped, &params, sizeof(ReduceParams));
+  wgpuBufferUnmap(uniform_buffer);
+  graph.add_uniform_buffer_bytes(sizeof(ReduceParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kReduceWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[3] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].binding = 1;
+  entries[1].visibility = WGPUShaderStage_Compute;
+  entries[1].buffer.type = WGPUBufferBindingType_Storage;
+  entries[2].binding = 2;
+  entries[2].visibility = WGPUShaderStage_Compute;
+  entries[2].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 3;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[3] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = in.buffer;
+  bg_entries[0].size = in.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = out.buffer;
+  bg_entries[1].size = out.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = uniform_buffer;
+  bg_entries[2].size = sizeof(ReduceParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 3;
+  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.x, op_name, workgroup_count.y});
+
+  // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch.
+  WGPUBuffer params_buf = uniform_buffer;
+  const uint32_t is_mean_u = is_mean ? 1u : 0u;
+  const uint64_t build_outputs = outputs;
+  graph.add_tensor_resize_hook(
+      in_id,
+      [in_id,
+       out_id,
+       dim,
+       keepdim,
+       is_mean_u,
+       build_outputs,
+       dispatch_idx,
+       params_buf](WebGPUGraph& g) {
+        const auto& d = g.cur_dims(in_id);
+        uint32_t o = 0, rr = 0, n = 0;
+        decompose(std::vector<int64_t>(d.begin(), d.end()), dim, o, rr, n);
+        if (o == 0u || rr == 0u || n == 0u) {
+          throw std::runtime_error("WebGPU reduce: live zero-sized reduction");
+        }
+        const uint64_t live_outputs = static_cast<uint64_t>(o) * n;
+        if (live_outputs > build_outputs) {
+          throw std::runtime_error(
+              "WebGPU reduce: live output count exceeds build max");
+        }
+        ReduceParams p = {};
+        p.outer = o;
+        p.r = rr;
+        p.inner = n;
+        p.is_mean = is_mean_u;
+        wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
+        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;
+        std::vector<int64_t> od;
+        for (int64_t i = 0; i < nd; ++i) {
+          if (i == rd) {
+            if (keepdim) {
+              od.push_back(1);
+            }
+          } else {
+            od.push_back(d[i]);
+          }
+        }
+        g.set_cur_dims(out_id, od);
+      });
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  // Graph owns it so the resize hook can rewrite it; freed in the dtor.
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+void sum_dim_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  reduce_impl(graph, args, /*is_mean=*/false, "sum.dim_IntList");
+}
+
+void mean_dim_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  reduce_impl(graph, args, /*is_mean=*/true, "mean.dim");
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.sum.dim_IntList, sum_dim_impl);
+  WEBGPU_REGISTER_OP(aten.mean.dim, mean_dim_impl);
+}
+
+} // namespace executorch::backends::webgpu

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

@@ -0,0 +1,56 @@
+struct Params {
+  outer_: u32,
+  r_: u32,
+  inner_: u32,
+  is_mean: u32,
+};
+
+@group(0) @binding(0) var<storage, read> inp: array<f32>;
+@group(0) @binding(1) var<storage, read_write> out: array<f32>;
+@group(0) @binding(2) var<uniform> params: 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(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;
+  var k: u32 = lid.x;
+  while (k < params.r_) {
+    acc = acc + inp[base + k * params.inner_];
+    k = k + wg_size;
+  }
+  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;
+  }
+}