Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Elu

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
eeb8454
·
verified ·
1 Parent(s): c8dc025

sync 6fdf6301e2bb

Browse files
Files changed (5) hide show
  1. README.md +2 -2
  2. build/webgpu/manifest.json +26 -23
  3. build/webgpu/metadata.json +6 -6
  4. build/webgpu/unary-scalar.wgsl.jinja +6 -10
  5. build/webgpu/unary-vec4.wgsl.jinja +20 -11
README.md CHANGED
@@ -47,14 +47,14 @@ Default values (overridable per request):
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
51
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
52
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
53
 
54
  ## Use with `@huggingface/kernels`
55
 
56
  ```sh
57
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
58
  ```
59
 
60
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
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 cases
51
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
52
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
53
 
54
  ## Use with `@huggingface/kernels`
55
 
56
  ```sh
57
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
58
  ```
59
 
60
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/manifest.json CHANGED
@@ -7,28 +7,23 @@
7
  "attributes": { "alpha": { "default": 1 } },
8
  "typeConstraints": { "T": ["float32", "float16"] },
9
  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
10
- "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
11
  "variants": [
12
  {
13
  "id": "same_layout_vec4",
14
  "priority": 20,
15
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
16
- "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
17
  "passes": [
18
  {
19
  "id": "main",
20
- "name": "Elu.vec4",
21
  "shader": "unary-vec4.wgsl.jinja",
22
  "derive": {
23
- "op": "\"elu\"",
24
  "alpha": "attrs.alpha",
25
  "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
26
  },
27
- "bindings": [
28
- { "arg": "x", "elementType": "$vectorScalar" },
29
- { "arg": "y", "elementType": "$vectorScalar" },
30
- { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
31
- ],
32
  "dispatch": {
33
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
34
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -39,24 +34,15 @@
39
  },
40
  {
41
  "id": "elementwise",
42
- "when": ["numel(shapes.x) == numel(shapes.y) and (f16Ok(dtypes.T))"],
 
43
  "passes": [
44
  {
45
  "id": "main",
46
  "name": "Elu",
47
  "shader": "unary-scalar.wgsl.jinja",
48
- "derive": { "op": "\"elu\"", "itemsPerInvocation": 4 },
49
- "bindings": [
50
- "x",
51
- "y",
52
- {
53
- "name": "params",
54
- "struct": [
55
- { "name": "count", "type": "u32", "value": "numel(shapes.y)" },
56
- { "name": "alpha", "type": "f32", "value": "attrs.alpha" }
57
- ]
58
- }
59
- ],
60
  "dispatch": {
61
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
62
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -65,5 +51,22 @@
65
  }
66
  ]
67
  }
68
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
69
  }
 
7
  "attributes": { "alpha": { "default": 1 } },
8
  "typeConstraints": { "T": ["float32", "float16"] },
9
  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
10
+ "derive": { "scalar": "dtypes.T" },
11
  "variants": [
12
  {
13
  "id": "same_layout_vec4",
14
  "priority": 20,
15
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
16
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
17
  "passes": [
18
  {
19
  "id": "main",
20
+ "name": "Elu.Vec4",
21
  "shader": "unary-vec4.wgsl.jinja",
22
  "derive": {
 
23
  "alpha": "attrs.alpha",
24
  "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
25
  },
26
+ "bindings": ["x", "y", "params_unary_alpha"],
 
 
 
 
27
  "dispatch": {
28
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
29
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
 
34
  },
35
  {
36
  "id": "elementwise",
37
+ "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
38
+ "derive": { "scalar": "dtypes.T" },
39
  "passes": [
40
  {
41
  "id": "main",
42
  "name": "Elu",
43
  "shader": "unary-scalar.wgsl.jinja",
44
+ "derive": { "itemsPerInvocation": 4 },
45
+ "bindings": ["x_scalar", "y_scalar", "params_2_unary_alpha"],
 
 
 
 
 
 
 
 
 
 
46
  "dispatch": {
47
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
48
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
51
  }
52
  ]
53
  }
54
+ ],
55
+ "bindings": {
56
+ "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
57
+ "y": { "buffer": "storage", "elementType": "$vectorScalar" },
58
+ "params_unary_alpha": {
59
+ "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }],
60
+ "name": "params"
61
+ },
62
+ "x_scalar": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
63
+ "y_scalar": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
64
+ "params_2_unary_alpha": {
65
+ "name": "params",
66
+ "struct": [
67
+ { "name": "count", "type": "u32", "value": "numel(shapes.y)" },
68
+ { "name": "alpha", "type": "f32", "value": "attrs.alpha" }
69
+ ]
70
+ }
71
+ }
72
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Elu",
3
- "id": "_ai_onnx_elu_webgpu_ec52d7b",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,15 +8,15 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "eiwKWjWWOzA7C9z26++rXtUz5yy870QjGR4DlTbyUxw=",
11
- "manifest.json": "OzK0Gx6etOts1GvfOigK0vvpHTSzMjS+/QIH16Y2wmw=",
12
  "test.json": "y6pMkqiL7rwwNdPkRUdOQj3ngy0de2tjjzkXYQirH7k=",
13
- "unary-scalar.wgsl.jinja": "fBv4sZpdbUvk1gkgyMWwk82P4NXrnUMyuMwGhA4CXUg=",
14
- "unary-vec4.wgsl.jinja": "JXyoLX0+h/n4xm23tuhcE47PvUrjjUk03Qo+Vuldn/s="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
21
  }
22
  }
 
1
  {
2
  "name": "ai.onnx.Elu",
3
+ "id": "_ai_onnx_elu_webgpu_76244e0",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "eiwKWjWWOzA7C9z26++rXtUz5yy870QjGR4DlTbyUxw=",
11
+ "manifest.json": "IvTthsR90xiWhC28mfqv1a7RGdJpANLjYbVgaky/kCQ=",
12
  "test.json": "y6pMkqiL7rwwNdPkRUdOQj3ngy0de2tjjzkXYQirH7k=",
13
+ "unary-scalar.wgsl.jinja": "XK3EIfMTR1w/ZxJCG20OhgY95ucK9G9+Bh3N9NgTxYk=",
14
+ "unary-vec4.wgsl.jinja": "muTLNKxvjWmxPhxDEnVm3irq5SE6t5PoGNJ7mmMb7wU="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
18
  "webgpu": {
19
+ "manifestSpec": "2.1",
20
  "variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
21
  }
22
  }
build/webgpu/unary-scalar.wgsl.jinja CHANGED
@@ -12,16 +12,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
12
  for (var i = begin; i < end; i = i + 1u) {
13
  {%- endmacro %}
14
  {% macro flat_tail_close() %}
15
- }
16
- {% endmacro %}
17
-
18
- // Scalar unary elementwise implementation. Specialization emits only the
19
- // selected operation and any numerical helper it requires.
20
- {% if usesF16 %}
21
- enable f16;
22
- {% endif %}
23
  {{ env.wgsl.resourceDeclarations }}
24
- {% set alpha = alpha | default(0) %}
25
  fn expm1_safe(x: f32) -> f32 {
26
  // exp(x) - 1 without the catastrophic cancellation of the literal exp(x) - 1.0,
27
  // which loses all significance for small |x| (exp(x) rounds to 1.0). A 6-term
@@ -32,8 +27,9 @@ fn expm1_safe(x: f32) -> f32 {
32
  }
33
  return exp(x) - 1.0;
34
  }
 
35
  {{ flat_tail_open() }}
36
  let v = f32(x[i]);
37
  y[i] = {{ scalar }}(select(params.alpha * expm1_safe(v), v, v >= 0.0));
38
- {{ flat_tail_close() -}}
39
  }
 
12
  for (var i = begin; i < end; i = i + 1u) {
13
  {%- endmacro %}
14
  {% macro flat_tail_close() %}
15
+ }{% endmacro %}
16
+ // Scalar unary elementwise kernel: applies the selected operation, with the
17
+ // numerical helpers it uses, to each element.
 
 
 
 
 
18
  {{ env.wgsl.resourceDeclarations }}
19
+
20
  fn expm1_safe(x: f32) -> f32 {
21
  // exp(x) - 1 without the catastrophic cancellation of the literal exp(x) - 1.0,
22
  // which loses all significance for small |x| (exp(x) rounds to 1.0). A 6-term
 
27
  }
28
  return exp(x) - 1.0;
29
  }
30
+
31
  {{ flat_tail_open() }}
32
  let v = f32(x[i]);
33
  y[i] = {{ scalar }}(select(params.alpha * expm1_safe(v), v, v >= 0.0));
34
+ {{ flat_tail_close() }}
35
  }
build/webgpu/unary-vec4.wgsl.jinja CHANGED
@@ -2,7 +2,6 @@
2
  // component.
3
  {{ env.wgsl.resourceDeclarations }}
4
 
5
- {% set alpha = alpha | default(0) %}
6
  fn expm1_safe(x: vec4<f32>) -> vec4<f32> {
7
  // exp(x) - 1 without the catastrophic cancellation of the literal exp(x) - 1.0,
8
  // which loses all significance for small |x| (exp(x) rounds to 1.0). A 6-term
@@ -12,35 +11,45 @@ fn expm1_safe(x: vec4<f32>) -> vec4<f32> {
12
  let direct = exp(x) - 1.0;
13
  return select(direct, series, abs(x) < vec4<f32>(0.125));
14
  }
 
15
  const ELU_ALPHA: f32 = f32({{ alpha }});
16
 
17
- {% set vec4PerThread = vec4PerThread %}
 
 
 
 
 
 
 
 
 
18
  {% if vec4PerThread > 1 %}
19
  const ITEMS: u32 = {{ vec4PerThread }}u;
20
- {% endif %}
21
 
 
22
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
23
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
24
- // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
25
- // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
26
  {% if vec4PerThread > 1 %}
27
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
28
  // access consecutive words on every step, while each lane can keep several
29
  // independent loads in flight.
30
- let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
31
  let span = (params.count + ITEMS - 1u) / ITEMS;
 
 
 
 
 
 
32
  for (var j = 0u; j < ITEMS; j = j + 1u) {
33
  let i = tid + j * span;
34
  if (i >= params.count) {
35
  break;
36
  }
37
  {% else %}
38
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
39
- if (i >= params.count) {
40
- return;
41
- }
42
  {% endif %}
43
-
44
  let xv = x[i];
45
  let fv = vec4<f32>(xv);
46
  y[i] = {{ vectorScalar }}(select(ELU_ALPHA * expm1_safe(fv), fv, fv >= vec4<f32>(0.0)));
 
2
  // component.
3
  {{ env.wgsl.resourceDeclarations }}
4
 
 
5
  fn expm1_safe(x: vec4<f32>) -> vec4<f32> {
6
  // exp(x) - 1 without the catastrophic cancellation of the literal exp(x) - 1.0,
7
  // which loses all significance for small |x| (exp(x) rounds to 1.0). A 6-term
 
11
  let direct = exp(x) - 1.0;
12
  return select(direct, series, abs(x) < vec4<f32>(0.125));
13
  }
14
+
15
  const ELU_ALPHA: f32 = f32({{ alpha }});
16
 
17
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
18
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
19
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
20
+ // per-axis workgroup fold width.
21
+ {% if bound == "" %}
22
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
23
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
24
+ if ({{ name }} >= {{ bound }}) {
25
+ return;
26
+ }{% endif %}{% endmacro %}
27
  {% if vec4PerThread > 1 %}
28
  const ITEMS: u32 = {{ vec4PerThread }}u;
 
29
 
30
+ {% endif %}
31
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
32
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 
33
  {% if vec4PerThread > 1 %}
34
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
35
  // access consecutive words on every step, while each lane can keep several
36
  // independent loads in flight.
37
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }}
38
  let span = (params.count + ITEMS - 1u) / ITEMS;
39
+ // Lanes at or past span exist only because the dispatch rounds up to whole
40
+ // workgroups. Lane span + k would start on lane k's second group and rewrite
41
+ // up to ITEMS - 1 groups another lane already stored.
42
+ if (tid >= span) {
43
+ return;
44
+ }
45
  for (var j = 0u; j < ITEMS; j = j + 1u) {
46
  let i = tid + j * span;
47
  if (i >= params.count) {
48
  break;
49
  }
50
  {% else %}
51
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
52
  {% endif %}
 
53
  let xv = x[i];
54
  let fv = vec4<f32>(xv);
55
  y[i] = {{ vectorScalar }}(select(ELU_ALPHA * expm1_safe(fv), fv, fv >= vec4<f32>(0.0)));