diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index b12763afd6a..441b572426b 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -52,6 +52,7 @@ set(WEBGPU_SRCS runtime/ops/cat/Cat.cpp runtime/ops/index/Index.cpp runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp + runtime/ops/gelu/Gelu.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/gelu/Gelu.cpp b/backends/webgpu/runtime/ops/gelu/Gelu.cpp new file mode 100644 index 00000000000..943012515ff --- /dev/null +++ b/backends/webgpu/runtime/ops/gelu/Gelu.cpp @@ -0,0 +1,100 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform buffer layout matching the WGSL Params struct; 16-byte aligned. +struct GeluParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; + +// aten.gelu.default args: [in, approximate, out] (mirrors Vulkan UnaryOp.cpp +// gelu — args[1] is the `approximate` string). approximate="none" selects the +// exact (erf) entry point; anything else (e.g. "tanh") selects the tanh +// approximation entry point. +void gelu_impl(WebGPUGraph& graph, const std::vector& args) { + const int in_id = args.at(0); + const int out_id = args.at(2); + const bool exact = graph.get_string(args.at(1)) == "none"; + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu"); + + uint32_t num_elements = + static_cast(out_tensor.nbytes / sizeof(float)); + + // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom). + uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); + uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_vec4_threads, wg_size, "gelu"); + + WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size); + + GeluParams params = {}; + params.num_elements = num_elements; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(GeluParams)); + graph.add_uniform_buffer_bytes(sizeof(GeluParams)); + + // input (read storage) + output (storage) + params. The exact/approximate + // choice is baked into the compiled pipeline via the entry point (mirrors + // onnxruntime's WebGPU EP), not a per-invocation select() — `main_tanh`/ + // `main_erf` each contain only their own formula, no double-eval. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kGeluWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(GeluParams)}, + }, + &wg_constant, + 1, + exact ? "main_erf" : "main_tanh"); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + // Drop our ref; the bind group keeps the uniform buffer alive until release. + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.gelu.default, gelu_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/gelu/gelu.wgsl b/backends/webgpu/runtime/ops/gelu/gelu.wgsl new file mode 100644 index 00000000000..4f7eb68bc96 --- /dev/null +++ b/backends/webgpu/runtime/ops/gelu/gelu.wgsl @@ -0,0 +1,65 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Abramowitz & Stegun 7.1.26 erf approximation (max abs err ~1.5e-7). +fn erf_approx4(x: vec4) -> vec4 { + let s = sign(x); + let a = abs(x); + let t = 1.0 / (1.0 + 0.3275911 * a); + let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-a * a); + return s * y; +} + +fn gelu_tanh4(x: vec4) -> vec4 { + let inner = clamp(0.7978845608028654 * (x + 0.044715 * x * x * x), vec4(-10.0), vec4(10.0)); + return 0.5 * x * (1.0 + tanh(inner)); +} + +fn gelu_erf4(x: vec4) -> vec4 { + return 0.5 * x * (1.0 + erf_approx4(x * 0.7071067811865476)); +} + +// Vec4 body + scalar-tail idiom (num_elements isn't guaranteed % 4 == 0 for +// arbitrary FFN activation shapes): each thread gathers up to 4 elements via +// select()-guarded scalar loads (the out-of-bounds lanes read a WebGPU-spec- +// clamped, NOT zero, value — but that value is discarded by the same select() +// before use), computes GELU as one vec4 op, then scatters back only the +// in-bounds lanes. +@compute @workgroup_size(wg_size, 1, 1) +fn main_tanh(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_tanh4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main_erf(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_erf4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} diff --git a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h new file mode 100644 index 00000000000..6da12e229af --- /dev/null +++ b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h @@ -0,0 +1,89 @@ +/* + * 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 gelu.wgsl - DO NOT EDIT. +// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f +inline constexpr const char* kGeluWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Abramowitz & Stegun 7.1.26 erf approximation (max abs err ~1.5e-7). +fn erf_approx4(x: vec4) -> vec4 { + let s = sign(x); + let a = abs(x); + let t = 1.0 / (1.0 + 0.3275911 * a); + let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-a * a); + return s * y; +} + +fn gelu_tanh4(x: vec4) -> vec4 { + let inner = clamp(0.7978845608028654 * (x + 0.044715 * x * x * x), vec4(-10.0), vec4(10.0)); + return 0.5 * x * (1.0 + tanh(inner)); +} + +fn gelu_erf4(x: vec4) -> vec4 { + return 0.5 * x * (1.0 + erf_approx4(x * 0.7071067811865476)); +} + +// Vec4 body + scalar-tail idiom (num_elements isn't guaranteed % 4 == 0 for +// arbitrary FFN activation shapes): each thread gathers up to 4 elements via +// select()-guarded scalar loads (the out-of-bounds lanes read a WebGPU-spec- +// clamped, NOT zero, value — but that value is discarded by the same select() +// before use), computes GELU as one vec4 op, then scatters back only the +// in-bounds lanes. +@compute @workgroup_size(wg_size, 1, 1) +fn main_tanh(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_tanh4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main_erf(@builtin(global_invocation_id) gid: vec3) { + let base = gid.x * 4u; + if (base >= params.num_elements) { + return; + } + let x1 = select(0.0, input[base + 1u], base + 1u < params.num_elements); + let x2 = select(0.0, input[base + 2u], base + 2u < params.num_elements); + let x3 = select(0.0, input[base + 3u], base + 3u < params.num_elements); + let y = gelu_erf4(vec4(input[base], x1, x2, x3)); + output[base] = y.x; + if (base + 1u < params.num_elements) { output[base + 1u] = y.y; } + if (base + 2u < params.num_elements) { output[base + 2u] = y.z; } + if (base + 3u < params.num_elements) { output[base + 3u] = y.w; } +} +)"; + +inline constexpr uint32_t kGeluWorkgroupSizeX = 64; +inline constexpr uint32_t kGeluWorkgroupSizeY = 1; +inline constexpr uint32_t kGeluWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu