Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.BitwiseOr

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

sync 6fdf6301e2bb

Browse files
Files changed (6) hide show
  1. README.md +2 -2
  2. build/webgpu/bitwise-binary-broadcast.wgsl.jinja +15 -16
  3. build/webgpu/bitwise-binary-vec4.wgsl.jinja +20 -10
  4. build/webgpu/manifest.json +11 -44
  5. build/webgpu/metadata.json +8 -9
  6. build/webgpu/test.json +2 -1
README.md CHANGED
@@ -40,14 +40,14 @@ See the [ONNX `BitwiseOr` spec](https://onnx.ai/onnx/operators/onnx__BitwiseOr.h
40
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
41
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
42
  - [`test.json`](build/webgpu/test.json) — correctness cases
43
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
44
  - [`bitwise-binary-broadcast.wgsl.jinja`](build/webgpu/bitwise-binary-broadcast.wgsl.jinja)
45
  - [`bitwise-binary-vec4.wgsl.jinja`](build/webgpu/bitwise-binary-vec4.wgsl.jinja)
46
 
47
  ## Use with `@huggingface/kernels`
48
 
49
  ```sh
50
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
51
  ```
52
 
53
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
40
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
41
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
42
  - [`test.json`](build/webgpu/test.json) — correctness cases
43
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
44
  - [`bitwise-binary-broadcast.wgsl.jinja`](build/webgpu/bitwise-binary-broadcast.wgsl.jinja)
45
  - [`bitwise-binary-vec4.wgsl.jinja`](build/webgpu/bitwise-binary-vec4.wgsl.jinja)
46
 
47
  ## Use with `@huggingface/kernels`
48
 
49
  ```sh
50
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
51
  ```
52
 
53
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/bitwise-binary-broadcast.wgsl.jinja CHANGED
@@ -12,9 +12,7 @@ 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
  {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
19
  fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
20
  {% if out_numel == 0 %}
@@ -50,14 +48,18 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
50
  {% endfor %}
51
  return offset;
52
  {% endif %}
53
- }
54
- {%- endmacro %}{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
55
  {% set op_numel = namespace(value=1) %}
56
- {% for d in opShape %}{% set op_numel.value = op_numel.value * d %}{% endfor %}
 
 
57
  {% set out_numel = namespace(value=1) %}
58
- {% for d in outShape %}{% set out_numel.value = out_numel.value * d %}{% endfor %}
59
- {{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %})
60
- {%- endmacro %}{% macro broadcast_offset_fn(fn_name, opShape, opRank, outShape, outRank) %}
 
 
61
  {% set op_numel = namespace(value=1) %}
62
  {% for d in opShape %}
63
  {% set op_numel.value = op_numel.value * d %}
@@ -74,8 +76,8 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
74
  {% endif %}
75
  {% endfor %}
76
  {% endif %}
77
- {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}
78
- {%- endmacro %}{% macro binary_broadcast_offsets() %}
79
  {% set aShape = aShape | default([]) %}
80
  {% set aRank = aRank | default(0) %}
81
  {% set bShape = bShape | default([]) %}
@@ -84,12 +86,9 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
84
  {% set cRank = cRank | default(0) %}
85
  {{ broadcast_offset_fn("a_offset", aShape, aRank, cShape, cRank) }}
86
 
87
- {{ broadcast_offset_fn("b_offset", bShape, bRank, cShape, cRank) }}
88
- {%- endmacro %}
89
-
90
  {{ env.wgsl.resourceDeclarations }}
91
 
92
-
93
  {{ binary_broadcast_offsets() }}
94
 
95
  {{ flat_tail_open() }}
@@ -98,5 +97,5 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
98
  value = value & 0xffu;
99
  {% endif %}
100
  c[i] = value;
101
- {{ flat_tail_close() -}}
102
  }
 
12
  for (var i = begin; i < end; i = i + 1u) {
13
  {%- endmacro %}
14
  {% macro flat_tail_close() %}
15
+ }{% endmacro %}
 
 
16
  {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
17
  fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
18
  {% if out_numel == 0 %}
 
48
  {% endfor %}
49
  return offset;
50
  {% endif %}
51
+ }{% endmacro %}
52
+ {% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
53
  {% set op_numel = namespace(value=1) %}
54
+ {% for d in opShape %}
55
+ {% set op_numel.value = op_numel.value * d %}
56
+ {% endfor %}
57
  {% set out_numel = namespace(value=1) %}
58
+ {% for d in outShape %}
59
+ {% set out_numel.value = out_numel.value * d %}
60
+ {% endfor %}
61
+ {{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %}){% endmacro %}
62
+ {% macro broadcast_offset_fn(fn_name, opShape, opRank, outShape, outRank) %}
63
  {% set op_numel = namespace(value=1) %}
64
  {% for d in opShape %}
65
  {% set op_numel.value = op_numel.value * d %}
 
76
  {% endif %}
77
  {% endfor %}
78
  {% endif %}
79
+ {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}{% endmacro %}
80
+ {% macro binary_broadcast_offsets() %}
81
  {% set aShape = aShape | default([]) %}
82
  {% set aRank = aRank | default(0) %}
83
  {% set bShape = bShape | default([]) %}
 
86
  {% set cRank = cRank | default(0) %}
87
  {{ broadcast_offset_fn("a_offset", aShape, aRank, cShape, cRank) }}
88
 
89
+ {{ broadcast_offset_fn("b_offset", bShape, bRank, cShape, cRank) }}{% endmacro %}
 
 
90
  {{ env.wgsl.resourceDeclarations }}
91
 
 
92
  {{ binary_broadcast_offsets() }}
93
 
94
  {{ flat_tail_open() }}
 
97
  value = value & 0xffu;
98
  {% endif %}
99
  c[i] = value;
100
+ {{ flat_tail_close() }}
101
  }
build/webgpu/bitwise-binary-vec4.wgsl.jinja CHANGED
@@ -4,33 +4,43 @@
4
  // (128-bit loads/stores). uint8 storage uses one u32 slot per element, so
5
  // the result is masked to the low byte per lane. Same semantics as the scalar
6
  // broadcast kernel when A, B, C share a shape.
7
- {% set vec4PerThread = vec4PerThread %}
 
 
 
 
 
 
 
 
 
 
8
  {% if vec4PerThread > 1 %}
9
  const ITEMS: u32 = {{ vec4PerThread }}u;
10
- {% endif %}
11
 
 
12
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
13
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
14
- // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
15
- // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
16
  {% if vec4PerThread > 1 %}
17
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
18
  // access consecutive words on every step, while each lane can keep several
19
  // independent loads in flight.
20
- let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
21
  let span = (params.count + ITEMS - 1u) / ITEMS;
 
 
 
 
 
 
22
  for (var j = 0u; j < ITEMS; j = j + 1u) {
23
  let i = tid + j * span;
24
  if (i >= params.count) {
25
  break;
26
  }
27
  {% else %}
28
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
29
- if (i >= params.count) {
30
- return;
31
- }
32
  {% endif %}
33
-
34
  var value = a[i] | b[i];
35
  {% if cDtype == "uint8" %}
36
  value = value & vec4<u32>(0xffu);
 
4
  // (128-bit loads/stores). uint8 storage uses one u32 slot per element, so
5
  // the result is masked to the low byte per lane. Same semantics as the scalar
6
  // broadcast kernel when A, B, C share a shape.
7
+
8
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
9
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
10
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
11
+ // per-axis workgroup fold width.
12
+ {% if bound == "" %}
13
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
14
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
15
+ if ({{ name }} >= {{ bound }}) {
16
+ return;
17
+ }{% endif %}{% endmacro %}
18
  {% if vec4PerThread > 1 %}
19
  const ITEMS: u32 = {{ vec4PerThread }}u;
 
20
 
21
+ {% endif %}
22
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
23
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 
24
  {% if vec4PerThread > 1 %}
25
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
26
  // access consecutive words on every step, while each lane can keep several
27
  // independent loads in flight.
28
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }}
29
  let span = (params.count + ITEMS - 1u) / ITEMS;
30
+ // Lanes at or past span exist only because the dispatch rounds up to whole
31
+ // workgroups. Lane span + k would start on lane k's second group and rewrite
32
+ // up to ITEMS - 1 groups another lane already stored.
33
+ if (tid >= span) {
34
+ return;
35
+ }
36
  for (var j = 0u; j < ITEMS; j = j + 1u) {
37
  let i = tid + j * span;
38
  if (i >= params.count) {
39
  break;
40
  }
41
  {% else %}
42
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
43
  {% endif %}
 
44
  var value = a[i] | b[i];
45
  {% if cDtype == "uint8" %}
46
  value = value & vec4<u32>(0xffu);
build/webgpu/manifest.json CHANGED
@@ -17,10 +17,9 @@
17
  "passes": [
18
  {
19
  "id": "main",
20
- "name": "BitwiseOr.vec4",
21
  "shader": "bitwise-binary-vec4.wgsl.jinja",
22
  "derive": {
23
- "op": "\"or\"",
24
  "cDtype": "tensorDtypes.c",
25
  "vec4PerThread": "4 if numel(shapes.c) * dtypeBytes(tensorDtypes.c) <= 16777216 else 1"
26
  },
@@ -36,7 +35,7 @@
36
  {
37
  "id": "broadcast",
38
  "when": ["tensorDtypes.a == tensorDtypes.b", "tensorDtypes.a == tensorDtypes.c", "ranks.a <= ranks.c", "ranks.b <= ranks.c"],
39
- "derive": { "bitwiseOp": "\"or\"", "logicalDtype": "tensorDtypes.c" },
40
  "passes": [
41
  {
42
  "id": "main",
@@ -51,35 +50,7 @@
51
  "cRank": "ranks.c",
52
  "itemsPerInvocation": 4
53
  },
54
- "bindings": ["a_2_bitwise", "b_2_bitwise", "c_2_bitwise", "params_2"],
55
- "dispatch": {
56
- "x": "min(ceilDiv((ceilDiv(numel(shapes.c), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
57
- "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.c), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
58
- "z": 1
59
- }
60
- }
61
- ]
62
- },
63
- {
64
- "id": "same_shape_scalar_x4",
65
- "priority": 15,
66
- "when": ["sameShape(shapes.a, shapes.c)", "sameShape(shapes.b, shapes.c)", "numel(shapes.c) > 0", "numel(shapes.c) % 4 != 0", "tensorDtypes.a == tensorDtypes.b", "tensorDtypes.a == tensorDtypes.c", "ranks.a <= ranks.c", "ranks.b <= ranks.c"],
67
- "derive": { "bitwiseOp": "\"or\"", "logicalDtype": "tensorDtypes.c" },
68
- "passes": [
69
- {
70
- "id": "main",
71
- "name": "BitwiseOr",
72
- "shader": "bitwise-binary-broadcast.wgsl.jinja",
73
- "derive": {
74
- "aShape": "shapes.a",
75
- "bShape": "shapes.b",
76
- "cShape": "shapes.c",
77
- "aRank": "ranks.a",
78
- "bRank": "ranks.b",
79
- "cRank": "ranks.c",
80
- "itemsPerInvocation": 4
81
- },
82
- "bindings": ["a_2_bitwise", "b_2_bitwise", "c_2_bitwise", "params_2"],
83
  "dispatch": {
84
  "x": "min(ceilDiv((ceilDiv(numel(shapes.c), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
85
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.c), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -90,17 +61,13 @@
90
  }
91
  ],
92
  "bindings": {
93
- "a": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
94
- "b": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
95
- "c_bitwise": { "buffer": "storage", "elementType": "$vectorScalar", "name": "c" },
96
- "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.c) / 4" }] },
97
- "a_2_bitwise": { "buffer": "read-only-storage", "name": "a", "elementType": "$T" },
98
- "b_2_bitwise": { "buffer": "read-only-storage", "name": "b", "elementType": "$T" },
99
- "c_2_bitwise": { "buffer": "storage", "name": "c", "elementType": "$T" },
100
- "params_2": {
101
- "buffer": "uniform",
102
- "name": "params",
103
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.c)" }]
104
- }
105
  }
106
  }
 
17
  "passes": [
18
  {
19
  "id": "main",
20
+ "name": "BitwiseOr.Vec4",
21
  "shader": "bitwise-binary-vec4.wgsl.jinja",
22
  "derive": {
 
23
  "cDtype": "tensorDtypes.c",
24
  "vec4PerThread": "4 if numel(shapes.c) * dtypeBytes(tensorDtypes.c) <= 16777216 else 1"
25
  },
 
35
  {
36
  "id": "broadcast",
37
  "when": ["tensorDtypes.a == tensorDtypes.b", "tensorDtypes.a == tensorDtypes.c", "ranks.a <= ranks.c", "ranks.b <= ranks.c"],
38
+ "derive": { "logicalDtype": "tensorDtypes.c" },
39
  "passes": [
40
  {
41
  "id": "main",
 
50
  "cRank": "ranks.c",
51
  "itemsPerInvocation": 4
52
  },
53
+ "bindings": ["a_2_bitwise", "b_2_bitwise", "c_2_bitwise", "params_count"],
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
54
  "dispatch": {
55
  "x": "min(ceilDiv((ceilDiv(numel(shapes.c), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
56
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.c), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
61
  }
62
  ],
63
  "bindings": {
64
+ "a": { "elementType": "$vectorScalar" },
65
+ "b": { "elementType": "$vectorScalar" },
66
+ "c_bitwise": { "elementType": "$vectorScalar", "name": "c" },
67
+ "params": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.c) / 4" }] },
68
+ "a_2_bitwise": { "name": "a", "elementType": "$T" },
69
+ "b_2_bitwise": { "name": "b", "elementType": "$T" },
70
+ "c_2_bitwise": { "name": "c", "elementType": "$T" },
71
+ "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.c)" }] }
 
 
 
 
72
  }
73
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.BitwiseOr",
3
- "id": "_ai_onnx_bitwiseor_webgpu_8e87398",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,19 +8,18 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "mERKODQEQmulEut7roCKqmrmyS+N1AQRVcIWVUTpnwc=",
11
- "bitwise-binary-broadcast.wgsl.jinja": "92rC9SLChfbxowLyL3FBnv8aF7q3Qz1MzTRuMJA0M9g=",
12
- "bitwise-binary-vec4.wgsl.jinja": "Ya480vMzE1FogZwUV0yNxmmoN8txJGwMw0hYks0P8V8=",
13
- "manifest.json": "g0N7NLtcwWg/wq+LSVjhGiCwg8rQcXh+St10paafCHI=",
14
- "test.json": "kEOiVM09gGxUF6F5gGIdczfCsGnnzBc0ZlPlcGFceHM="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": {
21
  "same_shape_vec4": ["bitwise-binary-vec4.wgsl.jinja"],
22
- "broadcast": ["bitwise-binary-broadcast.wgsl.jinja"],
23
- "same_shape_scalar_x4": ["bitwise-binary-broadcast.wgsl.jinja"]
24
  }
25
  }
26
  }
 
1
  {
2
  "name": "ai.onnx.BitwiseOr",
3
+ "id": "_ai_onnx_bitwiseor_webgpu_4aacf5f",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "mERKODQEQmulEut7roCKqmrmyS+N1AQRVcIWVUTpnwc=",
11
+ "bitwise-binary-broadcast.wgsl.jinja": "8Kpsu5+j9W4G5daZWs9Qt7p8btm35ZTZ+bCb3n6u4Fs=",
12
+ "bitwise-binary-vec4.wgsl.jinja": "mvuU+HilQUkhsL6C1zSIZNYXwhB008hau9kKUoaHKKc=",
13
+ "manifest.json": "adG6Jvfbrpp1JMSdDng3XWQuoqCBcxkWipnf73O0qto=",
14
+ "test.json": "fpTB2xj4oqysAwu5FTEzIpD5VBFBDkhdAg6W5jvQ7J0="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
18
  "webgpu": {
19
+ "manifestSpec": "2.1",
20
  "variants": {
21
  "same_shape_vec4": ["bitwise-binary-vec4.wgsl.jinja"],
22
+ "broadcast": ["bitwise-binary-broadcast.wgsl.jinja"]
 
23
  }
24
  }
25
  }
build/webgpu/test.json CHANGED
@@ -189,7 +189,8 @@
189
  "name": "onnx_backend_bitwise_or_i16_4d",
190
  "provenance": {
191
  "source": "cmake/external/onnx/onnx/backend/test/data/node/test_bitwise_or_i16_4d",
192
- "test": "test_bitwise_or_i16_4d"
 
193
  },
194
  "inputs": {
195
  "a": {
 
189
  "name": "onnx_backend_bitwise_or_i16_4d",
190
  "provenance": {
191
  "source": "cmake/external/onnx/onnx/backend/test/data/node/test_bitwise_or_i16_4d",
192
+ "test": "test_bitwise_or_i16_4d",
193
+ "notes": "The upstream test of this name stores int8 tensors despite the i16 in its name; the case keeps the upstream name and data."
194
  },
195
  "inputs": {
196
  "a": {