sync 6fdf6301e2bb

Browse files

Files changed (5) hide show

README.md +2 -2
build/webgpu/elementwise-bias-gelu.wgsl.jinja +16 -7
build/webgpu/manifest.json +17 -65
build/webgpu/metadata.json +7 -9
build/webgpu/test.json +49 -1

README.md CHANGED Viewed

@@ -40,13 +40,13 @@ See the [ONNX Runtime `FastGelu` contrib-operator spec](https://github.com/micro
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`elementwise-bias-gelu.wgsl.jinja`](build/webgpu/elementwise-bias-gelu.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`elementwise-bias-gelu.wgsl.jinja`](build/webgpu/elementwise-bias-gelu.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/elementwise-bias-gelu.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,11 @@
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
@@ -7,26 +15,27 @@
 // `vec4Tail` instead uses scalar bindings with four guarded lanes.
 // GELU uses the tanh approximation below, with its input clamped in the tails.
 fn tanh_safe(x: f32) -> f32 {
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
   return tanh(x);
 }
 fn gelu_value(v: f32) -> f32 {
   return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
 }
-{% if hasBias %}
 const HIDDEN: u32 = {{ hidden | default(0) }}u;
 {% endif %}
 @compute @workgroup_size({{ wg }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // per-axis dispatch fold width (outputs > 16.7M elements).
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wg }}u;
-  if (i >= params.count) {
-    return;
-  }
 {% if vec4Tail %}
   let base = i * 4u;
 {% for lane in range(4) %}

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
 // `vec4Tail` instead uses scalar bindings with four guarded lanes.
 // GELU uses the tanh approximation below, with its input clamped in the tails.
 fn tanh_safe(x: f32) -> f32 {
+  // tanh rounds to its saturated value for these tails in f32. Return that
+  // value directly, including for infinite input, before invoking the builtin.
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
+  // For tiny |x|, return x directly to preserve its sign and magnitude without
+  // relying on backend-specific builtin behavior near zero.
+  if (x > -1.0e-4 && x < 1.0e-4) { return x; }
   return tanh(x);
 }
 fn gelu_value(v: f32) -> f32 {
   return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
 }
+{% if hasBias %}
 const HIDDEN: u32 = {{ hidden | default(0) }}u;
 {% endif %}
 @compute @workgroup_size({{ wg }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(wg) }}
 {% if vec4Tail %}
   let base = i * 4u;
 {% for lane in range(4) %}

build/webgpu/manifest.json CHANGED Viewed

@@ -16,21 +16,16 @@
     "biasOk": "present.bias and ranks.bias == 1 and dim(shapes.bias, 0) == dim(shapes.X, ranks.X - 1)",
     "noBiasOk": "not present.bias",
     "vec4Ok": "numel(shapes.X) > 0 and numel(shapes.X) % 4 == 0 and dim(shapes.X, ranks.X - 1) % 4 == 0",
-    "scalar": "dtypes.T",
-    "approximate": "\"tanh\""
   },
   "when": ["baseOk"],
   "bindings": {
-    "x": { "arg": "X", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "y": { "arg": "Y", "buffer": "storage", "elementType": "$vectorScalar" },
-    "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X) / 4" }] },
-    "x_2": { "arg": "X", "name": "x", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "y_2": { "arg": "Y", "name": "y", "buffer": "storage", "elementType": "$scalar" },
-    "params_2": {
-      "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }]
-    }
   },
   "variants": [
     {
@@ -47,7 +42,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "FastGelu.vec4Bias",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
           "bindings": ["x", { "arg": "bias", "elementType": "$scalar", "length": "$hidden" }, "y", "params"],
           "dispatch": {
@@ -66,7 +61,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "FastGelu.vec4",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
           "bindings": ["x", "y", "params"],
           "dispatch": {
@@ -80,19 +75,19 @@
     {
       "id": "vec4_tail_bias",
       "priority": 20,
-      "when": ["biasOk", "numel(shapes.X) > 0"],
       "derive": {
         "vec4": false,
-        "vec4Tail": true,
         "hasBias": true,
         "hidden": "dim(shapes.X, ranks.X - 1) if dim(shapes.X, ranks.X - 1) > 0 else 1"
       },
       "passes": [
         {
           "id": "main",
-          "name": "FastGelu.vec4TailBias",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
-          "bindings": ["x_2", "bias", "y_2", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
@@ -104,14 +99,14 @@
     {
       "id": "vec4_tail_no_bias",
       "priority": 15,
-      "when": ["noBiasOk", "numel(shapes.X) > 0"],
-      "derive": { "vec4": false, "vec4Tail": true, "hasBias": false },
       "passes": [
         {
           "id": "main",
-          "name": "FastGelu.vec4Tail",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
-          "bindings": ["x_2", "y_2", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
@@ -119,49 +114,6 @@
           }
         }
       ]
-    },
-    {
-      "id": "scalar_bias",
-      "priority": 10,
-      "when": ["biasOk", "true"],
-      "derive": {
-        "vec4": false,
-        "vec4Tail": false,
-        "hasBias": true,
-        "hidden": "dim(shapes.X, ranks.X - 1) if dim(shapes.X, ranks.X - 1) > 0 else 1"
-      },
-      "passes": [
-        {
-          "id": "main",
-          "name": "FastGelu.scalarBias",
-          "shader": "elementwise-bias-gelu.wgsl.jinja",
-          "bindings": ["x_2", "bias", "y_2", "params_2"],
-          "dispatch": {
-            "x": "min(ceilDiv((numel(shapes.X)), (workgroupSize)), 65535)",
-            "y": "ceilDiv(ceilDiv((numel(shapes.X)), (workgroupSize)), 65535)",
-            "z": 1
-          }
-        }
-      ]
-    },
-    {
-      "id": "scalar_no_bias",
-      "priority": 0,
-      "when": ["noBiasOk", "true"],
-      "derive": { "vec4": false, "vec4Tail": false, "hasBias": false },
-      "passes": [
-        {
-          "id": "main",
-          "name": "FastGelu.scalar",
-          "shader": "elementwise-bias-gelu.wgsl.jinja",
-          "bindings": ["x_2", "y_2", "params_2"],
-          "dispatch": {
-            "x": "min(ceilDiv((numel(shapes.X)), (workgroupSize)), 65535)",
-            "y": "ceilDiv(ceilDiv((numel(shapes.X)), (workgroupSize)), 65535)",
-            "z": 1
-          }
-        }
-      ]
     }
   ]
 }

     "biasOk": "present.bias and ranks.bias == 1 and dim(shapes.bias, 0) == dim(shapes.X, ranks.X - 1)",
     "noBiasOk": "not present.bias",
     "vec4Ok": "numel(shapes.X) > 0 and numel(shapes.X) % 4 == 0 and dim(shapes.X, ranks.X - 1) % 4 == 0",
+    "scalar": "dtypes.T"
   },
   "when": ["baseOk"],
   "bindings": {
+    "x": { "arg": "X", "elementType": "$vectorScalar" },
+    "y": { "arg": "Y", "elementType": "$vectorScalar" },
+    "params": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X) / 4" }] },
+    "x_x": { "arg": "X", "name": "x", "elementType": "$scalar" },
+    "y_y": { "arg": "Y", "name": "y", "elementType": "$scalar" },
+    "params_main": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
   },
   "variants": [
     {
       "passes": [
         {
           "id": "main",
+          "name": "FastGelu.Vec4Bias",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
           "bindings": ["x", { "arg": "bias", "elementType": "$scalar", "length": "$hidden" }, "y", "params"],
           "dispatch": {
       "passes": [
         {
           "id": "main",
+          "name": "FastGelu.Vec4",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
           "bindings": ["x", "y", "params"],
           "dispatch": {
     {
       "id": "vec4_tail_bias",
       "priority": 20,
+      "when": ["biasOk"],
       "derive": {
         "vec4": false,
+        "vec4Tail": "numel(shapes.X) > 0",
         "hasBias": true,
         "hidden": "dim(shapes.X, ranks.X - 1) if dim(shapes.X, ranks.X - 1) > 0 else 1"
       },
       "passes": [
         {
           "id": "main",
+          "name": "FastGelu.Vec4TailBias",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
+          "bindings": ["x_x", "bias", "y_y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
     {
       "id": "vec4_tail_no_bias",
       "priority": 15,
+      "when": ["noBiasOk"],
+      "derive": { "vec4": false, "vec4Tail": "numel(shapes.X) > 0", "hasBias": false },
       "passes": [
         {
           "id": "main",
+          "name": "FastGelu.Vec4Tail",
           "shader": "elementwise-bias-gelu.wgsl.jinja",
+          "bindings": ["x_x", "y_y", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.X), 4)), (workgroupSize)), 65535)",
           }
         }
       ]
     }
   ]
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.FastGelu",
-  "id": "_com_microsoft_fastgelu_webgpu_da3aabd",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,21 +8,19 @@
     "algorithm": "sha256",
     "files": {
       "bench.json": "uob2jmkWOUzFyVqQy0sJhB7JkQGvh+5ux8/7f+O6DYA=",
-      "elementwise-bias-gelu.wgsl.jinja": "OW2nqbCYsSLMd8w5lKcNx4zjdkHi4i2kB+FFAKyHhFg=",
-      "manifest.json": "nHE7yhM3Q5obzXWntT8sXryEtqYicMVDVSqdP+kwbLs=",
-      "test.json": "oH8XR6itu84Ti3GZbKE8ImRXTl4N2njDh1qRw0+HhCs="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "vec4_bias": ["elementwise-bias-gelu.wgsl.jinja"],
       "vec4_no_bias": ["elementwise-bias-gelu.wgsl.jinja"],
       "vec4_tail_bias": ["elementwise-bias-gelu.wgsl.jinja"],
-      "vec4_tail_no_bias": ["elementwise-bias-gelu.wgsl.jinja"],
-      "scalar_bias": ["elementwise-bias-gelu.wgsl.jinja"],
-      "scalar_no_bias": ["elementwise-bias-gelu.wgsl.jinja"]
     }
   }
 }

 {
   "name": "com.microsoft.FastGelu",
+  "id": "_com_microsoft_fastgelu_webgpu_aa8ecda",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "uob2jmkWOUzFyVqQy0sJhB7JkQGvh+5ux8/7f+O6DYA=",
+      "elementwise-bias-gelu.wgsl.jinja": "RN55dM80C5gNCEsQC/n2ToecW/TW++sB2neIoraM05o=",
+      "manifest.json": "mUY8bymm2MAiLszkBwz6y1s+DT+ulKtK+MzMvKG8yyM=",
+      "test.json": "JgD1Mpetx475VrIbLzMW6x4yhiQIlCbYGKG1bTLptJ4="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "vec4_bias": ["elementwise-bias-gelu.wgsl.jinja"],
       "vec4_no_bias": ["elementwise-bias-gelu.wgsl.jinja"],
       "vec4_tail_bias": ["elementwise-bias-gelu.wgsl.jinja"],
+      "vec4_tail_no_bias": ["elementwise-bias-gelu.wgsl.jinja"]
     }
   }
 }

build/webgpu/test.json CHANGED Viewed

@@ -76,7 +76,7 @@
     {
       "name": "f32_scalar_no_bias_zero_sequence",
       "provenance": {
-        "notes": "An empty bias-free input bypasses both vec4 routes and selects the scalar no-bias kernel."
       },
       "inputs": { "X": { "dtype": "float32", "shape": [1, 0, 4], "data": { "kind": "values", "values": [] } } },
       "outputs": { "Y": { "dtype": "float32", "shape": [1, 0, 4], "data": { "kind": "values", "values": [] } } }
@@ -442,6 +442,54 @@
           "data": { "kind": "values", "values": [-0.04540231, -0.15880801, 0.0, 0.84119199, 1.95459769] }
         }
       }
     }
   ]
 }

     {
       "name": "f32_scalar_no_bias_zero_sequence",
       "provenance": {
+        "notes": "An empty float32 input with no bias must produce an empty output without reading an absent bias."
       },
       "inputs": { "X": { "dtype": "float32", "shape": [1, 0, 4], "data": { "kind": "values", "values": [] } } },
       "outputs": { "Y": { "dtype": "float32", "shape": [1, 0, 4], "data": { "kind": "values", "values": [] } } }
           "data": { "kind": "values", "values": [-0.04540231, -0.15880801, 0.0, 0.84119199, 1.95459769] }
         }
       }
+    },
+    {
+      "name": "tanh_near_zero_7",
+      "provenance": {
+        "notes": "Both sides of the tiny tanh input threshold, with scalar-tail and aligned vec4 shapes. Expected values use the mathematical tanh approximation."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [7],
+          "data": { "kind": "values", "values": [-0.000125, -0.00012, -1e-8, 0.0, 1e-8, 0.00012, 0.000125] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [7],
+          "data": {
+            "kind": "values",
+            "values": [-0.00006249376652688504, -0.000059994255231176074, -4.999999960105772e-9, 0.0, 5.000000039894228e-9, 0.00006000574476882393, 0.00006250623347311496]
+          },
+          "tolerance": 1e-10
+        }
+      }
+    },
+    {
+      "name": "tanh_near_zero_8",
+      "provenance": {
+        "notes": "Both sides of the tiny tanh input threshold, with scalar-tail and aligned vec4 shapes. Expected values use the mathematical tanh approximation."
+      },
+      "inputs": {
+        "X": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": { "kind": "values", "values": [-0.000125, -0.00012, -1e-8, 0.0, 1e-8, 0.00012, 0.000125, 0.00013] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": {
+            "kind": "values",
+            "values": [-0.00006249376652688504, -0.000059994255231176074, -4.999999960105772e-9, 0.0, 5.000000039894228e-9, 0.00006000574476882393, 0.00006250623347311496, 0.0000650067421245197]
+          },
+          "tolerance": 1e-10
+        }
+      }
     }
   ]
 }