Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Asinh

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
e05b2e2
·
verified ·
1 Parent(s): 46198c5

sync 6fdf6301e2bb

Browse files
Files changed (5) hide show
  1. README.md +2 -2
  2. build/webgpu/manifest.json +10 -25
  3. build/webgpu/metadata.json +6 -6
  4. build/webgpu/unary-scalar.wgsl.jinja +7 -10
  5. build/webgpu/unary-vec4.wgsl.jinja +21 -12
README.md CHANGED
@@ -39,14 +39,14 @@ See the [ONNX `Asinh` spec](https://onnx.ai/onnx/operators/onnx__Asinh.html) for
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
43
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
44
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
  ```sh
49
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
50
  ```
51
 
52
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
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 cases
43
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
44
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
  ```sh
49
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
50
  ```
51
 
52
  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
@@ -11,20 +11,13 @@
11
  "id": "same_layout_vec4",
12
  "priority": 20,
13
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
14
- "derive": {
15
- "scalar": "dtypes.T",
16
- "usesF16": "dtypes.T == \"f16\"",
17
- "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
18
- },
19
  "passes": [
20
  {
21
  "id": "main",
22
- "name": "Asinh.vec4",
23
  "shader": "unary-vec4.wgsl.jinja",
24
- "derive": {
25
- "op": "\"asinh\"",
26
- "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
27
- },
28
  "bindings": ["x", "y", "params_unary"],
29
  "dispatch": {
30
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -37,14 +30,14 @@
37
  {
38
  "id": "elementwise",
39
  "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
40
- "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
41
  "passes": [
42
  {
43
  "id": "main",
44
  "name": "Asinh",
45
  "shader": "unary-scalar.wgsl.jinja",
46
- "derive": { "op": "\"asinh\"", "itemsPerInvocation": 4 },
47
- "bindings": ["x_2", "y_2", "params_2_unary"],
48
  "dispatch": {
49
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
50
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -57,17 +50,9 @@
57
  "bindings": {
58
  "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
59
  "y": { "buffer": "storage", "elementType": "$vectorScalar" },
60
- "params_unary": {
61
- "buffer": "uniform",
62
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }],
63
- "name": "params"
64
- },
65
- "x_2": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
66
- "y_2": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
67
- "params_2_unary": {
68
- "buffer": "uniform",
69
- "name": "params",
70
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
71
- }
72
  }
73
  }
 
11
  "id": "same_layout_vec4",
12
  "priority": 20,
13
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
14
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
15
  "passes": [
16
  {
17
  "id": "main",
18
+ "name": "Asinh.Vec4",
19
  "shader": "unary-vec4.wgsl.jinja",
20
+ "derive": { "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1" },
 
 
 
21
  "bindings": ["x", "y", "params_unary"],
22
  "dispatch": {
23
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
 
30
  {
31
  "id": "elementwise",
32
  "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
33
+ "derive": { "scalar": "dtypes.T" },
34
  "passes": [
35
  {
36
  "id": "main",
37
  "name": "Asinh",
38
  "shader": "unary-scalar.wgsl.jinja",
39
+ "derive": { "itemsPerInvocation": 4 },
40
+ "bindings": ["x_scalar", "y_scalar", "params_2_unary"],
41
  "dispatch": {
42
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
43
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
50
  "bindings": {
51
  "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
52
  "y": { "buffer": "storage", "elementType": "$vectorScalar" },
53
+ "params_unary": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }], "name": "params" },
54
+ "x_scalar": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
55
+ "y_scalar": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
56
+ "params_2_unary": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
 
 
 
 
 
 
 
 
57
  }
58
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Asinh",
3
- "id": "_ai_onnx_asinh_webgpu_dff8eb8",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,15 +8,15 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "c30/glNqEipJKMSjT7MbC+8yHpEbRFD6M5xiCrvJaWU=",
11
- "manifest.json": "JTmqVV3cYRmFuzOTJhE5GMB3oa5mRnwOs7fvi6wpmjY=",
12
  "test.json": "qc/5lhbmU9ThJz0q4RmA5mXQYtfW/afhjeBHHiID6E8=",
13
- "unary-scalar.wgsl.jinja": "VKx4VV3z+wLxGBj1zs5LG/9ZhpuBP811fSR/0WmNUiA=",
14
- "unary-vec4.wgsl.jinja": "IvuDkPiD9xynHcuqKVS7biZK1Dfw8NKIuwbNeHd2bzc="
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.Asinh",
3
+ "id": "_ai_onnx_asinh_webgpu_cf120ab",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "c30/glNqEipJKMSjT7MbC+8yHpEbRFD6M5xiCrvJaWU=",
11
+ "manifest.json": "FIKdgQzweTBuLBa1NT2RQiWssOUnS78h23AQ0gzExNY=",
12
  "test.json": "qc/5lhbmU9ThJz0q4RmA5mXQYtfW/afhjeBHHiID6E8=",
13
+ "unary-scalar.wgsl.jinja": "CwoO4oudpVKYUVGaDc7f6VVTzYuFSBbs0qdScwIBW3k=",
14
+ "unary-vec4.wgsl.jinja": "YQ5RcWsHgUfeS/cbkZPcs/84jm45KPTnEc0L7tqqBW4="
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,15 +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
  {% macro emit_asinh_safe() %}
25
  fn asinh_safe(x: f32) -> f32 {
26
  // The log(x + sqrt(x*x + 1)) identity overflows in the x*x term for very
@@ -33,10 +29,11 @@ fn asinh_safe(x: f32) -> f32 {
33
  return sign(x) * (log(a) + 0.6931471805599453);
34
  }
35
  return asinh(x);
36
- }{% endmacro %}
 
37
  {{ emit_asinh_safe() }}
38
  {{ flat_tail_open() }}
39
  let v = f32(x[i]);
40
  y[i] = {{ scalar }}(asinh_safe(v));
41
- {{ flat_tail_close() -}}
42
  }
 
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
  {% macro emit_asinh_safe() %}
21
  fn asinh_safe(x: f32) -> f32 {
22
  // The log(x + sqrt(x*x + 1)) identity overflows in the x*x term for very
 
29
  return sign(x) * (log(a) + 0.6931471805599453);
30
  }
31
  return asinh(x);
32
+ }
33
+ {% endmacro %}
34
  {{ emit_asinh_safe() }}
35
  {{ flat_tail_open() }}
36
  let v = f32(x[i]);
37
  y[i] = {{ scalar }}(asinh_safe(v));
38
+ {{ flat_tail_close() }}
39
  }
build/webgpu/unary-vec4.wgsl.jinja CHANGED
@@ -14,36 +14,45 @@ fn asinh_safe(x: f32) -> f32 {
14
  return sign(x) * (log(a) + 0.6931471805599453);
15
  }
16
  return asinh(x);
17
- }{% endmacro %}
 
18
  {{ emit_asinh_safe() }}
19
-
20
- {% set vec4PerThread = vec4PerThread %}
 
 
 
 
 
 
 
 
21
  {% if vec4PerThread > 1 %}
22
  const ITEMS: u32 = {{ vec4PerThread }}u;
23
- {% endif %}
24
 
 
25
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
26
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
27
- // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
28
- // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
29
  {% if vec4PerThread > 1 %}
30
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
31
  // access consecutive words on every step, while each lane can keep several
32
  // independent loads in flight.
33
- let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
34
  let span = (params.count + ITEMS - 1u) / ITEMS;
 
 
 
 
 
 
35
  for (var j = 0u; j < ITEMS; j = j + 1u) {
36
  let i = tid + j * span;
37
  if (i >= params.count) {
38
  break;
39
  }
40
  {% else %}
41
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
42
- if (i >= params.count) {
43
- return;
44
- }
45
  {% endif %}
46
-
47
  let xv = x[i];
48
  let fv = vec4<f32>(xv);
49
  y[i] = {{ vectorScalar }}(vec4<f32>(asinh_safe(fv.x), asinh_safe(fv.y), asinh_safe(fv.z), asinh_safe(fv.w)));
 
14
  return sign(x) * (log(a) + 0.6931471805599453);
15
  }
16
  return asinh(x);
17
+ }
18
+ {% endmacro %}
19
  {{ emit_asinh_safe() }}
20
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
21
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
22
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
23
+ // per-axis workgroup fold width.
24
+ {% if bound == "" %}
25
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
26
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
27
+ if ({{ name }} >= {{ bound }}) {
28
+ return;
29
+ }{% endif %}{% endmacro %}
30
  {% if vec4PerThread > 1 %}
31
  const ITEMS: u32 = {{ vec4PerThread }}u;
 
32
 
33
+ {% endif %}
34
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
35
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 
36
  {% if vec4PerThread > 1 %}
37
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
38
  // access consecutive words on every step, while each lane can keep several
39
  // independent loads in flight.
40
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }}
41
  let span = (params.count + ITEMS - 1u) / ITEMS;
42
+ // Lanes at or past span exist only because the dispatch rounds up to whole
43
+ // workgroups. Lane span + k would start on lane k's second group and rewrite
44
+ // up to ITEMS - 1 groups another lane already stored.
45
+ if (tid >= span) {
46
+ return;
47
+ }
48
  for (var j = 0u; j < ITEMS; j = j + 1u) {
49
  let i = tid + j * span;
50
  if (i >= params.count) {
51
  break;
52
  }
53
  {% else %}
54
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
55
  {% endif %}
 
56
  let xv = x[i];
57
  let fv = vec4<f32>(xv);
58
  y[i] = {{ vectorScalar }}(vec4<f32>(asinh_safe(fv.x), asinh_safe(fv.y), asinh_safe(fv.z), asinh_safe(fv.w)));