Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.BitShift

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

sync 6fdf6301e2bb

Browse files
Files changed (5) hide show
  1. README.md +2 -2
  2. build/webgpu/bitshift-vec4.wgsl.jinja +9 -6
  3. build/webgpu/bitshift.wgsl.jinja +20 -16
  4. build/webgpu/manifest.json +1 -1
  5. build/webgpu/metadata.json +6 -6
README.md CHANGED
@@ -48,14 +48,14 @@ Attributes and default values (overridable per request):
48
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
49
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
50
  - [`test.json`](build/webgpu/test.json) — correctness cases
51
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
52
  - [`bitshift-vec4.wgsl.jinja`](build/webgpu/bitshift-vec4.wgsl.jinja)
53
  - [`bitshift.wgsl.jinja`](build/webgpu/bitshift.wgsl.jinja)
54
 
55
  ## Use with `@huggingface/kernels`
56
 
57
  ```sh
58
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
59
  ```
60
 
61
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
48
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
49
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
50
  - [`test.json`](build/webgpu/test.json) — correctness cases
51
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
52
  - [`bitshift-vec4.wgsl.jinja`](build/webgpu/bitshift-vec4.wgsl.jinja)
53
  - [`bitshift.wgsl.jinja`](build/webgpu/bitshift.wgsl.jinja)
54
 
55
  ## Use with `@huggingface/kernels`
56
 
57
  ```sh
58
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
59
  ```
60
 
61
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/bitshift-vec4.wgsl.jinja CHANGED
@@ -1,3 +1,11 @@
 
 
 
 
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
 
3
  const BIT_WIDTH: u32 = {{ bitWidth }}u;
@@ -19,12 +27,7 @@ fn shift_one(in_value: u32, amount: u32) -> u32 {
19
 
20
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
21
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
22
- // 2D-folded flat index: gid.y carries the high bits past the
23
- // per-axis dispatch fold width.
24
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
25
- if (i >= params.count) {
26
- return;
27
- }
28
  let xv = x[i];
29
  let amount = shift[i];
30
  z[i] = vec4<u32>(
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
6
+ if ({{ name }} >= {{ bound }}) {
7
+ return;
8
+ }{% endmacro %}
9
  {{ env.wgsl.resourceDeclarations }}
10
 
11
  const BIT_WIDTH: u32 = {{ bitWidth }}u;
 
27
 
28
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
29
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
30
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
 
 
31
  let xv = x[i];
32
  let amount = shift[i];
33
  z[i] = vec4<u32>(
build/webgpu/bitshift.wgsl.jinja CHANGED
@@ -33,14 +33,18 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
33
  {% endfor %}
34
  return offset;
35
  {% endif %}
36
- }
37
- {%- endmacro %}{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
38
  {% set op_numel = namespace(value=1) %}
39
- {% for d in opShape %}{% set op_numel.value = op_numel.value * d %}{% endfor %}
 
 
40
  {% set out_numel = namespace(value=1) %}
41
- {% for d in outShape %}{% set out_numel.value = out_numel.value * d %}{% endfor %}
42
- {{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %})
43
- {%- endmacro %}{% macro broadcast_offset_fn(fn_name, opShape, opRank, outShape, outRank) %}
 
 
44
  {% set op_numel = namespace(value=1) %}
45
  {% for d in opShape %}
46
  {% set op_numel.value = op_numel.value * d %}
@@ -57,26 +61,26 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
57
  {% endif %}
58
  {% endfor %}
59
  {% endif %}
60
- {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}
61
- {%- endmacro %}
62
-
 
 
 
 
 
 
63
  {{ env.wgsl.resourceDeclarations }}
64
 
65
  const BIT_WIDTH: u32 = {{ bitWidth }}u;
66
 
67
-
68
  {{ broadcast_offset_fn("x_offset", xShape, xRank, zShape, zRank) }}
69
 
70
  {{ broadcast_offset_fn("y_offset", yShape, yRank, zShape, zRank) }}
71
 
72
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
73
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
74
- // 2D-folded flat index: gid.y carries the high bits past the
75
- // per-axis dispatch fold width.
76
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
77
- if (i >= params.count) {
78
- return;
79
- }
80
  let amount = shift[{{ broadcast_offset_call("y_offset", yShape, zShape, "i") }}];
81
  var value = 0u;
82
  if (amount < BIT_WIDTH) {
 
33
  {% endfor %}
34
  return offset;
35
  {% endif %}
36
+ }{% endmacro %}
37
+ {% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
38
  {% set op_numel = namespace(value=1) %}
39
+ {% for d in opShape %}
40
+ {% set op_numel.value = op_numel.value * d %}
41
+ {% endfor %}
42
  {% set out_numel = namespace(value=1) %}
43
+ {% for d in outShape %}
44
+ {% set out_numel.value = out_numel.value * d %}
45
+ {% endfor %}
46
+ {{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %}){% endmacro %}
47
+ {% macro broadcast_offset_fn(fn_name, opShape, opRank, outShape, outRank) %}
48
  {% set op_numel = namespace(value=1) %}
49
  {% for d in opShape %}
50
  {% set op_numel.value = op_numel.value * d %}
 
61
  {% endif %}
62
  {% endfor %}
63
  {% endif %}
64
+ {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}{% endmacro %}
65
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
66
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
67
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
68
+ // per-axis workgroup fold width.
69
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
70
+ if ({{ name }} >= {{ bound }}) {
71
+ return;
72
+ }{% endmacro %}
73
  {{ env.wgsl.resourceDeclarations }}
74
 
75
  const BIT_WIDTH: u32 = {{ bitWidth }}u;
76
 
 
77
  {{ broadcast_offset_fn("x_offset", xShape, xRank, zShape, zRank) }}
78
 
79
  {{ broadcast_offset_fn("y_offset", yShape, yRank, zShape, zRank) }}
80
 
81
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
82
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
83
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
 
 
84
  let amount = shift[{{ broadcast_offset_call("y_offset", yShape, zShape, "i") }}];
85
  var value = 0u;
86
  if (amount < BIT_WIDTH) {
build/webgpu/manifest.json CHANGED
@@ -24,7 +24,7 @@
24
  "passes": [
25
  {
26
  "id": "main",
27
- "name": "BitShift.vec4",
28
  "shader": "bitshift-vec4.wgsl.jinja",
29
  "bindings": [
30
  { "arg": "x", "elementType": "$vectorScalar" },
 
24
  "passes": [
25
  {
26
  "id": "main",
27
+ "name": "BitShift.Vec4",
28
  "shader": "bitshift-vec4.wgsl.jinja",
29
  "bindings": [
30
  { "arg": "x", "elementType": "$vectorScalar" },
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.BitShift",
3
- "id": "_ai_onnx_bitshift_webgpu_b738951",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,15 +8,15 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "pCMEZApV+9hKgXsgY8QFt30AzJB6jYpkK9vUoboxcv4=",
11
- "bitshift-vec4.wgsl.jinja": "t6oiDS+UC3rc8/U4pGJABtsCNpNI5Uh+EFEePIcoEUc=",
12
- "bitshift.wgsl.jinja": "w5jaDNLdgC/51OhsXtWNIdZTBf7mFnx/et01uoYsFQY=",
13
- "manifest.json": "I9WKns2iHR6NHc7PkEz4QzVO8xQibqM8ZZdedOSA4v0=",
14
  "test.json": "AaUcb5BUmYyvUXpso+6EIwZo3/LLieW55+9CV2pPXlg="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": { "same_shape_vec4": ["bitshift-vec4.wgsl.jinja"], "broadcast": ["bitshift.wgsl.jinja"] }
21
  }
22
  }
 
1
  {
2
  "name": "ai.onnx.BitShift",
3
+ "id": "_ai_onnx_bitshift_webgpu_9ba94ad",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "pCMEZApV+9hKgXsgY8QFt30AzJB6jYpkK9vUoboxcv4=",
11
+ "bitshift-vec4.wgsl.jinja": "3TWMsISu+BxfUW03yuKmF2qkvFpSvTXizQ1PnVLZVII=",
12
+ "bitshift.wgsl.jinja": "fiOkuARxBLyx/F2oMFr8AniXUyxZN8+zuI4qCXw6i3c=",
13
+ "manifest.json": "Zq+jV5XZGcYYOjd8UcqVORcbvcBEdx4sAEQskWqU+wY=",
14
  "test.json": "AaUcb5BUmYyvUXpso+6EIwZo3/LLieW55+9CV2pPXlg="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
18
  "webgpu": {
19
+ "manifestSpec": "2.1",
20
  "variants": { "same_shape_vec4": ["bitshift-vec4.wgsl.jinja"], "broadcast": ["bitshift.wgsl.jinja"] }
21
  }
22
  }