[ExecuTorch][WebGPU] Add optimized gelu op

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})

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 <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/gelu/gelu_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+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<int>& 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<uint32_t>(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, &params, 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

backends/webgpu/runtime/ops/gelu/gelu.wgsl

@@ -0,0 +1,65 @@
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+}
+@group(0) @binding(2) var<uniform> 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<f32>) -> vec4<f32> {
+    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<f32>) -> vec4<f32> {
+    let inner = clamp(0.7978845608028654 * (x + 0.044715 * x * x * x), vec4<f32>(-10.0), vec4<f32>(10.0));
+    return 0.5 * x * (1.0 + tanh(inner));
+}
+
+fn gelu_erf4(x: vec4<f32>) -> vec4<f32> {
+    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<u32>) {
+    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<f32>(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<u32>) {
+    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<f32>(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; }
+}

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 <cstdint>
+
+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<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+}
+@group(0) @binding(2) var<uniform> 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<f32>) -> vec4<f32> {
+    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<f32>) -> vec4<f32> {
+    let inner = clamp(0.7978845608028654 * (x + 0.044715 * x * x * x), vec4<f32>(-10.0), vec4<f32>(10.0));
+    return 0.5 * x * (1.0 + tanh(inner));
+}
+
+fn gelu_erf4(x: vec4<f32>) -> vec4<f32> {
+    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<u32>) {
+    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<f32>(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<u32>) {
+    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<f32>(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