Xenova HF Staff commited on
Commit
53d66ee
·
verified ·
1 Parent(s): 9ca4c42

sync 91d990483a17

Browse files
README.md CHANGED
@@ -18,15 +18,15 @@ See the [ONNX Runtime `Gelu` contrib-operator spec](https://github.com/microsoft
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
- | --- | --- | --- | --- | --- | --- | --- |
23
- | `X` | `X` | `T` | — | — | Values transformed elementwise by the exact GELU activation. | required |
24
 
25
  ## Outputs
26
 
27
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
28
- | --- | --- | --- | --- | --- | --- | --- |
29
- | `Y` | `Y` | `T` | same as `X` | same as `X` | Output tensor after applying GELU; same shape as the input. | required |
30
 
31
  ## Type constraints
32
 
@@ -36,7 +36,7 @@ See the [ONNX Runtime `Gelu` contrib-operator spec](https://github.com/microsoft
36
 
37
  ## Files
38
 
39
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
40
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
41
  - [`test.json`](build/webgpu/test.json) — correctness cases
42
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -44,10 +44,14 @@ See the [ONNX Runtime `Gelu` contrib-operator spec](https://github.com/microsoft
44
 
45
  ## Use with `@huggingface/kernels`
46
 
47
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
48
- It then allocates the result tensors automatically.
 
 
 
49
 
50
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
51
 
52
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
53
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Logical dtype | Rank | Shape | Description | Presence |
22
+ | --- | --- | --- | --- | --- | --- |
23
+ | `X` | `T` | — | — | Values transformed elementwise by the exact GELU activation. | required |
24
 
25
  ## Outputs
26
 
27
+ | Name | Logical dtype | Rank | Shape | Description | Presence |
28
+ | --- | --- | --- | --- | --- | --- |
29
+ | `Y` | `T` | same as `X` | same as `X` | Output tensor after applying GELU; same shape as the input. | required |
30
 
31
  ## Type constraints
32
 
 
36
 
37
  ## Files
38
 
39
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
40
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
41
  - [`test.json`](build/webgpu/test.json) — correctness cases
42
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
44
 
45
  ## Use with `@huggingface/kernels`
46
 
47
+ ```sh
48
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
49
+ ```
50
+
51
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
52
 
53
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
54
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
55
 
56
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
57
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.Gelu",
3
  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
4
  "cases": [
5
  {
 
1
  {
 
2
  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
3
  "cases": [
4
  {
build/webgpu/elementwise-bias-gelu.wgsl.jinja CHANGED
@@ -1,26 +1,17 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {{ env.wgsl.resourceDeclarations }}
5
  {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
6
 
7
- // Bias plus GELU, with a specialization-selected tanh or erf approximation.
8
- // The optional bias is a rank-1 vector broadcast over the innermost (hidden)
9
- // axis: bias index = element_index % HIDDEN. The `vec4` path requires
10
- // HIDDEN % 4 == 0 and numel % 4 == 0 so a vec4 group never crosses the hidden
11
- // axis (the bias slice is then contiguous). `vec4Tail` keeps scalar bindings but
12
- // evaluates four guarded lanes per invocation, so odd hidden sizes retain the
13
- // same parallel efficiency without crossing row/bias boundaries. Gelu math and overflow guards match
14
- // the vectorized unary implementation: tanh saturates to +/-1 by |x|~9, and the
15
- // erf path uses the same rational approximation as Gelu.
16
- {% if approximate == "erf" %}
17
  fn erf_approx(x: f32) -> f32 {
18
  let ax = abs(x);
19
  // The polynomial has a small nonzero floor near zero. Use erf(x) ~=
20
  // 2/sqrt(pi)*x below 2^-20 to preserve erf(0) == 0, odd symmetry, and the
21
- // correctly rounded f32 result. The exactly representable threshold keeps
22
- // scalar and vector branching identical. NaN falls through to the polynomial
23
- // and propagates.
24
  if (ax < 9.5367431640625e-7) {
25
  return 1.1283791670955126 * x;
26
  }
@@ -29,62 +20,24 @@ fn erf_approx(x: f32) -> f32 {
29
  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-(ax * ax));
30
  return sign * y;
31
  }
32
- {% else %}
33
- fn tanh_safe(x: f32) -> f32 {
34
- if (x > 10.0) { return 1.0; }
35
- if (x < -10.0) { return -1.0; }
36
- return tanh(x);
37
- }
38
- {% endif %}
39
  fn gelu_value(v: f32) -> f32 {
40
- {% if approximate == "erf" %}
41
  return 0.5 * v * (1.0 + erf_approx(v * 0.7071067811865476));
42
- {% else %}
43
- return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
44
- {% endif %}
45
  }
46
- {% if hasBias %}
47
-
48
- const HIDDEN: u32 = {{ hidden }}u;
49
-
50
- {% endif %}
51
  @compute @workgroup_size({{ wg }})
52
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
53
  // 2D-folded flat index: gid.y carries the high bits past the
54
- // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
55
- let i = gid.x + gid.y * nwg.x * {{ wg }}u;
56
  if (i >= params.count) {
57
  return;
58
  }
59
- {% if vec4Tail %}
60
- let base = i * 4u;
61
- {% for lane in range(4) %}
62
- if (base + {{ lane }}u < params.count) {
63
- let xv{{ lane }} = f32(x[base + {{ lane }}u]);
64
- {% if hasBias %}
65
- let v{{ lane }} = xv{{ lane }} + f32(bias[(base + {{ lane }}u) % HIDDEN]);
66
- {% else %}
67
- let v{{ lane }} = xv{{ lane }};
68
- {% endif %}
69
- y[base + {{ lane }}u] = {{ scalar }}(gelu_value(v{{ lane }}));
70
- }
71
- {% endfor %}
72
- {% elif vec4 %}
73
  let xv = vec4<f32>(x[i]);
74
- {% if hasBias %}
75
- let bcol = (i * 4u) % HIDDEN;
76
- let v = xv + vec4<f32>(f32(bias[bcol]), f32(bias[bcol + 1u]), f32(bias[bcol + 2u]), f32(bias[bcol + 3u]));
77
- {% else %}
78
  let v = xv;
79
- {% endif %}
80
  y[i] = vec4<{{ scalar }}>(vec4<f32>(gelu_value(v.x), gelu_value(v.y), gelu_value(v.z), gelu_value(v.w)));
81
  {% else %}
82
  let xv = f32(x[i]);
83
- {% if hasBias %}
84
- let v = xv + f32(bias[i % HIDDEN]);
85
- {% else %}
86
  let v = xv;
87
- {% endif %}
88
  y[i] = {{ scalar }}(gelu_value(v));
89
  {% endif %}
90
  }
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
  {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
3
 
4
+ // Computes GELU, optionally after adding a rank-1 bias broadcast over the
5
+ // innermost axis. A present bias uses element_index % hidden_extent. The `vec4`
6
+ // path requires the innermost extent and element count to be divisible by four;
7
+ // `vec4Tail` instead uses scalar bindings with four guarded lanes.
8
+ // GELU uses the rational erf approximation defined below.
 
 
 
 
 
9
  fn erf_approx(x: f32) -> f32 {
10
  let ax = abs(x);
11
  // The polynomial has a small nonzero floor near zero. Use erf(x) ~=
12
  // 2/sqrt(pi)*x below 2^-20 to preserve erf(0) == 0, odd symmetry, and the
13
+ // small-input linear behavior. The threshold is exactly representable in
14
+ // f32. NaN falls through to the polynomial and propagates.
 
15
  if (ax < 9.5367431640625e-7) {
16
  return 1.1283791670955126 * x;
17
  }
 
20
  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-(ax * ax));
21
  return sign * y;
22
  }
 
 
 
 
 
 
 
23
  fn gelu_value(v: f32) -> f32 {
 
24
  return 0.5 * v * (1.0 + erf_approx(v * 0.7071067811865476));
 
 
 
25
  }
 
 
 
 
 
26
  @compute @workgroup_size({{ wg }})
27
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
28
  // 2D-folded flat index: gid.y carries the high bits past the
29
+ // per-axis dispatch fold width (outputs > 16.7M elements).
30
+ let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wg }}u;
31
  if (i >= params.count) {
32
  return;
33
  }
34
+ {% if vec4 %}
 
 
 
 
 
 
 
 
 
 
 
 
 
35
  let xv = vec4<f32>(x[i]);
 
 
 
 
36
  let v = xv;
 
37
  y[i] = vec4<{{ scalar }}>(vec4<f32>(gelu_value(v.x), gelu_value(v.y), gelu_value(v.z), gelu_value(v.w)));
38
  {% else %}
39
  let xv = f32(x[i]);
 
 
 
40
  let v = xv;
 
41
  y[i] = {{ scalar }}(gelu_value(v));
42
  {% endif %}
43
  }
build/webgpu/manifest.json CHANGED
@@ -2,99 +2,55 @@
2
  "domain": "com.microsoft",
3
  "name": "Gelu",
4
  "sinceVersion": 1,
5
- "description": "Applies the Gaussian Error Linear Unit (GELU) activation elementwise: `Y = 0.5 * X * (1 + erf(X / sqrt(2)))`. The output has the same shape as the input. Float16 and float32 are supported; the schema's double and bfloat16 types are not.",
6
- "inputs": [
7
- { "role": "X", "dtype": "T", "description": "Values transformed elementwise by the exact GELU activation." }
8
- ],
9
- "outputs": [
10
- {
11
- "role": "Y",
12
- "dtype": "T",
13
- "rank": "ranks.X",
14
- "shape": "shapes.X",
15
- "description": "Output tensor after applying GELU; same shape as the input."
16
- }
17
- ],
18
  "typeConstraints": { "T": ["float32", "float16"] },
19
- "args": {
20
- "X": { "kind": "tensor", "semantic": "X", "role": "input" },
21
- "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
22
- },
23
- "tunables": { "WORKGROUP_SIZE": 256 },
24
- "constants": {
25
- "scalar": "dtypes.T",
26
- "usesF16": "dtypes.T == \"f16\"",
27
- "approximate": "\"erf\"",
28
- "vec4Tail": false,
29
- "hasBias": false
30
- },
31
  "variants": [
32
  {
33
  "id": "vec4",
34
  "priority": 20,
35
- "when": ["numel(shapes.X) == numel(shapes.Y)", "f16Ok(dtypes.T)", "numel(shapes.X) > 0", "numel(shapes.X) % 4 == 0"],
36
- "constants": { "vec4": true, "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
37
  "passes": [
38
  {
39
  "id": "main",
40
  "name": "Gelu.vec4",
41
  "shader": "elementwise-bias-gelu.wgsl.jinja",
42
  "bindings": [
43
- {
44
- "name": "x",
45
- "arg": "X",
46
- "semantic": "X",
47
- "buffer": { "type": "read-only-storage" },
48
- "elementType": "$vectorScalar"
49
- },
50
- {
51
- "name": "y",
52
- "arg": "Y",
53
- "semantic": "Y",
54
- "buffer": { "type": "storage" },
55
- "elementType": "$vectorScalar"
56
- },
57
- {
58
- "name": "params",
59
- "semantic": "kernel.params",
60
- "buffer": { "type": "uniform" },
61
- "struct": {
62
- "name": "Params",
63
- "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X) / 4" }]
64
- }
65
- }
66
  ],
67
- "dispatch": { "threads": "numel(shapes.X) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
 
 
 
 
68
  }
69
  ]
70
  },
71
  {
72
  "id": "scalar",
73
  "priority": 0,
74
- "when": ["numel(shapes.X) == numel(shapes.Y)", "f16Ok(dtypes.T)"],
75
- "constants": { "vec4": false },
76
  "passes": [
77
  {
78
  "id": "main",
79
  "name": "Gelu.scalar",
80
  "shader": "elementwise-bias-gelu.wgsl.jinja",
81
  "bindings": [
82
- {
83
- "name": "x",
84
- "arg": "X",
85
- "semantic": "X",
86
- "buffer": { "type": "read-only-storage" },
87
- "elementType": "$scalar"
88
- },
89
- { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
90
- {
91
- "name": "params",
92
- "semantic": "kernel.params",
93
- "buffer": { "type": "uniform" },
94
- "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
95
- }
96
  ],
97
- "dispatch": { "threads": "numel(shapes.X)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
 
 
 
 
98
  }
99
  ]
100
  }
 
2
  "domain": "com.microsoft",
3
  "name": "Gelu",
4
  "sinceVersion": 1,
5
+ "inputs": { "X": { "dtype": "T" } },
6
+ "outputs": { "Y": { "dtype": "T", "rank": "ranks.X", "shape": "shapes.X" } },
 
 
 
 
 
 
 
 
 
 
 
7
  "typeConstraints": { "T": ["float32", "float16"] },
8
+ "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
9
+ "derive": { "scalar": "dtypes.T", "approximate": "\"erf\"", "vec4Tail": false, "hasBias": false },
10
+ "when": ["numel(shapes.X) == numel(shapes.Y)", "f16Ok(dtypes.T)"],
 
 
 
 
 
 
 
 
 
11
  "variants": [
12
  {
13
  "id": "vec4",
14
  "priority": 20,
15
+ "when": ["numel(shapes.X) > 0", "numel(shapes.X) % 4 == 0"],
16
+ "derive": { "vec4": true, "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
17
  "passes": [
18
  {
19
  "id": "main",
20
  "name": "Gelu.vec4",
21
  "shader": "elementwise-bias-gelu.wgsl.jinja",
22
  "bindings": [
23
+ { "arg": "X", "name": "x", "elementType": "$vectorScalar" },
24
+ { "arg": "Y", "name": "y", "elementType": "$vectorScalar" },
25
+ { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X) / 4" }] }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
  ],
27
+ "dispatch": {
28
+ "x": "min(ceilDiv((numel(shapes.X) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
29
+ "y": "ceilDiv(ceilDiv((numel(shapes.X) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
30
+ "z": 1
31
+ }
32
  }
33
  ]
34
  },
35
  {
36
  "id": "scalar",
37
  "priority": 0,
38
+ "derive": { "vec4": false },
 
39
  "passes": [
40
  {
41
  "id": "main",
42
  "name": "Gelu.scalar",
43
  "shader": "elementwise-bias-gelu.wgsl.jinja",
44
  "bindings": [
45
+ { "arg": "X", "name": "x", "elementType": "$scalar" },
46
+ { "arg": "Y", "name": "y", "elementType": "$scalar" },
47
+ { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
 
 
 
 
 
 
 
 
 
 
 
48
  ],
49
+ "dispatch": {
50
+ "x": "min(ceilDiv((numel(shapes.X)), (tunables.WORKGROUP_SIZE)), 65535)",
51
+ "y": "ceilDiv(ceilDiv((numel(shapes.X)), (tunables.WORKGROUP_SIZE)), 65535)",
52
+ "z": 1
53
+ }
54
  }
55
  ]
56
  }
build/webgpu/metadata.json CHANGED
@@ -1,18 +1,21 @@
1
  {
2
  "name": "com.microsoft.Gelu",
3
- "id": "_com_microsoft_gelu_webgpu_6511522",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "cFEvXUyEhaDZGPv9ppB2H/jtirbnrkEOd8ngzRycX+k=",
11
- "elementwise-bias-gelu.wgsl.jinja": "Eg8N2jJCMce+IsYNcCzuxvs8CSv7yTQXdorvv+c0L58=",
12
- "manifest.json": "89wktvMtLewo412huOWKAmpUsL1Q9k9RYkeWTxrhepE=",
13
- "test.json": "CmfJcfaIcytJIPtgyCKKLh72PzJvUlrpycr/FB0EZr8="
14
  }
15
  },
16
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
17
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.Gelu" }
 
 
 
18
  }
 
1
  {
2
  "name": "com.microsoft.Gelu",
3
+ "id": "_com_microsoft_gelu_webgpu_9075789",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "zb+BUPCMklZQithfY94RzWttjR/SZK0PIWgZjexuaxs=",
11
+ "elementwise-bias-gelu.wgsl.jinja": "K4LR9m7TibnO/vuJpCLk9+CTk1slr+E/XCGlmgH7hww=",
12
+ "manifest.json": "R8nYWq7dFuIme6gF1PUJWYKRv8u6u67dxzGzsZFLs38=",
13
+ "test.json": "K4VvkDT8J9wTZh+uYL3GSIAzvOw0PuV04RAIyhxmgxM="
14
  }
15
  },
16
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
17
+ "webgpu": {
18
+ "manifestSpec": "2.0",
19
+ "variants": { "vec4": ["elementwise-bias-gelu.wgsl.jinja"], "scalar": ["elementwise-bias-gelu.wgsl.jinja"] }
20
+ }
21
  }
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.Gelu",
3
  "cases": [
4
  {
5
  "name": "dispatch_cliff_scalar_f32",
@@ -105,7 +104,7 @@
105
  "provenance": {
106
  "source": "onnxruntime/test/contrib_ops/activation_op_test.cc",
107
  "test": "ActivationOpTest.Gelu",
108
- "notes": "Scalar-path companion for com.microsoft.Gelu subnormal linear-region behavior."
109
  },
110
  "inputs": {
111
  "X": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1e-40, -1e-40, 1e-38] } }
 
1
  {
 
2
  "cases": [
3
  {
4
  "name": "dispatch_cliff_scalar_f32",
 
104
  "provenance": {
105
  "source": "onnxruntime/test/contrib_ops/activation_op_test.cc",
106
  "test": "ActivationOpTest.Gelu",
107
+ "notes": "Subnormal inputs exercise the exact GELU linear region on the scalar path."
108
  },
109
  "inputs": {
110
  "X": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1e-40, -1e-40, 1e-38] } }