diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 48337d18da1..36626b173e0 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/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 diff --git a/backends/webgpu/runtime/ops/reduce/Reduce.cpp b/backends/webgpu/runtime/ops/reduce/Reduce.cpp new file mode 100644 index 00000000000..b276354d2a8 --- /dev/null +++ b/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 +#include +#include +#include + +#include + +#include +#include +#include +#include + +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& dims, + int64_t dim, + uint32_t& outer, + uint32_t& r, + uint32_t& inner) { + const int64_t ndim = static_cast(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(dims[d]); + } + for (int64_t d = dim + 1; d < ndim; ++d) { + in *= static_cast(dims[d]); + } + outer = static_cast(o); + r = static_cast(dims[dim]); + inner = static_cast(in); +} + +void reduce_impl( + WebGPUGraph& graph, + const std::vector& 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& 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(d); + } + const uint64_t outputs = static_cast(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(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, ¶ms, 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(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(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(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(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(d.size()); + int64_t rd = dim < 0 ? dim + nd : dim; + std::vector 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& args) { + reduce_impl(graph, args, /*is_mean=*/false, "sum.dim_IntList"); +} + +void mean_dim_impl(WebGPUGraph& graph, const std::vector& 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 diff --git a/backends/webgpu/runtime/ops/reduce/reduce.wgsl b/backends/webgpu/runtime/ops/reduce/reduce.wgsl new file mode 100644 index 00000000000..81c4f63c753 --- /dev/null +++ b/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 inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var 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 partials: array; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 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; + } +} diff --git a/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h b/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h new file mode 100644 index 00000000000..f3b06ca8b33 --- /dev/null +++ b/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h @@ -0,0 +1,80 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from reduce.wgsl - DO NOT EDIT. +// wgsl-sha256: b4a40b67af55986ea5136f06c260cf3625e4db7809975fc616aeace16b479b2d +inline constexpr const char* kReduceWGSL = R"( +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + is_mean: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var 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 partials: array; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 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; + } +} +)"; + +inline constexpr uint32_t kReduceWorkgroupSizeX = 256; +inline constexpr uint32_t kReduceWorkgroupSizeY = 1; +inline constexpr uint32_t kReduceWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu