Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Abs

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
56326e7
·
verified ·
1 Parent(s): 11ade4e

sync 6fdf6301e2bb

Browse files
Files changed (5) hide show
  1. README.md +3 -3
  2. build/webgpu/manifest.json +37 -52
  3. build/webgpu/metadata.json +8 -8
  4. build/webgpu/unary-scalar.wgsl.jinja +8 -10
  5. build/webgpu/unary-vec4.wgsl.jinja +23 -12
README.md CHANGED
@@ -39,14 +39,14 @@ See the [ONNX `Abs` spec](https://onnx.ai/onnx/operators/onnx__Abs.html) for the
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.
@@ -60,5 +60,5 @@ Replace each `*Data` placeholder with a typed array containing the corresponding
60
  import { getKernel } from "@huggingface/kernels";
61
 
62
  const kernel = await getKernel("webgpu-kernels/ai.onnx.Abs", { version: 1 });
63
- const { y } = await kernel({ x: { data: xData, shape: [] } });
64
  ```
 
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.
 
60
  import { getKernel } from "@huggingface/kernels";
61
 
62
  const kernel = await getKernel("webgpu-kernels/ai.onnx.Abs", { version: 1 });
63
+ const { y } = await kernel({ x: { data: xData, shape: [3] } });
64
  ```
build/webgpu/manifest.json CHANGED
@@ -6,27 +6,20 @@
6
  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
7
  "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int8", "uint8"] },
8
  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
9
- "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "isInt8": "tensorDtypes.x == \"int8\"" },
10
  "variants": [
11
  {
12
  "id": "same_layout_vec4",
13
  "priority": 20,
14
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
15
- "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
16
  "passes": [
17
  {
18
  "id": "main",
19
- "name": "Abs.vec4",
20
  "shader": "unary-vec4.wgsl.jinja",
21
- "derive": {
22
- "op": "\"abs\"",
23
- "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
24
- },
25
- "bindings": [
26
- { "arg": "x", "elementType": "$vectorScalar" },
27
- { "arg": "y", "elementType": "$vectorScalar" },
28
- { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }] }
29
- ],
30
  "dispatch": {
31
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
32
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -35,37 +28,51 @@
35
  }
36
  ]
37
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
  {
39
  "id": "same_layout_vec4_tail",
40
  "priority": 19,
41
- "when": ["numel(shapes.x) > 4", "numel(shapes.x) % 4 != 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
42
- "derive": {
43
- "scalar": "dtypes.T",
44
- "usesF16": "dtypes.T == \"f16\"",
45
- "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
46
- },
47
  "passes": [
48
  {
49
  "id": "bulk",
50
- "name": "Abs.vec4Bulk",
51
  "shader": "unary-vec4.wgsl.jinja",
52
  "derive": {
53
- "op": "\"abs\"",
54
- "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
55
  },
56
  "bindings": ["x", "y", "params_unary_tail"],
57
  "dispatch": {
58
- "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
59
- "y": "ceilDiv(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
60
  "z": 1
61
  }
62
  },
63
  {
64
  "id": "tail",
65
- "name": "Abs.tail",
66
  "shader": "unary-scalar.wgsl.jinja",
67
- "derive": { "op": "\"abs\"", "itemsPerInvocation": 4, "tailOnly": true },
68
- "bindings": ["x_2", "y_2", "params_2_unary_tail"],
69
  "dispatch": {
70
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
71
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -73,42 +80,20 @@
73
  }
74
  }
75
  ]
76
- },
77
- {
78
- "id": "elementwise",
79
- "when": ["numel(shapes.x) == numel(shapes.y) and (f16Ok(dtypes.T))"],
80
- "passes": [
81
- {
82
- "id": "main",
83
- "name": "Abs",
84
- "shader": "unary-scalar.wgsl.jinja",
85
- "derive": { "op": "\"abs\"", "itemsPerInvocation": 4 },
86
- "bindings": [
87
- { "arg": "x" },
88
- { "arg": "y" },
89
- { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
90
- ],
91
- "dispatch": {
92
- "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
93
- "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
94
- "z": 1
95
- }
96
- }
97
- ]
98
  }
99
  ],
100
  "bindings": {
101
  "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
102
  "y": { "buffer": "storage", "elementType": "$vectorScalar" },
 
 
 
 
103
  "params_unary_tail": {
104
- "buffer": "uniform",
105
  "struct": [{ "name": "count", "type": "u32", "value": "floor(numel(shapes.y) / 4)" }],
106
  "name": "params"
107
  },
108
- "x_2": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
109
- "y_2": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
110
  "params_2_unary_tail": {
111
- "buffer": "uniform",
112
  "name": "params",
113
  "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
114
  }
 
6
  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
7
  "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int8", "uint8"] },
8
  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
9
+ "derive": { "scalar": "dtypes.T", "isInt8": "tensorDtypes.x == \"int8\"" },
10
  "variants": [
11
  {
12
  "id": "same_layout_vec4",
13
  "priority": 20,
14
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
15
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
16
  "passes": [
17
  {
18
  "id": "main",
19
+ "name": "Abs.Vec4",
20
  "shader": "unary-vec4.wgsl.jinja",
21
+ "derive": { "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1" },
22
+ "bindings": ["x", "y", "params_unary"],
 
 
 
 
 
 
 
23
  "dispatch": {
24
  "x": "min(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
25
  "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
 
28
  }
29
  ]
30
  },
31
+ {
32
+ "id": "elementwise",
33
+ "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
34
+ "derive": { "scalar": "dtypes.T" },
35
+ "passes": [
36
+ {
37
+ "id": "main",
38
+ "name": "Abs",
39
+ "shader": "unary-scalar.wgsl.jinja",
40
+ "derive": { "itemsPerInvocation": 4 },
41
+ "bindings": ["x_scalar", "y_scalar", "params_2_unary"],
42
+ "dispatch": {
43
+ "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
44
+ "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
45
+ "z": 1
46
+ }
47
+ }
48
+ ]
49
+ },
50
  {
51
  "id": "same_layout_vec4_tail",
52
  "priority": 19,
53
+ "when": ["numel(shapes.x) > 4", "numel(shapes.x) % 4 != 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)", "ceilDiv(numel(shapes.y), 4) >= ceilDiv(tunables.MIN_TAIL_ELEMENTS | default(0), 4)"],
54
+ "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
55
  "passes": [
56
  {
57
  "id": "bulk",
58
+ "name": "Abs.Vec4Bulk",
59
  "shader": "unary-vec4.wgsl.jinja",
60
  "derive": {
61
+ "vec4PerThread": "tunables.TAIL_VEC4_PER_THREAD | default(4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)"
 
62
  },
63
  "bindings": ["x", "y", "params_unary_tail"],
64
  "dispatch": {
65
+ "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), vec4PerThread)), (tunables.WORKGROUP_SIZE)), 65535)",
66
+ "y": "ceilDiv(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), vec4PerThread)), (tunables.WORKGROUP_SIZE)), 65535)",
67
  "z": 1
68
  }
69
  },
70
  {
71
  "id": "tail",
72
+ "name": "Abs.Tail",
73
  "shader": "unary-scalar.wgsl.jinja",
74
+ "derive": { "itemsPerInvocation": 4, "tailOnly": true },
75
+ "bindings": ["x_scalar", "y_scalar", "params_2_unary_tail"],
76
  "dispatch": {
77
  "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
78
  "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
 
80
  }
81
  }
82
  ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
  }
84
  ],
85
  "bindings": {
86
  "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
87
  "y": { "buffer": "storage", "elementType": "$vectorScalar" },
88
+ "params_unary": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }], "name": "params" },
89
+ "x_scalar": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
90
+ "y_scalar": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
91
+ "params_2_unary": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] },
92
  "params_unary_tail": {
 
93
  "struct": [{ "name": "count", "type": "u32", "value": "floor(numel(shapes.y) / 4)" }],
94
  "name": "params"
95
  },
 
 
96
  "params_2_unary_tail": {
 
97
  "name": "params",
98
  "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
99
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Abs",
3
- "id": "_ai_onnx_abs_webgpu_d58a189",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,19 +8,19 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "Ih1NAfwzMaCcTx87u2cQkpumJrkFlhSFUWTLMJKRlv8=",
11
- "manifest.json": "ydWYMT92F2Ji3aEFK746rb5M+PZGZA5Z3KrvtTnMVb8=",
12
  "test.json": "c5cMh22zfZDCfTTdEBgUb79TeUUQ7YS/oIvSzjP78J8=",
13
- "unary-scalar.wgsl.jinja": "6GB3YhDxyYXLECp5KVTt8/fYEvXP7/M5CkKuLQJATgc=",
14
- "unary-vec4.wgsl.jinja": "SNF1L8DMgZZ81K7jiPxDX62hNrEexAvUq5Lmt0uUoS8="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": {
21
  "same_layout_vec4": ["unary-vec4.wgsl.jinja"],
22
- "same_layout_vec4_tail": ["unary-scalar.wgsl.jinja", "unary-vec4.wgsl.jinja"],
23
- "elementwise": ["unary-scalar.wgsl.jinja"]
24
  }
25
  }
26
  }
 
1
  {
2
  "name": "ai.onnx.Abs",
3
+ "id": "_ai_onnx_abs_webgpu_84f9b8f",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "Ih1NAfwzMaCcTx87u2cQkpumJrkFlhSFUWTLMJKRlv8=",
11
+ "manifest.json": "uwBUephmnQNSAjs0kgmJ4LCddzZ3lS9qagmOV/9IUcQ=",
12
  "test.json": "c5cMh22zfZDCfTTdEBgUb79TeUUQ7YS/oIvSzjP78J8=",
13
+ "unary-scalar.wgsl.jinja": "IeupqMA2qhmyjdRGcpHQHvttij5VtYsdcmHFYv6X/Wo=",
14
+ "unary-vec4.wgsl.jinja": "vyG2isFT7EWrvfkPK90own2TFW7xS/qedhPxr9amAzI="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
18
  "webgpu": {
19
+ "manifestSpec": "2.1",
20
  "variants": {
21
  "same_layout_vec4": ["unary-vec4.wgsl.jinja"],
22
+ "elementwise": ["unary-scalar.wgsl.jinja"],
23
+ "same_layout_vec4_tail": ["unary-scalar.wgsl.jinja", "unary-vec4.wgsl.jinja"]
24
  }
25
  }
26
  }
build/webgpu/unary-scalar.wgsl.jinja CHANGED
@@ -22,16 +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 isInt8 = isInt8 is defined and isInt8 %}
29
- // Scalar unary elementwise implementation. Specialization emits only the
30
- // selected operation and any numerical helper it requires.
31
- {% if usesF16 %}
32
- enable f16;
33
- {% endif %}
34
  {{ env.wgsl.resourceDeclarations }}
 
35
  {{ flat_tail_open() }}
36
  {% if scalar == "i32" %}
37
  let value = x[i];
@@ -45,8 +40,11 @@ enable f16;
45
  {% endif %}
46
  {% elif scalar == "u32" %}
47
  y[i] = x[i];
 
 
 
48
  {% else %}
49
  y[i] = {{ scalar }}(abs(f32(x[i])));
50
  {% endif %}
51
- {{ flat_tail_close() -}}
52
  }
 
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 scalar == "i32" %}
32
  let value = x[i];
 
40
  {% endif %}
41
  {% elif scalar == "u32" %}
42
  y[i] = x[i];
43
+ {% elif scalar == "f32" %}
44
+ // Bitwise magnitude avoids flushing subnormal inputs through f32 arithmetic.
45
+ y[i] = bitcast<f32>(bitcast<u32>(x[i]) & 0x7fffffffu);
46
  {% else %}
47
  y[i] = {{ scalar }}(abs(f32(x[i])));
48
  {% endif %}
49
+ {{ flat_tail_close() }}
50
  }
build/webgpu/unary-vec4.wgsl.jinja CHANGED
@@ -1,36 +1,43 @@
1
- {% set isInt8 = isInt8 is defined and isInt8 %}
2
  // Loads and stores vec4<T> while evaluating the selected unary operation per
3
  // component.
4
  {{ env.wgsl.resourceDeclarations }}
5
 
6
-
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
  let xv = x[i];
35
  {% if scalar == "i32" %}
36
  let absv = select(xv, -xv, xv < vec4<i32>(0i));
@@ -43,6 +50,10 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
43
  {% endif %}
44
  {% elif scalar == "u32" %}
45
  y[i] = xv;
 
 
 
 
46
  {% else %}
47
  y[i] = {{ vectorScalar }}(abs(vec4<f32>(xv)));
48
  {% endif %}
 
 
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 scalar == "i32" %}
43
  let absv = select(xv, -xv, xv < vec4<i32>(0i));
 
50
  {% endif %}
51
  {% elif scalar == "u32" %}
52
  y[i] = xv;
53
+ {% elif scalar == "f32" %}
54
+ // Clearing the sign bit preserves subnormal magnitudes on GPUs that flush
55
+ // f32 arithmetic to zero; it also implements IEEE abs for NaNs and infinities.
56
+ y[i] = bitcast<vec4<f32>>(bitcast<vec4<u32>>(xv) & vec4<u32>(0x7fffffffu));
57
  {% else %}
58
  y[i] = {{ vectorScalar }}(abs(vec4<f32>(xv)));
59
  {% endif %}