Add an aten.erf kernel (#24908)

backends/webgpu/partitioner/partitioner.py

@@ -116,6 +116,7 @@ def _fp32_unary_features() -> OpFeatures:
         dtypes={torch.float32},
     ),
     exir_ops.edge.aten.elu.default: _fp32_unary_features(),
+    exir_ops.edge.aten.erf.default: _fp32_unary_features(),
     exir_ops.edge.aten.silu.default: _fp32_unary_features(),
 }
 

backends/webgpu/runtime/WebGPUShaderRegistry.cpp

@@ -136,6 +136,7 @@
 #include <executorch/backends/webgpu/runtime/ops/unary/clamp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/cos_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/elu_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/erf_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/exp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/neg_wgsl.h>
@@ -159,7 +160,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 141> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 142> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -489,6 +490,13 @@ constexpr std::array<WebGPUShaderInfo, 141> kShaderRegistry = {{
         kEmbeddingQ4gswWorkgroupSizeY,
         kEmbeddingQ4gswWorkgroupSizeZ,
     },
+    {
+        "erf",
+        kErfWGSL,
+        kErfWorkgroupSizeX,
+        kErfWorkgroupSizeY,
+        kErfWorkgroupSizeZ,
+    },
     {
         "et_vk_sdpa_av",
         kEtVkSdpaAvWGSL,

backends/webgpu/runtime/ops/unary/Activations.cpp

@@ -13,6 +13,7 @@
 #include <executorch/backends/webgpu/runtime/ops/unary/clamp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/cos_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/elu_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/erf_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/exp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/neg_wgsl.h>
@@ -129,6 +130,11 @@ void elu_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       input_scale);
 }
 
+void erf_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_unary_op(
+      graph, args.at(0), args.at(1), kErfWGSL, kErfWorkgroupSizeX, "erf");
+}
+
 void clamp_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   // aten.clamp.default args: [in, min, max, out]; min/max None -> +/-inf.
   const float lo = get_val_or_inf(graph, args.at(1), /*is_max=*/false);
@@ -187,6 +193,7 @@ WEBGPU_REGISTER_OPERATORS {
   WEBGPU_REGISTER_OP(aten.hardswish.default, hardswish_impl);
   WEBGPU_REGISTER_OP(aten.silu.default, silu_impl);
   WEBGPU_REGISTER_OP(aten.elu.default, elu_impl);
+  WEBGPU_REGISTER_OP(aten.erf.default, erf_impl);
   WEBGPU_REGISTER_OP(aten.clamp.default, clamp_impl);
   WEBGPU_REGISTER_OP(aten.hardtanh.default, hardtanh_impl);
   WEBGPU_REGISTER_OP(aten.pow.Tensor_Scalar, pow_scalar_impl);

backends/webgpu/runtime/ops/unary/erf.wgsl

@@ -0,0 +1,44 @@
+@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 = 256u;
+
+// erf(x) as an odd/even rational fit on [-4, 4] (Eigen's fast erf for float),
+// within 5e-7 of the exact value with small relative error near 0. Past |x| = 4
+// erf rounds to +-1 in fp32. NaN is passed through explicitly: clamp() of NaN
+// is implementation-defined, and WGSL may assume float compares never see NaN,
+// so the check reads the bits.
+fn erf_rational(x: f32) -> f32 {
+    let c = clamp(x, -4.0, 4.0);
+    let c2 = c * c;
+    var p = -2.72614225801306e-10;
+    p = p * c2 + 2.77068142495902e-08;
+    p = p * c2 - 2.10102402082508e-06;
+    p = p * c2 - 5.69250639462346e-05;
+    p = p * c2 - 7.34990630326855e-04;
+    p = p * c2 - 2.95459980854025e-03;
+    p = p * c2 - 1.60960333262415e-02;
+    var q = -1.45660718464996e-05;
+    q = q * c2 - 2.13374055278905e-04;
+    q = q * c2 - 1.68282697438203e-03;
+    q = q * c2 - 7.37332916720468e-03;
+    q = q * c2 - 1.42647390514189e-02;
+    let is_nan = (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
+    return select(c * p / q, x, is_nan);
+}
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    output[idx] = erf_rational(input[idx]);
+}

backends/webgpu/runtime/ops/unary/erf_wgsl.h

@@ -0,0 +1,68 @@
+/*
+ * 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 erf.wgsl - DO NOT EDIT.
+// wgsl-sha256: ad547be9895644875ffdc4db9cdf28373ea2125f4c9ca54bf613c6d62fffc87a
+inline constexpr const char* kErfWGSL = 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 = 256u;
+
+// erf(x) as an odd/even rational fit on [-4, 4] (Eigen's fast erf for float),
+// within 5e-7 of the exact value with small relative error near 0. Past |x| = 4
+// erf rounds to +-1 in fp32. NaN is passed through explicitly: clamp() of NaN
+// is implementation-defined, and WGSL may assume float compares never see NaN,
+// so the check reads the bits.
+fn erf_rational(x: f32) -> f32 {
+    let c = clamp(x, -4.0, 4.0);
+    let c2 = c * c;
+    var p = -2.72614225801306e-10;
+    p = p * c2 + 2.77068142495902e-08;
+    p = p * c2 - 2.10102402082508e-06;
+    p = p * c2 - 5.69250639462346e-05;
+    p = p * c2 - 7.34990630326855e-04;
+    p = p * c2 - 2.95459980854025e-03;
+    p = p * c2 - 1.60960333262415e-02;
+    var q = -1.45660718464996e-05;
+    q = q * c2 - 2.13374055278905e-04;
+    q = q * c2 - 1.68282697438203e-03;
+    q = q * c2 - 7.37332916720468e-03;
+    q = q * c2 - 1.42647390514189e-02;
+    let is_nan = (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
+    return select(c * p / q, x, is_nan);
+}
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    output[idx] = erf_rational(input[idx]);
+}
+)";
+
+inline constexpr uint32_t kErfWorkgroupSizeX = 256;
+inline constexpr uint32_t kErfWorkgroupSizeY = 1;
+inline constexpr uint32_t kErfWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu

backends/webgpu/test/BUCK

@@ -39,6 +39,7 @@ fbcode_target(
     srcs = [
         "ops/test_conv_transpose1d.py",
         "ops/test_elu_silu.py",
+        "ops/test_erf.py",
         "test_webgpu_partitioner.py",
     ],
     deps = [

backends/webgpu/test/native/test_dynamic_shape.cpp

@@ -21,6 +21,7 @@
 //   O  ONE dyn_conv1d graph reused across live input lengths
 //   P  ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap
 //   Q  dyn_layer_norm(_rows)   R dyn_mean_keepdim   S dyn_elu_silu (reused)
+//   S2 dyn_erf_gelu (reused)
 // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py.
 //
 // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else
@@ -1960,6 +1961,16 @@ TEST(DynamicShape, EluSiluReusedGraph) {
   }
 }
 
+// S2: x * 0.5 * (1 + erf(x / sqrt(2))) on ONE loaded graph, shrinking and
+// growing S.
+TEST(DynamicShape, ErfGeluReusedGraph) {
+  Module m(g_dir + "/dyn_erf_gelu.pte");
+  ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_erf_gelu.pte";
+  for (int s : {128, 7, 1, 125, 7}) {
+    check_s(m, "dyn_erf_gelu", s);
+  }
+}
+
 TEST(DynamicShape, SwiGluReusedGraph) {
   Module module(g_dir + "/dyn_swiglu.pte");
   ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_swiglu.pte";

backends/webgpu/test/op_tests/cases.py

@@ -72,6 +72,12 @@
     signed_ramp,
     SiluModule,
 )
+from executorch.backends.webgpu.test.ops.test_erf import (
+    ErfGeluModule,
+    ErfModule,
+    signed_small,
+    wide_ramp,
+)
 from executorch.backends.webgpu.test.ops.test_expand_copy import (
     CONFIGS as _EXPAND_COPY_CONFIGS,
     ExpandCopyModule,
@@ -2632,6 +2638,49 @@ def _silu_suite() -> WebGPUTestSuite:
     )
 
 
+@register_op_test("erf")
+def _erf_suite() -> WebGPUTestSuite:
+    return WebGPUTestSuite(
+        module_factory=ErfModule,
+        cases=[
+            Case(name="wide", inputs=(InputSpec(shape=(4, 257), gen=wide_ramp),)),
+            # atol=0: near 0 only a relative-accurate kernel passes.
+            Case(
+                name="small_relative",
+                inputs=(InputSpec(shape=(2, 300), gen=signed_small),),
+                atol=0.0,
+            ),
+            Case(name="rank4_randn", inputs=((S1, S2, S2, M2),)),
+            Case(
+                name="param_input",
+                construct={"param_shape": (4, 257)},
+                inputs=((4, 257),),
+            ),
+            Case(
+                name="extreme_values",
+                inputs=(InputSpec(shape=(4, 10), gen=extreme_values),),
+                equal_nan=True,
+            ),
+        ],
+        atol=1e-6,
+        rtol=1e-5,
+    )
+
+
+@register_op_test("erf_gelu")
+def _erf_gelu_suite() -> WebGPUTestSuite:
+    # The hand-written exact GELU of HF ESM-2 / BERT: div, erf, add, mul, mul.
+    return WebGPUTestSuite(
+        module_factory=ErfGeluModule,
+        cases=[
+            Case(name="range", inputs=(InputSpec(shape=(3, 480), gen=wide_ramp),)),
+            Case(name="randn", inputs=((1, 37, 480),)),
+        ],
+        atol=1e-5,
+        rtol=1e-4,
+    )
+
+
 @register_op_test("conv_transpose1d")
 def _conv_transpose1d_suite() -> WebGPUTestSuite:
     # 3-D transposed aten.convolution, run on the transposed 2-D kernels (H = 1).

backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py

@@ -22,6 +22,7 @@
 from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
 from executorch.backends.webgpu.partitioner import WebGPUPartitioner
 from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule
+from executorch.backends.webgpu.test.ops.test_erf import ErfGeluModule
 from executorch.backends.webgpu.test.ops.test_gelu import GeluModule
 from executorch.exir import to_edge_transform_and_lower
 from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes
@@ -440,6 +441,18 @@ def export_dynamic_expand_copy_rejection_case(out_dir: str) -> None:
     )
 
 
+def export_dynamic_erf_gelu_cases(out_dir: str, s_dim) -> None:
+    erf_gelu = ErfGeluModule()
+    _export(
+        erf_gelu,
+        (_ramp((1, 1, MAXS, HIDDEN)),),
+        {"x": {2: s_dim}},
+        os.path.join(out_dir, "dyn_erf_gelu.pte"),
+        partitioner_cls=WebGPUPartitioner,
+    )
+    _write_goldens(erf_gelu, "dyn_erf_gelu", out_dir, [MAXS, 125, 7, 1])
+
+
 def export_dynamic_shape_cases(out_dir: str) -> None:
     """Write the dynamic + static .pte's and per-S goldens for the native test."""
     os.makedirs(out_dir, exist_ok=True)
@@ -555,6 +568,9 @@ def export_dynamic_shape_cases(out_dir: str) -> None:
     )
     _write_goldens(elu_silu, "dyn_elu_silu", out_dir, [MAXS, 32, 1])
 
+    # 2h.4) The hand-written exact GELU of ESM-2 / BERT (erf) with a DYNAMIC S.
+    export_dynamic_erf_gelu_cases(out_dir, s_dim)
+
     # 2i) Dynamic SwiGLU graph patterns. The Llama-width fixture forces a 2D
     # fused dispatch at M=512; compact variants isolate commutative matching and
     # graph/intermediate ownership guards around same-input q4 projections.
@@ -1769,6 +1785,9 @@ def test_export_dynamic_rms(self) -> None:
                 "dyn_elu_silu.pte",
                 "dyn_elu_silu.S32.input.bin",
                 "dyn_elu_silu.S32.golden.bin",
+                "dyn_erf_gelu.pte",
+                "dyn_erf_gelu.S125.input.bin",
+                "dyn_erf_gelu.S7.golden.bin",
                 "dyn_linear_bk64.pte",
                 "dyn_linear_bk64.S512.input.bin",
                 "dyn_linear_bk64.S512.golden.bin",

backends/webgpu/test/ops/test_erf.py

@@ -0,0 +1,61 @@
+# 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.
+
+"""`aten.erf.default` modules + inputs for the WebGPU tests.
+
+Imported by `cases.py`, `test_webgpu_partitioner.py` and the dynamic-shape
+export. `ErfGeluModule` is the hand-written exact GELU of Hugging Face ESM-2
+and BERT-style models, x * 0.5 * (1 + erf(x / sqrt(2))).
+"""
+
+import math
+
+import torch
+
+
+def _numel(shape) -> int:
+    n = 1
+    for d in shape:
+        n *= d
+    return n
+
+
+def wide_ramp(shape) -> torch.Tensor:
+    """Deterministic linspace over [-10, 10]: both tails saturate to +-1."""
+    return torch.linspace(-10.0, 10.0, _numel(shape), dtype=torch.float32).reshape(
+        shape
+    )
+
+
+def signed_small(shape) -> torch.Tensor:
+    """+-geomspace(1e-30, 1), alternating sign, where erf(x) ~ 2x / sqrt(pi)
+    and only a relative-accurate kernel matches."""
+    n = _numel(shape)
+    mag = torch.logspace(-30.0, 0.0, n, dtype=torch.float64)
+    sign = torch.where(torch.arange(n) % 2 == 0, 1.0, -1.0).double()
+    return (sign * mag).float().reshape(shape)
+
+
+class ErfModule(torch.nn.Module):
+    """aten.erf; with `param_shape`, on a parameter added to the input."""
+
+    def __init__(self, param_shape=None) -> None:
+        super().__init__()
+        self.param = (
+            torch.nn.Parameter(wide_ramp(param_shape), requires_grad=False)
+            if param_shape is not None
+            else None
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.param is not None:
+            return torch.erf(self.param) + x
+        return torch.erf(x)
+
+
+class ErfGeluModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))