Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.CastLike

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
37cd06e
·
verified ·
1 Parent(s): 1056d1c

sync 6fdf6301e2bb

Browse files
Files changed (6) hide show
  1. README.md +2 -2
  2. build/webgpu/manifest.json +22 -42
  3. build/webgpu/metadata.json +7 -7
  4. build/webgpu/test.json +1 -1
  5. build/webgpu/unary-scalar.wgsl.jinja +5 -14
  6. build/webgpu/unary-vec4.wgsl.jinja +19 -14
README.md CHANGED
@@ -50,7 +50,7 @@ Default values (overridable per request):
50
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
51
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
52
  - [`test.json`](build/webgpu/test.json) — correctness cases
53
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
54
  - [`cast-scalar-x4.wgsl.jinja`](build/webgpu/cast-scalar-x4.wgsl.jinja)
55
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
56
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
@@ -58,7 +58,7 @@ Default values (overridable per request):
58
  ## Use with `@huggingface/kernels`
59
 
60
  ```sh
61
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
62
  ```
63
 
64
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
50
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
51
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
52
  - [`test.json`](build/webgpu/test.json) — correctness cases
53
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
54
  - [`cast-scalar-x4.wgsl.jinja`](build/webgpu/cast-scalar-x4.wgsl.jinja)
55
  - [`unary-scalar.wgsl.jinja`](build/webgpu/unary-scalar.wgsl.jinja)
56
  - [`unary-vec4.wgsl.jinja`](build/webgpu/unary-vec4.wgsl.jinja)
 
58
  ## Use with `@huggingface/kernels`
59
 
60
  ```sh
61
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
62
  ```
63
 
64
  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
@@ -12,25 +12,21 @@
12
  },
13
  "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "VEC4_TAIL_MIN_ELEMENTS": { "default": 4096 } },
14
  "derive": {
15
- "usesF16": "tensorDtypes.x == \"float16\" or tensorDtypes.y == \"float16\"",
16
  "wrapNarrowInt": "(tensorDtypes.x == \"float32\" or tensorDtypes.x == \"float16\") and (tensorDtypes.y == \"int8\" or tensorDtypes.y == \"uint8\")",
17
- "wrapSigned": "tensorDtypes.y == \"int8\""
 
 
 
 
18
  },
19
  "when": ["numel(shapes.x) == numel(shapes.y)", "ranks.y == ranks.x", "f16Ok(dtypes.T1) and f16Ok(dtypes.T2)"],
20
  "bindings": {
21
- "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
22
- "y": { "buffer": "storage", "elementType": "$outVec" },
23
- "params": {
24
- "buffer": "uniform",
25
- "struct": [{ "name": "count", "type": "u32", "value": "floor(numel(shapes.y) / 4)" }]
26
- },
27
- "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$inScalar" },
28
- "y_2": { "name": "y", "buffer": "storage", "elementType": "$outScalar" },
29
- "params_2": {
30
- "name": "params",
31
- "buffer": "uniform",
32
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
33
- }
34
  },
35
  "variants": [
36
  {
@@ -38,19 +34,14 @@
38
  "priority": 20,
39
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "tensorDtypes.y != \"bool\""],
40
  "derive": {
41
- "scalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))",
42
- "vectorScalar": "\"vec4<\" ~ (\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))) ~ \">\"",
43
- "outVec": "\"vec4<\" ~ (\"f16\" if tensorDtypes.y == \"float16\" else (\"f32\" if tensorDtypes.y == \"float32\" else (\"i32\" if tensorDtypes.y == \"int32\" or tensorDtypes.y == \"int8\" else \"u32\"))) ~ \">\""
44
  },
45
  "passes": [
46
  {
47
  "id": "main",
48
- "name": "CastLike.vec4",
49
  "shader": "unary-vec4.wgsl.jinja",
50
- "derive": {
51
- "op": "\"castlike\"",
52
- "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
53
- },
54
  "bindings": ["x", "y", "params"],
55
  "dispatch": {
56
  "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -65,21 +56,14 @@
65
  "priority": 10,
66
  "when": ["numel(shapes.x) >= tunables.VEC4_TAIL_MIN_ELEMENTS", "numel(shapes.x) % 4 != 0", "tensorDtypes.y != \"bool\""],
67
  "derive": {
68
- "inScalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))",
69
- "outScalar": "\"f16\" if tensorDtypes.y == \"float16\" else (\"f32\" if tensorDtypes.y == \"float32\" else (\"i32\" if tensorDtypes.y == \"int32\" or tensorDtypes.y == \"int8\" else \"u32\"))",
70
- "scalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))",
71
- "vectorScalar": "\"vec4<\" ~ (\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))) ~ \">\"",
72
- "outVec": "\"vec4<\" ~ (\"f16\" if tensorDtypes.y == \"float16\" else (\"f32\" if tensorDtypes.y == \"float32\" else (\"i32\" if tensorDtypes.y == \"int32\" or tensorDtypes.y == \"int8\" else \"u32\"))) ~ \">\""
73
  },
74
  "passes": [
75
  {
76
  "id": "bulk",
77
- "name": "CastLike.vec4Bulk",
78
  "shader": "unary-vec4.wgsl.jinja",
79
- "derive": {
80
- "op": "\"castlike\"",
81
- "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
82
- },
83
  "bindings": ["x", "y", "params"],
84
  "dispatch": {
85
  "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -89,10 +73,10 @@
89
  },
90
  {
91
  "id": "tail",
92
- "name": "CastLike.scalarTail",
93
  "shader": "cast-scalar-x4.wgsl.jinja",
94
  "derive": { "tailOnly": true },
95
- "bindings": ["x_2", "y_2", "params_2"],
96
  "dispatch": {
97
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
98
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -103,18 +87,14 @@
103
  },
104
  {
105
  "id": "elementwise",
106
- "derive": {
107
- "inScalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))",
108
- "outScalar": "\"f16\" if tensorDtypes.y == \"float16\" else (\"f32\" if tensorDtypes.y == \"float32\" else (\"i32\" if tensorDtypes.y == \"int32\" or tensorDtypes.y == \"int8\" else \"u32\"))",
109
- "toBool": "tensorDtypes.y == \"bool\""
110
- },
111
  "passes": [
112
  {
113
  "id": "main",
114
  "name": "CastLike",
115
  "shader": "unary-scalar.wgsl.jinja",
116
- "derive": { "op": "\"castlike\"", "itemsPerInvocation": 4 },
117
- "bindings": ["x_2", "y_2", "params_2"],
118
  "dispatch": {
119
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
120
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
12
  },
13
  "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "VEC4_TAIL_MIN_ELEMENTS": { "default": 4096 } },
14
  "derive": {
 
15
  "wrapNarrowInt": "(tensorDtypes.x == \"float32\" or tensorDtypes.x == \"float16\") and (tensorDtypes.y == \"int8\" or tensorDtypes.y == \"uint8\")",
16
+ "wrapSigned": "tensorDtypes.y == \"int8\"",
17
+ "vectorScalar": "\"vec4<\" ~ (\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))) ~ \">\"",
18
+ "outVec": "\"vec4<\" ~ (\"f16\" if tensorDtypes.y == \"float16\" else (\"f32\" if tensorDtypes.y == \"float32\" else (\"i32\" if tensorDtypes.y == \"int32\" or tensorDtypes.y == \"int8\" else \"u32\"))) ~ \">\"",
19
+ "inScalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))",
20
+ "outScalar": "\"f16\" if tensorDtypes.y == \"float16\" else (\"f32\" if tensorDtypes.y == \"float32\" else (\"i32\" if tensorDtypes.y == \"int32\" or tensorDtypes.y == \"int8\" else \"u32\"))"
21
  },
22
  "when": ["numel(shapes.x) == numel(shapes.y)", "ranks.y == ranks.x", "f16Ok(dtypes.T1) and f16Ok(dtypes.T2)"],
23
  "bindings": {
24
+ "x": { "elementType": "$vectorScalar" },
25
+ "y": { "elementType": "$outVec" },
26
+ "params": { "struct": [{ "name": "count", "type": "u32", "value": "floor(numel(shapes.y) / 4)" }] },
27
+ "x_tail": { "name": "x", "elementType": "$inScalar" },
28
+ "y_tail": { "name": "y", "elementType": "$outScalar" },
29
+ "params_tail": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
 
 
 
 
 
 
 
30
  },
31
  "variants": [
32
  {
 
34
  "priority": 20,
35
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "tensorDtypes.y != \"bool\""],
36
  "derive": {
37
+ "scalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))"
 
 
38
  },
39
  "passes": [
40
  {
41
  "id": "main",
42
+ "name": "CastLike.Vec4",
43
  "shader": "unary-vec4.wgsl.jinja",
44
+ "derive": { "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1" },
 
 
 
45
  "bindings": ["x", "y", "params"],
46
  "dispatch": {
47
  "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
 
56
  "priority": 10,
57
  "when": ["numel(shapes.x) >= tunables.VEC4_TAIL_MIN_ELEMENTS", "numel(shapes.x) % 4 != 0", "tensorDtypes.y != \"bool\""],
58
  "derive": {
59
+ "scalar": "\"f16\" if tensorDtypes.x == \"float16\" else (\"f32\" if tensorDtypes.x == \"float32\" else (\"i32\" if tensorDtypes.x == \"int32\" or tensorDtypes.x == \"int8\" else \"u32\"))"
 
 
 
 
60
  },
61
  "passes": [
62
  {
63
  "id": "bulk",
64
+ "name": "CastLike.Vec4Bulk",
65
  "shader": "unary-vec4.wgsl.jinja",
66
+ "derive": { "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1" },
 
 
 
67
  "bindings": ["x", "y", "params"],
68
  "dispatch": {
69
  "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
 
73
  },
74
  {
75
  "id": "tail",
76
+ "name": "CastLike.ScalarTail",
77
  "shader": "cast-scalar-x4.wgsl.jinja",
78
  "derive": { "tailOnly": true },
79
+ "bindings": ["x_tail", "y_tail", "params_tail"],
80
  "dispatch": {
81
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
82
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
 
87
  },
88
  {
89
  "id": "elementwise",
90
+ "derive": { "toBool": "tensorDtypes.y == \"bool\"" },
 
 
 
 
91
  "passes": [
92
  {
93
  "id": "main",
94
  "name": "CastLike",
95
  "shader": "unary-scalar.wgsl.jinja",
96
+ "derive": { "itemsPerInvocation": 4 },
97
+ "bindings": ["x_tail", "y_tail", "params_tail"],
98
  "dispatch": {
99
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
100
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.CastLike",
3
- "id": "_ai_onnx_castlike_webgpu_b3b5a5c",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -9,15 +9,15 @@
9
  "files": {
10
  "bench.json": "HvPktrO4Hl0z7C35f4x48CIdfhgr5BXEPqDjvZWDBbw=",
11
  "cast-scalar-x4.wgsl.jinja": "ZEmWwTu9NiMzEKa+PuUlxatkSfJ8WLLYD78x0ay8l/U=",
12
- "manifest.json": "omTaO2Iw4AaBsEmMGrnaTWCUpxDFQ1mQQtk2ivwV7Yo=",
13
- "test.json": "uXxmLeBAgMyAvKw77XQjlYD+MzcMu9YPDObTu9xM9CI=",
14
- "unary-scalar.wgsl.jinja": "zBUAwYvX5TCH+YwAb8iuo9hhR7g51kBVeFGOMjxn8oo=",
15
- "unary-vec4.wgsl.jinja": "vndf17SxewEvqPOmYVKn4LQTBbVuBNjQ3cDs+olqyTo="
16
  }
17
  },
18
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
19
  "webgpu": {
20
- "manifestSpec": "2.0",
21
  "variants": {
22
  "same_layout_vec4": ["unary-vec4.wgsl.jinja"],
23
  "vec4_tail": ["cast-scalar-x4.wgsl.jinja", "unary-vec4.wgsl.jinja"],
 
1
  {
2
  "name": "ai.onnx.CastLike",
3
+ "id": "_ai_onnx_castlike_webgpu_2586498",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
9
  "files": {
10
  "bench.json": "HvPktrO4Hl0z7C35f4x48CIdfhgr5BXEPqDjvZWDBbw=",
11
  "cast-scalar-x4.wgsl.jinja": "ZEmWwTu9NiMzEKa+PuUlxatkSfJ8WLLYD78x0ay8l/U=",
12
+ "manifest.json": "qHRXKFzEsOnCP5gUdwhNxzoJ0MhJJtUCycKt1D1sUE8=",
13
+ "test.json": "1sAn8cqmTJT30RUXHrVDIB/E18sbrVmqljzFmHBQfxE=",
14
+ "unary-scalar.wgsl.jinja": "5wku4obJp+j904VWGZHGzYen00Ns4I3YCZ4DgEuYD/Y=",
15
+ "unary-vec4.wgsl.jinja": "EtHp0Kb/QUFl7cAuTzxm+hWe7aOk7h1w++zLhD1ort8="
16
  }
17
  },
18
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
19
  "webgpu": {
20
+ "manifestSpec": "2.1",
21
  "variants": {
22
  "same_layout_vec4": ["unary-vec4.wgsl.jinja"],
23
  "vec4_tail": ["cast-scalar-x4.wgsl.jinja", "unary-vec4.wgsl.jinja"],
build/webgpu/test.json CHANGED
@@ -265,7 +265,7 @@
265
  },
266
  "outputs": { "y": { "dtype": "int8", "shape": [4099], "tolerance": 0 } },
267
  "provenance": {
268
- "notes": "CastLike to an int8 target above the vec4-tail crossover (4099 = 4096 packed elements plus a 3-element scalar tail), selecting the vec4 bulk plus scalar-tail path with its signed int8 storage conversion. Every value is within int8's representable range, so the fixture asserts ONNX-defined truncation toward zero without depending on undefined out-of-range float-to-fixed-point behavior."
269
  }
270
  },
271
  {
 
265
  },
266
  "outputs": { "y": { "dtype": "int8", "shape": [4099], "tolerance": 0 } },
267
  "provenance": {
268
+ "notes": "A 4099-element float32-to-int8 CastLike (target dtype int8) is not four-aligned, leaving a 3-element tail. All values are within int8 range, verifying truncation toward zero rather than undefined out-of-range behavior."
269
  }
270
  },
271
  {
build/webgpu/unary-scalar.wgsl.jinja CHANGED
@@ -22,20 +22,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
22
  for (var i = begin; i < end; i = i + 1u) {
23
  {%- endmacro %}
24
  {% macro flat_tail_close() %}
25
- }
26
- {% endmacro %}
27
-
28
- {% set toBool = toBool is defined and toBool %}
29
- {% set wrapNarrowInt = wrapNarrowInt is defined and wrapNarrowInt %}
30
- {% set wrapSigned = wrapSigned is defined and wrapSigned %}
31
- {% set inScalar = inScalar | default(scalar) %}
32
- {% set outScalar = outScalar | default(scalar) %}
33
- // Scalar unary elementwise implementation. Specialization emits only the
34
- // selected operation and any numerical helper it requires.
35
- {% if usesF16 %}
36
- enable f16;
37
- {% endif %}
38
  {{ env.wgsl.resourceDeclarations }}
 
39
  {{ flat_tail_open() }}
40
  {% if toBool %}
41
  {% if inScalar == "f16" or inScalar == "f32" %}
@@ -54,5 +45,5 @@ enable f16;
54
  {% else %}
55
  y[i] = {{ outScalar }}(x[i]);
56
  {% endif %}
57
- {{ flat_tail_close() -}}
58
  }
 
22
  for (var i = begin; i < end; i = i + 1u) {
23
  {%- endmacro %}
24
  {% macro flat_tail_close() %}
25
+ }{% endmacro %}
26
+ // Scalar unary elementwise kernel: applies the selected operation, with the
27
+ // numerical helpers it uses, to each element.
 
 
 
 
 
 
 
 
 
 
28
  {{ env.wgsl.resourceDeclarations }}
29
+
30
  {{ flat_tail_open() }}
31
  {% if toBool %}
32
  {% if inScalar == "f16" or inScalar == "f32" %}
 
45
  {% else %}
46
  y[i] = {{ outScalar }}(x[i]);
47
  {% endif %}
48
+ {{ flat_tail_close() }}
49
  }
build/webgpu/unary-vec4.wgsl.jinja CHANGED
@@ -1,38 +1,43 @@
1
- {% set wrapNarrowInt = wrapNarrowInt is defined and wrapNarrowInt %}
2
- {% set wrapSigned = wrapSigned is defined and wrapSigned %}
3
- {% set outVec = outVec | default(vectorScalar) %}
4
  // Loads and stores vec4<T> while evaluating the selected unary operation per
5
  // component.
6
  {{ env.wgsl.resourceDeclarations }}
7
 
8
-
9
- {% set vec4PerThread = vec4PerThread %}
 
 
 
 
 
 
 
 
10
  {% if vec4PerThread > 1 %}
11
  const ITEMS: u32 = {{ vec4PerThread }}u;
12
- {% endif %}
13
 
 
14
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
15
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
16
- // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
17
- // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
18
  {% if vec4PerThread > 1 %}
19
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
20
  // access consecutive words on every step, while each lane can keep several
21
  // independent loads in flight.
22
- let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
23
  let span = (params.count + ITEMS - 1u) / ITEMS;
 
 
 
 
 
 
24
  for (var j = 0u; j < ITEMS; j = j + 1u) {
25
  let i = tid + j * span;
26
  if (i >= params.count) {
27
  break;
28
  }
29
  {% else %}
30
- let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
31
- if (i >= params.count) {
32
- return;
33
- }
34
  {% endif %}
35
-
36
  let xv = x[i];
37
  {% if wrapNarrowInt %}
38
  // Float-to-int8/uint8 Cast and CastLike wrap modulo 256 rather than saturate:
 
 
 
 
1
  // Loads and stores vec4<T> while evaluating the selected unary operation per
2
  // component.
3
  {{ env.wgsl.resourceDeclarations }}
4
 
5
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
6
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
7
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
8
+ // per-axis workgroup fold width.
9
+ {% if bound == "" %}
10
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
11
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
12
+ if ({{ name }} >= {{ bound }}) {
13
+ return;
14
+ }{% endif %}{% endmacro %}
15
  {% if vec4PerThread > 1 %}
16
  const ITEMS: u32 = {{ vec4PerThread }}u;
 
17
 
18
+ {% endif %}
19
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
20
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 
21
  {% if vec4PerThread > 1 %}
22
  // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
23
  // access consecutive words on every step, while each lane can keep several
24
  // independent loads in flight.
25
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }}
26
  let span = (params.count + ITEMS - 1u) / ITEMS;
27
+ // Lanes at or past span exist only because the dispatch rounds up to whole
28
+ // workgroups. Lane span + k would start on lane k's second group and rewrite
29
+ // up to ITEMS - 1 groups another lane already stored.
30
+ if (tid >= span) {
31
+ return;
32
+ }
33
  for (var j = 0u; j < ITEMS; j = j + 1u) {
34
  let i = tid + j * span;
35
  if (i >= params.count) {
36
  break;
37
  }
38
  {% else %}
39
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
 
 
 
40
  {% endif %}
 
41
  let xv = x[i];
42
  {% if wrapNarrowInt %}
43
  // Float-to-int8/uint8 Cast and CastLike wrap modulo 256 rather than saturate: