sync 91d990483a17
Browse files- README.md +13 -9
- build/webgpu/bench.json +0 -1
- build/webgpu/manifest.json +45 -70
- build/webgpu/metadata.json +14 -7
- build/webgpu/quick-gelu.wgsl.jinja +20 -22
- build/webgpu/test.json +1 -2
README.md
CHANGED
|
@@ -18,15 +18,15 @@ See the [ONNX Runtime `QuickGelu` contrib-operator spec](https://github.com/micr
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
-
| Name |
|
| 22 |
-
| --- | --- | --- | --- | --- | --- |
|
| 23 |
-
| `X` | `
|
| 24 |
|
| 25 |
## Outputs
|
| 26 |
|
| 27 |
-
| Name |
|
| 28 |
-
| --- | --- | --- | --- | --- | --- |
|
| 29 |
-
| `Y` | `
|
| 30 |
|
| 31 |
## Attributes
|
| 32 |
|
|
@@ -44,7 +44,7 @@ Default values (overridable per request):
|
|
| 44 |
|
| 45 |
## Files
|
| 46 |
|
| 47 |
-
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
|
| 48 |
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
|
| 49 |
- [`test.json`](build/webgpu/test.json) — correctness cases
|
| 50 |
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
|
|
@@ -52,10 +52,14 @@ Default values (overridable per request):
|
|
| 52 |
|
| 53 |
## Use with `@huggingface/kernels`
|
| 54 |
|
| 55 |
-
|
| 56 |
-
|
|
|
|
|
|
|
|
|
|
| 57 |
|
| 58 |
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
|
|
|
|
| 59 |
|
| 60 |
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
|
| 61 |
|
|
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
+
| Name | Logical dtype | Rank | Shape | Description | Presence |
|
| 22 |
+
| --- | --- | --- | --- | --- | --- |
|
| 23 |
+
| `X` | `T` | — | — | Input tensor of any shape. | 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; same shape as the input. | required |
|
| 30 |
|
| 31 |
## Attributes
|
| 32 |
|
|
|
|
| 44 |
|
| 45 |
## Files
|
| 46 |
|
| 47 |
+
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
|
| 48 |
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
|
| 49 |
- [`test.json`](build/webgpu/test.json) — correctness cases
|
| 50 |
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
|
|
|
|
| 52 |
|
| 53 |
## Use with `@huggingface/kernels`
|
| 54 |
|
| 55 |
+
```sh
|
| 56 |
+
npm install --save-exact @huggingface/kernels@0.0.1-preview.2
|
| 57 |
+
```
|
| 58 |
+
|
| 59 |
+
Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
|
| 60 |
|
| 61 |
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
|
| 62 |
+
It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
|
| 63 |
|
| 64 |
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
|
| 65 |
|
build/webgpu/bench.json
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
{
|
| 2 |
-
"op": "com.microsoft.QuickGelu",
|
| 3 |
"tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
|
| 4 |
"cases": [
|
| 5 |
{
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
|
| 3 |
"cases": [
|
| 4 |
{
|
build/webgpu/manifest.json
CHANGED
|
@@ -2,53 +2,25 @@
|
|
| 2 |
"domain": "com.microsoft",
|
| 3 |
"name": "QuickGelu",
|
| 4 |
"sinceVersion": 1,
|
| 5 |
-
"
|
| 6 |
-
"
|
| 7 |
-
"
|
| 8 |
-
{
|
| 9 |
-
"role": "Y",
|
| 10 |
-
"dtype": "T",
|
| 11 |
-
"rank": "ranks.X",
|
| 12 |
-
"shape": "shapes.X",
|
| 13 |
-
"description": "Output tensor; same shape as the input."
|
| 14 |
-
}
|
| 15 |
-
],
|
| 16 |
-
"attributes": { "alpha": 1.702 },
|
| 17 |
-
"attributeDescriptions": {
|
| 18 |
-
"alpha": "Scalar multiplier applied to `x` inside the sigmoid; defaults to 1.702, which approximates GELU."
|
| 19 |
-
},
|
| 20 |
"typeConstraints": { "T": ["float32", "float16"] },
|
| 21 |
-
"
|
| 22 |
-
"X": { "kind": "tensor", "semantic": "X", "role": "input" },
|
| 23 |
-
"Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
|
| 24 |
-
},
|
| 25 |
-
"tunables": { "WORKGROUP_SIZE": 256 },
|
| 26 |
"derive": {
|
| 27 |
"deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
|
| 28 |
"workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
|
| 29 |
"baseOk": "workgroupOk and numel(shapes.X) == numel(shapes.Y) and f16Ok(dtypes.T)",
|
| 30 |
"vec4Ok": "numel(shapes.X) > 0 and numel(shapes.X) % 4 == 0"
|
| 31 |
},
|
| 32 |
-
"
|
| 33 |
-
"scalarTail": [
|
| 34 |
-
{ "name": "x", "arg": "X", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
|
| 35 |
-
{ "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
|
| 36 |
-
{
|
| 37 |
-
"name": "params",
|
| 38 |
-
"semantic": "kernel.params",
|
| 39 |
-
"buffer": { "type": "uniform" },
|
| 40 |
-
"struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
|
| 41 |
-
}
|
| 42 |
-
]
|
| 43 |
-
},
|
| 44 |
"variants": [
|
| 45 |
{
|
| 46 |
"id": "vec4",
|
| 47 |
"priority": 20,
|
| 48 |
-
"when": ["
|
| 49 |
-
"
|
| 50 |
"scalar": "dtypes.T",
|
| 51 |
-
"usesF16": "dtypes.T == \"f16\"",
|
| 52 |
"vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
|
| 53 |
"vec4": true,
|
| 54 |
"vec4Tail": false
|
|
@@ -57,63 +29,66 @@
|
|
| 57 |
{
|
| 58 |
"id": "main",
|
| 59 |
"name": "QuickGelu.vec4",
|
| 60 |
-
"
|
|
|
|
| 61 |
"bindings": [
|
| 62 |
-
{
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
"semantic": "X",
|
| 66 |
-
"buffer": { "type": "read-only-storage" },
|
| 67 |
-
"elementType": "$vectorScalar"
|
| 68 |
-
},
|
| 69 |
-
{
|
| 70 |
-
"name": "y",
|
| 71 |
-
"arg": "Y",
|
| 72 |
-
"semantic": "Y",
|
| 73 |
-
"buffer": { "type": "storage" },
|
| 74 |
-
"elementType": "$vectorScalar"
|
| 75 |
-
},
|
| 76 |
-
{
|
| 77 |
-
"name": "params",
|
| 78 |
-
"semantic": "kernel.params",
|
| 79 |
-
"buffer": { "type": "uniform" },
|
| 80 |
-
"struct": {
|
| 81 |
-
"name": "Params",
|
| 82 |
-
"fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X) / 4" }]
|
| 83 |
-
}
|
| 84 |
-
}
|
| 85 |
],
|
| 86 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 87 |
}
|
| 88 |
]
|
| 89 |
},
|
| 90 |
{
|
| 91 |
"id": "vec4_tail",
|
| 92 |
"priority": 10,
|
| 93 |
-
"when": ["
|
| 94 |
-
"
|
| 95 |
"passes": [
|
| 96 |
{
|
| 97 |
"id": "main",
|
| 98 |
"name": "QuickGelu.vec4Tail",
|
| 99 |
-
"
|
| 100 |
-
"
|
| 101 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 102 |
}
|
| 103 |
]
|
| 104 |
},
|
| 105 |
{
|
| 106 |
"id": "scalar",
|
| 107 |
"priority": 0,
|
| 108 |
-
"when": ["
|
| 109 |
-
"
|
| 110 |
"passes": [
|
| 111 |
{
|
| 112 |
"id": "main",
|
| 113 |
"name": "QuickGelu.scalar",
|
| 114 |
-
"
|
| 115 |
-
"
|
| 116 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 117 |
}
|
| 118 |
]
|
| 119 |
}
|
|
|
|
| 2 |
"domain": "com.microsoft",
|
| 3 |
"name": "QuickGelu",
|
| 4 |
"sinceVersion": 1,
|
| 5 |
+
"inputs": { "X": { "dtype": "T" } },
|
| 6 |
+
"outputs": { "Y": { "dtype": "T", "rank": "ranks.X", "shape": "shapes.X" } },
|
| 7 |
+
"attributes": { "alpha": { "default": 1.702 } },
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
"typeConstraints": { "T": ["float32", "float16"] },
|
| 9 |
+
"tunables": { "WORKGROUP_SIZE": { "default": 256 } },
|
|
|
|
|
|
|
|
|
|
|
|
|
| 10 |
"derive": {
|
| 11 |
"deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
|
| 12 |
"workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
|
| 13 |
"baseOk": "workgroupOk and numel(shapes.X) == numel(shapes.Y) and f16Ok(dtypes.T)",
|
| 14 |
"vec4Ok": "numel(shapes.X) > 0 and numel(shapes.X) % 4 == 0"
|
| 15 |
},
|
| 16 |
+
"when": ["baseOk"],
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 17 |
"variants": [
|
| 18 |
{
|
| 19 |
"id": "vec4",
|
| 20 |
"priority": 20,
|
| 21 |
+
"when": ["vec4Ok"],
|
| 22 |
+
"derive": {
|
| 23 |
"scalar": "dtypes.T",
|
|
|
|
| 24 |
"vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
|
| 25 |
"vec4": true,
|
| 26 |
"vec4Tail": false
|
|
|
|
| 29 |
{
|
| 30 |
"id": "main",
|
| 31 |
"name": "QuickGelu.vec4",
|
| 32 |
+
"shader": "quick-gelu.wgsl.jinja",
|
| 33 |
+
"derive": { "alpha": "attrs.alpha" },
|
| 34 |
"bindings": [
|
| 35 |
+
{ "arg": "X", "name": "x", "elementType": "$vectorScalar" },
|
| 36 |
+
{ "arg": "Y", "name": "y", "elementType": "$vectorScalar" },
|
| 37 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X) / 4" }] }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 38 |
],
|
| 39 |
+
"dispatch": {
|
| 40 |
+
"x": "min(ceilDiv((numel(shapes.X) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 41 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.X) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 42 |
+
"z": 1
|
| 43 |
+
}
|
| 44 |
}
|
| 45 |
]
|
| 46 |
},
|
| 47 |
{
|
| 48 |
"id": "vec4_tail",
|
| 49 |
"priority": 10,
|
| 50 |
+
"when": ["numel(shapes.X) > 0"],
|
| 51 |
+
"derive": { "scalar": "dtypes.T", "vec4": false, "vec4Tail": true },
|
| 52 |
"passes": [
|
| 53 |
{
|
| 54 |
"id": "main",
|
| 55 |
"name": "QuickGelu.vec4Tail",
|
| 56 |
+
"shader": "quick-gelu.wgsl.jinja",
|
| 57 |
+
"derive": { "alpha": "attrs.alpha" },
|
| 58 |
+
"bindings": [
|
| 59 |
+
{ "arg": "X", "name": "x", "elementType": "$scalar" },
|
| 60 |
+
{ "arg": "Y", "name": "y", "elementType": "$scalar" },
|
| 61 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
|
| 62 |
+
],
|
| 63 |
+
"dispatch": {
|
| 64 |
+
"x": "min(ceilDiv((ceilDiv(numel(shapes.X), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 65 |
+
"y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.X), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 66 |
+
"z": 1
|
| 67 |
+
}
|
| 68 |
}
|
| 69 |
]
|
| 70 |
},
|
| 71 |
{
|
| 72 |
"id": "scalar",
|
| 73 |
"priority": 0,
|
| 74 |
+
"when": ["true"],
|
| 75 |
+
"derive": { "scalar": "dtypes.T", "vec4": false, "vec4Tail": false },
|
| 76 |
"passes": [
|
| 77 |
{
|
| 78 |
"id": "main",
|
| 79 |
"name": "QuickGelu.scalar",
|
| 80 |
+
"shader": "quick-gelu.wgsl.jinja",
|
| 81 |
+
"derive": { "alpha": "attrs.alpha" },
|
| 82 |
+
"bindings": [
|
| 83 |
+
{ "arg": "X", "name": "x", "elementType": "$scalar" },
|
| 84 |
+
{ "arg": "Y", "name": "y", "elementType": "$scalar" },
|
| 85 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
|
| 86 |
+
],
|
| 87 |
+
"dispatch": {
|
| 88 |
+
"x": "min(ceilDiv((numel(shapes.X)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 89 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.X)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 90 |
+
"z": 1
|
| 91 |
+
}
|
| 92 |
}
|
| 93 |
]
|
| 94 |
}
|
build/webgpu/metadata.json
CHANGED
|
@@ -1,18 +1,25 @@
|
|
| 1 |
{
|
| 2 |
"name": "com.microsoft.QuickGelu",
|
| 3 |
-
"id": "
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
-
"bench.json": "
|
| 11 |
-
"manifest.json": "
|
| 12 |
-
"quick-gelu.wgsl.jinja": "
|
| 13 |
-
"test.json": "
|
| 14 |
}
|
| 15 |
},
|
| 16 |
-
"provenance": { "kernel": { "sha": "
|
| 17 |
-
"webgpu": {
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
}
|
|
|
|
| 1 |
{
|
| 2 |
"name": "com.microsoft.QuickGelu",
|
| 3 |
+
"id": "_com_microsoft_quickgelu_webgpu_57d870f",
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
+
"bench.json": "UkzjTMKrQBdgOKbmqRYDd53jtc70n2SBbpJvjok4i60=",
|
| 11 |
+
"manifest.json": "BwEkPwdzXku28VYZ+dMC7eKDL4HzAlj6H2Iwx61Rmw8=",
|
| 12 |
+
"quick-gelu.wgsl.jinja": "/+KTUHzwuZwWV4Q2df6YEdgv2b5wu6EBdI1eZUCUj10=",
|
| 13 |
+
"test.json": "yMUdwnbXzdJruLruAgEJ6bcftDu0hE1av8ovE2UsXQs="
|
| 14 |
}
|
| 15 |
},
|
| 16 |
+
"provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
|
| 17 |
+
"webgpu": {
|
| 18 |
+
"manifestSpec": "2.0",
|
| 19 |
+
"variants": {
|
| 20 |
+
"vec4": ["quick-gelu.wgsl.jinja"],
|
| 21 |
+
"vec4_tail": ["quick-gelu.wgsl.jinja"],
|
| 22 |
+
"scalar": ["quick-gelu.wgsl.jinja"]
|
| 23 |
+
}
|
| 24 |
+
}
|
| 25 |
}
|
build/webgpu/quick-gelu.wgsl.jinja
CHANGED
|
@@ -1,50 +1,48 @@
|
|
| 1 |
{% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
|
| 2 |
{% if note == "dispatch-limit" %}
|
| 3 |
-
// 2D-folded flat index: gid.y carries the high bits past the
|
| 4 |
-
//
|
| 5 |
{% elif note == "limit" %}
|
| 6 |
-
// 2D-folded flat index: gid.y carries the high bits past the
|
| 7 |
-
//
|
| 8 |
{% elif note == "device-axis" %}
|
| 9 |
-
// The flat dispatch is folded across x/y at
|
| 10 |
-
//
|
| 11 |
{% elif note == "vec4-limit" %}
|
| 12 |
-
// 2D-folded flat vec4 index: gid.y carries the high bits past the
|
| 13 |
-
//
|
| 14 |
{% elif note == "element-limit" %}
|
| 15 |
// 2D-folded flat element index: gid.y carries the high bits past the
|
| 16 |
-
//
|
| 17 |
{% elif note == "dispatch" %}
|
| 18 |
-
// 2D-folded flat index: gid.y carries the high bits past the
|
| 19 |
-
//
|
| 20 |
{% endif %}
|
| 21 |
{% if bound == "" %}
|
| 22 |
-
let {{ name }} = gid.x + gid.y *
|
| 23 |
{%- elif guardInline %}
|
| 24 |
-
let {{ name }} = gid.x + gid.y *
|
| 25 |
if ({{ name }} >= {{ bound }}) { return; }
|
| 26 |
{%- else %}
|
| 27 |
-
let {{ name }} = gid.x + gid.y *
|
| 28 |
if ({{ name }} >= {{ bound }}) {
|
| 29 |
return;
|
| 30 |
}
|
| 31 |
{%- endif %}
|
| 32 |
{% endmacro %}
|
| 33 |
|
| 34 |
-
{% if usesF16 %}
|
| 35 |
-
enable f16;
|
| 36 |
-
{% endif %}
|
| 37 |
{{ env.wgsl.resourceDeclarations }}
|
| 38 |
|
| 39 |
// com.microsoft.QuickGelu : Y = X * sigmoid(alpha * X)
|
| 40 |
// sigmoid here uses the stable two-branch form so the gate never overflows
|
| 41 |
// for extreme magnitudes (alpha*x = +/-1702 for x = -/+1000 with the default
|
| 42 |
-
// alpha): the naive 1/(1+exp(-alpha*x)) computes exp(+1702)
|
| 43 |
-
//
|
| 44 |
-
//
|
|
|
|
| 45 |
// as a constant; alpha == 0 collapses to sigmoid(0) = 0.5 exactly (the
|
| 46 |
// z >= 0 branch: 1/(1+exp(0))).
|
| 47 |
-
const ALPHA: f32 = f32({{
|
| 48 |
|
| 49 |
fn sigmoid_stable(z: f32) -> f32 {
|
| 50 |
if (z >= 0.0) {
|
|
@@ -59,7 +57,7 @@ fn quick_gelu(v: f32) -> f32 {
|
|
| 59 |
}
|
| 60 |
|
| 61 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 62 |
-
fn main(@builtin(global_invocation_id) gid: vec3<u32>
|
| 63 |
{{ flat_index_2d() }}
|
| 64 |
{% if vec4Tail %}
|
| 65 |
let base = i * 4u;
|
|
|
|
| 1 |
{% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
|
| 2 |
{% if note == "dispatch-limit" %}
|
| 3 |
+
// 2D-folded flat index: gid.y carries the high bits past the dispatch's
|
| 4 |
+
// per-axis workgroup fold width (outputs > 16.7M elements).
|
| 5 |
{% elif note == "limit" %}
|
| 6 |
+
// 2D-folded flat index: gid.y carries the high bits past the dispatch's
|
| 7 |
+
// per-axis workgroup fold width.
|
| 8 |
{% elif note == "device-axis" %}
|
| 9 |
+
// The flat dispatch is folded across x/y at a fixed per-axis workgroup
|
| 10 |
+
// width; gid.y carries the high portion of the output index.
|
| 11 |
{% elif note == "vec4-limit" %}
|
| 12 |
+
// 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
|
| 13 |
+
// per-axis workgroup fold width (the dispatch caps x and spills into y).
|
| 14 |
{% elif note == "element-limit" %}
|
| 15 |
// 2D-folded flat element index: gid.y carries the high bits past the
|
| 16 |
+
// dispatch's per-axis workgroup fold width.
|
| 17 |
{% elif note == "dispatch" %}
|
| 18 |
+
// 2D-folded flat index: gid.y carries the high bits past the dispatch's
|
| 19 |
+
// per-axis workgroup fold width.
|
| 20 |
{% endif %}
|
| 21 |
{% if bound == "" %}
|
| 22 |
+
let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 23 |
{%- elif guardInline %}
|
| 24 |
+
let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 25 |
if ({{ name }} >= {{ bound }}) { return; }
|
| 26 |
{%- else %}
|
| 27 |
+
let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 28 |
if ({{ name }} >= {{ bound }}) {
|
| 29 |
return;
|
| 30 |
}
|
| 31 |
{%- endif %}
|
| 32 |
{% endmacro %}
|
| 33 |
|
|
|
|
|
|
|
|
|
|
| 34 |
{{ env.wgsl.resourceDeclarations }}
|
| 35 |
|
| 36 |
// com.microsoft.QuickGelu : Y = X * sigmoid(alpha * X)
|
| 37 |
// sigmoid here uses the stable two-branch form so the gate never overflows
|
| 38 |
// for extreme magnitudes (alpha*x = +/-1702 for x = -/+1000 with the default
|
| 39 |
+
// alpha): the naive 1/(1+exp(-alpha*x)) computes exp(+1702), overflowing
|
| 40 |
+
// its intermediate even though the sigmoid result is finite. The forms
|
| 41 |
+
// exp(z)/(1+exp(z)) (z <= 0) and 1/(1+exp(-z)) (z >= 0) each only evaluate
|
| 42 |
+
// exp of a non-positive argument. ALPHA is compiled
|
| 43 |
// as a constant; alpha == 0 collapses to sigmoid(0) = 0.5 exactly (the
|
| 44 |
// z >= 0 branch: 1/(1+exp(0))).
|
| 45 |
+
const ALPHA: f32 = f32({{ alpha }});
|
| 46 |
|
| 47 |
fn sigmoid_stable(z: f32) -> f32 {
|
| 48 |
if (z >= 0.0) {
|
|
|
|
| 57 |
}
|
| 58 |
|
| 59 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 60 |
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
| 61 |
{{ flat_index_2d() }}
|
| 62 |
{% if vec4Tail %}
|
| 63 |
let base = i * 4u;
|
build/webgpu/test.json
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
{
|
| 2 |
-
"op": "com.microsoft.QuickGelu",
|
| 3 |
"cases": [
|
| 4 |
{
|
| 5 |
"name": "dispatch_cliff_scalar_over_16M",
|
|
@@ -205,7 +204,7 @@
|
|
| 205 |
"provenance": {
|
| 206 |
"source": "onnxruntime/test/contrib_ops/activation_op_test.cc",
|
| 207 |
"test": "ActivationOpTest.QuickGelu",
|
| 208 |
-
"notes": "
|
| 209 |
},
|
| 210 |
"attrs": { "alpha": 0 },
|
| 211 |
"inputs": {
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"cases": [
|
| 3 |
{
|
| 4 |
"name": "dispatch_cliff_scalar_over_16M",
|
|
|
|
| 204 |
"provenance": {
|
| 205 |
"source": "onnxruntime/test/contrib_ops/activation_op_test.cc",
|
| 206 |
"test": "ActivationOpTest.QuickGelu",
|
| 207 |
+
"notes": "With `alpha = 0`, subnormal inputs exercise exact half-input behavior on the scalar path."
|
| 208 |
},
|
| 209 |
"attrs": { "alpha": 0 },
|
| 210 |
"inputs": {
|