Xenova HF Staff commited on
Commit
c0e2f2a
·
verified ·
1 Parent(s): ebc1b30

sync 91d990483a17

Browse files
README.md CHANGED
@@ -18,15 +18,15 @@ See the [ONNX `Log` spec](https://onnx.ai/onnx/operators/onnx__Log.html) for the
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
- | `input` | `x` | `T` | — | — | Values whose natural logarithms are computed elementwise. | required |
24
 
25
  ## Outputs
26
 
27
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
28
  | --- | --- | --- | --- | --- | --- | --- |
29
- | `output` | `y` | `T` | same as `input` | same as `input` | The natural log of the input tensor, computed elementwise. | required |
30
 
31
  ## Type constraints
32
 
@@ -36,7 +36,7 @@ See the [ONNX `Log` spec](https://onnx.ai/onnx/operators/onnx__Log.html) for the
36
 
37
  ## Files
38
 
39
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -45,10 +45,14 @@ See the [ONNX `Log` spec](https://onnx.ai/onnx/operators/onnx__Log.html) for the
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
49
- It then allocates the result tensors automatically.
 
 
 
50
 
51
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
52
 
53
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
54
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
+ | `x` | `input` | `T` | — | — | Values whose natural logarithms are computed elementwise. | required |
24
 
25
  ## Outputs
26
 
27
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
28
  | --- | --- | --- | --- | --- | --- | --- |
29
+ | `y` | `output` | `T` | same as `x` | same as `x` | The natural log of the input tensor, computed elementwise. | required |
30
 
31
  ## Type constraints
32
 
 
36
 
37
  ## Files
38
 
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
 
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.
53
 
54
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
55
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
56
 
57
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
58
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "ai.onnx.Log",
3
  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
4
  "cases": [
5
  {
 
1
  {
 
2
  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
3
  "cases": [
4
  {
build/webgpu/manifest.json CHANGED
@@ -2,30 +2,16 @@
2
  "domain": "ai.onnx",
3
  "name": "Log",
4
  "sinceVersion": 13,
5
- "description": "Computes the natural logarithm of each element of the input tensor, producing an output of the same shape.",
6
- "inputs": [
7
- { "role": "input", "dtype": "T", "description": "Values whose natural logarithms are computed elementwise." }
8
- ],
9
- "outputs": [
10
- {
11
- "role": "output",
12
- "dtype": "T",
13
- "rank": "ranks.input",
14
- "description": "The natural log of the input tensor, computed elementwise.",
15
- "shape": "shapes.input"
16
- }
17
- ],
18
  "typeConstraints": { "T": ["float32", "float16"] },
19
- "args": {
20
- "x": { "kind": "tensor", "semantic": "input", "role": "input" },
21
- "y": { "kind": "tensor", "semantic": "output", "role": "output" }
22
- },
23
- "tunables": { "WORKGROUP_SIZE": 256 },
24
  "variants": [
25
  {
26
  "id": "same_layout_vec4",
 
27
  "when": ["numel(shapes.x) > 0", "numel(shapes.x) % 4 == 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
28
- "constants": {
29
  "scalar": "dtypes.T",
30
  "usesF16": "dtypes.T == \"f16\"",
31
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
@@ -34,47 +20,103 @@
34
  {
35
  "id": "main",
36
  "name": "Log.vec4",
37
- "source": { "shader": "unary-vec4.wgsl.jinja", "inputs": { "op": "\"log\"" } },
38
- "bindings": [
39
- { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$vectorScalar" },
40
- { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
41
- {
42
- "name": "params",
43
- "semantic": "kernel.params",
44
- "buffer": { "type": "uniform" },
45
- "struct": {
46
- "name": "Params",
47
- "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }]
48
- }
49
- }
50
- ],
51
- "dispatch": { "threads": "numel(shapes.y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
52
  }
53
- ],
54
- "priority": 20
55
  },
56
  {
57
  "id": "scalar",
58
  "when": ["numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
59
- "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
60
  "passes": [
61
  {
62
  "id": "main",
63
  "name": "Log",
64
- "source": { "shader": "unary-scalar.wgsl.jinja", "inputs": { "op": "\"log\"", "itemsPerInvocation": 4 } },
65
- "bindings": [
66
- { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
67
- { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
68
- {
69
- "name": "params",
70
- "semantic": "kernel.params",
71
- "buffer": { "type": "uniform" },
72
- "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
73
- }
74
- ],
75
- "dispatch": { "threads": "ceilDiv(numel(shapes.y), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
76
  }
77
  ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
78
  }
79
- ]
80
  }
 
2
  "domain": "ai.onnx",
3
  "name": "Log",
4
  "sinceVersion": 13,
5
+ "inputs": { "x": { "onnx": "input", "dtype": "T" } },
6
+ "outputs": { "y": { "onnx": "output", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
 
 
 
 
 
 
 
 
 
 
 
7
  "typeConstraints": { "T": ["float32", "float16"] },
8
+ "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
 
 
 
 
9
  "variants": [
10
  {
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 ~ \">\""
 
20
  {
21
  "id": "main",
22
  "name": "Log.vec4",
23
+ "shader": "unary-vec4.wgsl.jinja",
24
+ "derive": {
25
+ "op": "\"log\"",
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)",
31
+ "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y) / 4, 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
32
+ "z": 1
33
+ }
 
 
 
 
34
  }
35
+ ]
 
36
  },
37
  {
38
  "id": "scalar",
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": "Log",
45
+ "shader": "unary-scalar.wgsl.jinja",
46
+ "derive": { "op": "\"log\"", "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)",
51
+ "z": 1
52
+ }
 
 
 
 
53
  }
54
  ]
55
+ },
56
+ {
57
+ "id": "same_layout_vec4_tail",
58
+ "priority": 19,
59
+ "when": ["numel(shapes.x) > 4", "numel(shapes.x) % 4 != 0", "numel(shapes.x) == numel(shapes.y)", "f16Ok(dtypes.T)"],
60
+ "derive": {
61
+ "scalar": "dtypes.T",
62
+ "usesF16": "dtypes.T == \"f16\"",
63
+ "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
64
+ },
65
+ "passes": [
66
+ {
67
+ "id": "bulk",
68
+ "name": "Log.vec4Bulk",
69
+ "shader": "unary-vec4.wgsl.jinja",
70
+ "derive": {
71
+ "op": "\"log\"",
72
+ "vec4PerThread": "4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1"
73
+ },
74
+ "bindings": ["x", "y", "params_unary_tail"],
75
+ "dispatch": {
76
+ "x": "min(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
77
+ "y": "ceilDiv(ceilDiv((ceilDiv(floor(numel(shapes.y) / 4), 4 if numel(shapes.y) * dtypeBytes(tensorDtypes.y) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
78
+ "z": 1
79
+ }
80
+ },
81
+ {
82
+ "id": "tail",
83
+ "name": "Log.tail",
84
+ "shader": "unary-scalar.wgsl.jinja",
85
+ "derive": { "op": "\"log\"", "itemsPerInvocation": 4, "tailOnly": true },
86
+ "bindings": ["x_2", "y_2", "params_2_unary_tail"],
87
+ "dispatch": {
88
+ "x": "min(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
89
+ "y": "ceilDiv(ceilDiv((1), (tunables.WORKGROUP_SIZE)), 65535)",
90
+ "z": 1
91
+ }
92
+ }
93
+ ]
94
+ }
95
+ ],
96
+ "bindings": {
97
+ "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
98
+ "y": { "buffer": "storage", "elementType": "$vectorScalar" },
99
+ "params_unary": {
100
+ "buffer": "uniform",
101
+ "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y) / 4" }],
102
+ "name": "params"
103
+ },
104
+ "x_2": { "buffer": "read-only-storage", "name": "x", "elementType": "$scalar" },
105
+ "y_2": { "buffer": "storage", "name": "y", "elementType": "$scalar" },
106
+ "params_2_unary": {
107
+ "buffer": "uniform",
108
+ "name": "params",
109
+ "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
110
+ },
111
+ "params_unary_tail": {
112
+ "buffer": "uniform",
113
+ "struct": [{ "name": "count", "type": "u32", "value": "floor(numel(shapes.y) / 4)" }],
114
+ "name": "params"
115
+ },
116
+ "params_2_unary_tail": {
117
+ "buffer": "uniform",
118
+ "name": "params",
119
+ "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }]
120
  }
121
+ }
122
  }
build/webgpu/metadata.json CHANGED
@@ -1,19 +1,26 @@
1
  {
2
  "name": "ai.onnx.Log",
3
- "id": "_ai_onnx_log_webgpu_2cf8591",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "r7qXaSJXGGynT+ARGoJs6iYyKU+pTj7T7TbCcS1LRLs=",
11
- "manifest.json": "8qaXwo8hQlTvqDgxd6nCxHGAf8k11eDjoUNfKRvG164=",
12
- "test.json": "B1RS6+XfT6gcxEtTg2ZGEnEkyMawcYaeCBTIuntIGtc=",
13
- "unary-scalar.wgsl.jinja": "mfdeF8kKBQgPX6wT3gRXd/fnkMK6WBCOHOO7Y+rNFVY=",
14
- "unary-vec4.wgsl.jinja": "1Kt6VZ/JRGUyA4Im6S/4ND8TGN3YtreG3gTtpiyTKEg="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
18
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Log" }
 
 
 
 
 
 
 
19
  }
 
1
  {
2
  "name": "ai.onnx.Log",
3
+ "id": "_ai_onnx_log_webgpu_c292e60",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "rBR41F3ZZnKOmPZu2RyG+Ubcy0fLTaNUQEFGa6BXdL0=",
11
+ "manifest.json": "h1/OmDD030GidTDxDfADacxt03KmZDberTGWKczOE4U=",
12
+ "test.json": "NbHcFJZksvqv48iHLQdE9fxVRS0odNEI2kCy4+zZasI=",
13
+ "unary-scalar.wgsl.jinja": "3OAztNgDK1dZP6smLXu7cf2KioCvEPmwFfijRJU6Y+8=",
14
+ "unary-vec4.wgsl.jinja": "6MeMmNGtoJNgLkZ3XECQPtYQYdcFuANJ3xWRKt4NzKM="
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
+ "scalar": ["unary-scalar.wgsl.jinja"],
23
+ "same_layout_vec4_tail": ["unary-scalar.wgsl.jinja", "unary-vec4.wgsl.jinja"]
24
+ }
25
+ }
26
  }
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "ai.onnx.Log",
3
  "cases": [
4
  {
5
  "name": "vector_positive",
@@ -47,7 +46,7 @@
47
  "provenance": {
48
  "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
49
  "test": "MathOpTest.Log",
50
- "notes": "Scalar-path companion for positive subnormal Log inputs."
51
  },
52
  "inputs": {
53
  "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1e-45, 1e-40, 1e-38] } }
@@ -97,7 +96,7 @@
97
  "provenance": {
98
  "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
99
  "test": "MathOpTest.Log",
100
- "notes": "Scalar-path companion: positive subnormal inputs are valid and should produce finite logs."
101
  },
102
  "inputs": {
103
  "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1e-45, 1e-40, 1e-39] } }
 
1
  {
 
2
  "cases": [
3
  {
4
  "name": "vector_positive",
 
46
  "provenance": {
47
  "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
48
  "test": "MathOpTest.Log",
49
+ "notes": "Unaligned scalar-path positive subnormal inputs exercise finite Log results."
50
  },
51
  "inputs": {
52
  "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1e-45, 1e-40, 1e-38] } }
 
96
  "provenance": {
97
  "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
98
  "test": "MathOpTest.Log",
99
+ "notes": "On the scalar path, positive subnormal inputs are valid and must produce finite logarithms."
100
  },
101
  "inputs": {
102
  "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1e-45, 1e-40, 1e-39] } }
build/webgpu/unary-scalar.wgsl.jinja CHANGED
@@ -1,21 +1,32 @@
1
  {% macro flat_tail_open() %}
2
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
3
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
4
  // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
5
- // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y).
6
- let invocation = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
7
  // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
8
  // the logical tensor length (or its storage binding) to be vec4 aligned.
9
- let begin = invocation * {{ source.itemsPerInvocation }}u;
10
- let end = min(begin + {{ source.itemsPerInvocation }}u, params.count);
 
 
 
 
 
 
 
 
 
 
 
11
  for (var i = begin; i < end; i = i + 1u) {
12
  {%- endmacro %}
13
  {% macro flat_tail_close() %}
14
  }
15
  {% endmacro %}
16
 
17
- // Scalar unary fallback. Each branch retains the operation's numeric hardening,
18
- // including Payne-Hanek trigonometric range reduction and NaN/overflow guards.
19
  {% if usesF16 %}
20
  enable f16;
21
  {% endif %}
 
1
  {% macro flat_tail_open() %}
2
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
3
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
4
  // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
5
+ // dispatch's per-axis workgroup fold width (the dispatch caps x and spills the rest into y).
6
+ let invocation = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
7
  // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
8
  // the logical tensor length (or its storage binding) to be vec4 aligned.
9
+ {% set itemsPerInvocation = itemsPerInvocation if itemsPerInvocation is defined else 4 %}
10
+ {% if tailOnly is defined and tailOnly %}
11
+ // Tail of a vec4 bulk pass: lane zero alone covers the elements the packed
12
+ // pass left behind, from the last multiple of the item count to the end.
13
+ if (invocation != 0u) {
14
+ return;
15
+ }
16
+ let begin = params.count - params.count % {{ itemsPerInvocation }}u;
17
+ let end = params.count;
18
+ {% else %}
19
+ let begin = invocation * {{ itemsPerInvocation }}u;
20
+ let end = min(begin + {{ itemsPerInvocation }}u, params.count);
21
+ {% endif %}
22
  for (var i = begin; i < end; i = i + 1u) {
23
  {%- endmacro %}
24
  {% macro flat_tail_close() %}
25
  }
26
  {% endmacro %}
27
 
28
+ // Scalar unary elementwise implementation. Specialization emits only the
29
+ // selected operation and any numerical helper it requires.
30
  {% if usesF16 %}
31
  enable f16;
32
  {% endif %}
build/webgpu/unary-vec4.wgsl.jinja CHANGED
@@ -1,8 +1,5 @@
1
- // Loads and stores vec4<T> (128 bits) while retaining scalar per-component
2
- // arithmetic, including per-component helper calls for guard-heavy operations.
3
- {% if usesF16 %}
4
- enable f16;
5
- {% endif %}
6
  {{ env.wgsl.resourceDeclarations }}
7
 
8
  {% set unaryDomainGuard = device.adapterInfo.architecture == "" or device.adapterInfo.architecture == "apple" %}
@@ -33,16 +30,38 @@ fn log_safe(x: f32) -> f32 {
33
  {% endif %}
34
  {{ emit_log_safe(unaryDomainGuard) }}
35
 
 
 
 
 
 
36
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
37
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
38
  // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
39
- // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y).
40
- let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 
 
 
 
 
 
 
 
 
 
 
 
41
  if (i >= params.count) {
42
  return;
43
  }
 
 
44
  let xv = x[i];
45
  let fv = vec4<f32>(xv);
46
  y[i] = {{ vectorScalar }}(vec4<f32>(
47
  log_safe(fv.x), log_safe(fv.y), log_safe(fv.z), log_safe(fv.w)));
 
 
 
48
  }
 
1
+ // Loads and stores vec4<T> while evaluating the selected unary operation per
2
+ // component.
 
 
 
3
  {{ env.wgsl.resourceDeclarations }}
4
 
5
  {% set unaryDomainGuard = device.adapterInfo.architecture == "" or device.adapterInfo.architecture == "apple" %}
 
30
  {% endif %}
31
  {{ emit_log_safe(unaryDomainGuard) }}
32
 
33
+ {% set vec4PerThread = vec4PerThread %}
34
+ {% if vec4PerThread > 1 %}
35
+ const ITEMS: u32 = {{ vec4PerThread }}u;
36
+ {% endif %}
37
+
38
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
39
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
40
  // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
41
+ // per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
42
+ {% if vec4PerThread > 1 %}
43
+ // Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
44
+ // access consecutive words on every step, while each lane can keep several
45
+ // independent loads in flight.
46
+ let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
47
+ let span = (params.count + ITEMS - 1u) / ITEMS;
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
+ let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
55
  if (i >= params.count) {
56
  return;
57
  }
58
+ {% endif %}
59
+
60
  let xv = x[i];
61
  let fv = vec4<f32>(xv);
62
  y[i] = {{ vectorScalar }}(vec4<f32>(
63
  log_safe(fv.x), log_safe(fv.y), log_safe(fv.z), log_safe(fv.w)));
64
+ {% if vec4PerThread > 1 %}
65
+ }
66
+ {% endif %}
67
  }