sync 91d990483a17
Browse files- README.md +13 -9
- build/webgpu/bench.json +0 -1
- build/webgpu/manifest.json +31 -71
- build/webgpu/metadata.json +11 -8
- build/webgpu/test.json +3 -6
- build/webgpu/unary-scalar.wgsl.jinja +8 -7
- build/webgpu/unary-vec4.wgsl.jinja +27 -8
README.md
CHANGED
|
@@ -18,15 +18,15 @@ See the [ONNX `Sign` spec](https://onnx.ai/onnx/operators/onnx__Sign.html) for t
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
-
| Name |
|
| 22 |
-
| --- | --- | --- | --- | --- | --- |
|
| 23 |
-
| `input` | `
|
| 24 |
|
| 25 |
## Outputs
|
| 26 |
|
| 27 |
-
| Name |
|
| 28 |
-
| --- | --- | --- | --- | --- | --- |
|
| 29 |
-
| `output` | `
|
| 30 |
|
| 31 |
## Type constraints
|
| 32 |
|
|
@@ -36,7 +36,7 @@ See the [ONNX `Sign` spec](https://onnx.ai/onnx/operators/onnx__Sign.html) for t
|
|
| 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 `Sign` spec](https://onnx.ai/onnx/operators/onnx__Sign.html) for t
|
|
| 45 |
|
| 46 |
## Use with `@huggingface/kernels`
|
| 47 |
|
| 48 |
-
|
| 49 |
-
|
|
|
|
|
|
|
|
|
|
| 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 | Logical dtype | Rank | Shape | Description | Presence |
|
| 22 |
+
| --- | --- | --- | --- | --- | --- |
|
| 23 |
+
| `input` | `T` | — | — | Input tensor of any numeric type. | required |
|
| 24 |
|
| 25 |
## Outputs
|
| 26 |
|
| 27 |
+
| Name | Logical dtype | Rank | Shape | Description | Presence |
|
| 28 |
+
| --- | --- | --- | --- | --- | --- |
|
| 29 |
+
| `output` | `T` | same as `input` | same as `input` | Elementwise sign of the input; same shape and type as the input. | 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.Sign",
|
| 3 |
"cases": [
|
| 4 |
{
|
| 5 |
"name": "1m_f32",
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"cases": [
|
| 3 |
{
|
| 4 |
"name": "1m_f32",
|
build/webgpu/manifest.json
CHANGED
|
@@ -2,98 +2,58 @@
|
|
| 2 |
"domain": "ai.onnx",
|
| 3 |
"name": "Sign",
|
| 4 |
"sinceVersion": 13,
|
| 5 |
-
"
|
| 6 |
-
"
|
| 7 |
-
"outputs": [
|
| 8 |
-
{
|
| 9 |
-
"role": "output",
|
| 10 |
-
"dtype": "T",
|
| 11 |
-
"rank": "ranks.input",
|
| 12 |
-
"description": "Elementwise sign of the input; same shape and type as the input.",
|
| 13 |
-
"shape": "shapes.input"
|
| 14 |
-
}
|
| 15 |
-
],
|
| 16 |
"typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"] },
|
| 17 |
-
"
|
| 18 |
-
|
| 19 |
-
"output": { "kind": "tensor", "semantic": "output", "role": "output" }
|
| 20 |
-
},
|
| 21 |
-
"tunables": { "WORKGROUP_SIZE": 256 },
|
| 22 |
-
"constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
|
| 23 |
"variants": [
|
| 24 |
{
|
| 25 |
"id": "same_layout_vec4",
|
|
|
|
| 26 |
"when": ["numel(shapes.input) > 0", "numel(shapes.input) % 4 == 0", "numel(shapes.input) == numel(shapes.output)", "f16Ok(dtypes.T)"],
|
| 27 |
-
"
|
| 28 |
"passes": [
|
| 29 |
{
|
| 30 |
"id": "main",
|
| 31 |
"name": "Sign.vec4",
|
| 32 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
"bindings": [
|
| 34 |
-
{
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
"semantic": "input",
|
| 38 |
-
"buffer": { "type": "read-only-storage" },
|
| 39 |
-
"elementType": "$vectorScalar"
|
| 40 |
-
},
|
| 41 |
-
{
|
| 42 |
-
"name": "y",
|
| 43 |
-
"arg": "output",
|
| 44 |
-
"semantic": "output",
|
| 45 |
-
"buffer": { "type": "storage" },
|
| 46 |
-
"elementType": "$vectorScalar"
|
| 47 |
-
},
|
| 48 |
-
{
|
| 49 |
-
"name": "params",
|
| 50 |
-
"semantic": "kernel.params",
|
| 51 |
-
"buffer": { "type": "uniform" },
|
| 52 |
-
"struct": {
|
| 53 |
-
"name": "Params",
|
| 54 |
-
"fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" }]
|
| 55 |
-
}
|
| 56 |
-
}
|
| 57 |
],
|
| 58 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 59 |
}
|
| 60 |
-
]
|
| 61 |
-
"priority": 20
|
| 62 |
},
|
| 63 |
{
|
| 64 |
"id": "elementwise",
|
| 65 |
-
"when": "numel(shapes.input) == numel(shapes.output) and (f16Ok(dtypes.T))",
|
| 66 |
"passes": [
|
| 67 |
{
|
| 68 |
"id": "main",
|
| 69 |
"name": "Sign",
|
| 70 |
-
"
|
|
|
|
| 71 |
"bindings": [
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
"semantic": "input",
|
| 76 |
-
"buffer": { "type": "read-only-storage" },
|
| 77 |
-
"elementType": "$scalar"
|
| 78 |
-
},
|
| 79 |
-
{
|
| 80 |
-
"name": "output",
|
| 81 |
-
"arg": "output",
|
| 82 |
-
"semantic": "output",
|
| 83 |
-
"buffer": { "type": "storage" },
|
| 84 |
-
"elementType": "$scalar"
|
| 85 |
-
},
|
| 86 |
-
{
|
| 87 |
-
"name": "params",
|
| 88 |
-
"semantic": "kernel.params",
|
| 89 |
-
"buffer": { "type": "uniform" },
|
| 90 |
-
"struct": {
|
| 91 |
-
"name": "Params",
|
| 92 |
-
"fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.output)" }]
|
| 93 |
-
}
|
| 94 |
-
}
|
| 95 |
],
|
| 96 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 97 |
}
|
| 98 |
]
|
| 99 |
}
|
|
|
|
| 2 |
"domain": "ai.onnx",
|
| 3 |
"name": "Sign",
|
| 4 |
"sinceVersion": 13,
|
| 5 |
+
"inputs": { "input": { "dtype": "T" } },
|
| 6 |
+
"outputs": { "output": { "dtype": "T", "rank": "ranks.input", "shape": "shapes.input" } },
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 7 |
"typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"] },
|
| 8 |
+
"tunables": { "WORKGROUP_SIZE": { "default": 256 } },
|
| 9 |
+
"derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
|
|
|
|
|
|
|
|
|
|
|
|
|
| 10 |
"variants": [
|
| 11 |
{
|
| 12 |
"id": "same_layout_vec4",
|
| 13 |
+
"priority": 20,
|
| 14 |
"when": ["numel(shapes.input) > 0", "numel(shapes.input) % 4 == 0", "numel(shapes.input) == numel(shapes.output)", "f16Ok(dtypes.T)"],
|
| 15 |
+
"derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
|
| 16 |
"passes": [
|
| 17 |
{
|
| 18 |
"id": "main",
|
| 19 |
"name": "Sign.vec4",
|
| 20 |
+
"shader": "unary-vec4.wgsl.jinja",
|
| 21 |
+
"derive": {
|
| 22 |
+
"op": "\"sign\"",
|
| 23 |
+
"vec4PerThread": "4 if numel(shapes.output) * dtypeBytes(tensorDtypes.output) <= 16777216 else 1"
|
| 24 |
+
},
|
| 25 |
"bindings": [
|
| 26 |
+
{ "arg": "input", "name": "x", "elementType": "$vectorScalar" },
|
| 27 |
+
{ "arg": "output", "name": "y", "elementType": "$vectorScalar" },
|
| 28 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" }] }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 29 |
],
|
| 30 |
+
"dispatch": {
|
| 31 |
+
"x": "min(ceilDiv((ceilDiv(numel(shapes.output) / 4, 4 if numel(shapes.output) * dtypeBytes(tensorDtypes.output) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 32 |
+
"y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.output) / 4, 4 if numel(shapes.output) * dtypeBytes(tensorDtypes.output) <= 16777216 else 1)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 33 |
+
"z": 1
|
| 34 |
+
}
|
| 35 |
}
|
| 36 |
+
]
|
|
|
|
| 37 |
},
|
| 38 |
{
|
| 39 |
"id": "elementwise",
|
| 40 |
+
"when": ["numel(shapes.input) == numel(shapes.output) and (f16Ok(dtypes.T))"],
|
| 41 |
"passes": [
|
| 42 |
{
|
| 43 |
"id": "main",
|
| 44 |
"name": "Sign",
|
| 45 |
+
"shader": "unary-scalar.wgsl.jinja",
|
| 46 |
+
"derive": { "op": "\"sign\"", "itemsPerInvocation": 4 },
|
| 47 |
"bindings": [
|
| 48 |
+
"input",
|
| 49 |
+
"output",
|
| 50 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.output)" }] }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 51 |
],
|
| 52 |
+
"dispatch": {
|
| 53 |
+
"x": "min(ceilDiv((ceilDiv(numel(shapes.output), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 54 |
+
"y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.output), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 55 |
+
"z": 1
|
| 56 |
+
}
|
| 57 |
}
|
| 58 |
]
|
| 59 |
}
|
build/webgpu/metadata.json
CHANGED
|
@@ -1,19 +1,22 @@
|
|
| 1 |
{
|
| 2 |
"name": "ai.onnx.Sign",
|
| 3 |
-
"id": "
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
-
"bench.json": "
|
| 11 |
-
"manifest.json": "
|
| 12 |
-
"test.json": "
|
| 13 |
-
"unary-scalar.wgsl.jinja": "
|
| 14 |
-
"unary-vec4.wgsl.jinja": "
|
| 15 |
}
|
| 16 |
},
|
| 17 |
-
"provenance": { "kernel": { "sha": "
|
| 18 |
-
"webgpu": {
|
|
|
|
|
|
|
|
|
|
| 19 |
}
|
|
|
|
| 1 |
{
|
| 2 |
"name": "ai.onnx.Sign",
|
| 3 |
+
"id": "_ai_onnx_sign_webgpu_7e75e78",
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
+
"bench.json": "BeK/pm3J1N63lmR/lECkTp1EWNnUeMFlM0F3OdgdhPk=",
|
| 11 |
+
"manifest.json": "EvDo7Z3CuI8sTL6W9FMajRDcM3UUbuipgGymQxGltlM=",
|
| 12 |
+
"test.json": "wqsQX1gxCRJBO5+/hR/fIFpxvnVtWZirFG1q+C+tF5I=",
|
| 13 |
+
"unary-scalar.wgsl.jinja": "ZmU+8SDLS8BCho41AuG6LIxEFTVx+xq8TBb3kyuTC+s=",
|
| 14 |
+
"unary-vec4.wgsl.jinja": "5eC2XytxRMc9QBwFeqlwy9ubxkWVr9aClnppDxHbvWQ="
|
| 15 |
}
|
| 16 |
},
|
| 17 |
+
"provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
|
| 18 |
+
"webgpu": {
|
| 19 |
+
"manifestSpec": "2.0",
|
| 20 |
+
"variants": { "same_layout_vec4": ["unary-vec4.wgsl.jinja"], "elementwise": ["unary-scalar.wgsl.jinja"] }
|
| 21 |
+
}
|
| 22 |
}
|
build/webgpu/test.json
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
{
|
| 2 |
-
"op": "ai.onnx.Sign",
|
| 3 |
"cases": [
|
| 4 |
{
|
| 5 |
"name": "int16_scalar_boundaries",
|
|
@@ -44,7 +43,7 @@
|
|
| 44 |
"provenance": {
|
| 45 |
"source": "onnxruntime/test/providers/cpu/math/sign_test.cc",
|
| 46 |
"test": "MathOpTest.Sign_float",
|
| 47 |
-
"notes": "
|
| 48 |
},
|
| 49 |
"inputs": {
|
| 50 |
"input": {
|
|
@@ -236,7 +235,7 @@
|
|
| 236 |
"provenance": {
|
| 237 |
"source": "onnxruntime/test/providers/cpu/math/sign_test.cc",
|
| 238 |
"test": "MathOpTest.Sign_float",
|
| 239 |
-
"notes": "Scalar
|
| 240 |
},
|
| 241 |
"inputs": {
|
| 242 |
"input": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": ["NaN", -2.0, 2.0] } }
|
|
@@ -301,9 +300,7 @@
|
|
| 301 |
},
|
| 302 |
{
|
| 303 |
"name": "f32_scalar_tail_4095",
|
| 304 |
-
"provenance": {
|
| 305 |
-
"notes": "Compact sibling for the large f32 Sign scalar-fallback benchmark; odd numel forces the elementwise path while preserving exact sign semantics."
|
| 306 |
-
},
|
| 307 |
"inputs": {
|
| 308 |
"input": {
|
| 309 |
"dtype": "float32",
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"cases": [
|
| 3 |
{
|
| 4 |
"name": "int16_scalar_boundaries",
|
|
|
|
| 43 |
"provenance": {
|
| 44 |
"source": "onnxruntime/test/providers/cpu/math/sign_test.cc",
|
| 45 |
"test": "MathOpTest.Sign_float",
|
| 46 |
+
"notes": "Vectorized Sign must preserve the signs of nonzero subnormal inputs."
|
| 47 |
},
|
| 48 |
"inputs": {
|
| 49 |
"input": {
|
|
|
|
| 235 |
"provenance": {
|
| 236 |
"source": "onnxruntime/test/providers/cpu/math/sign_test.cc",
|
| 237 |
"test": "MathOpTest.Sign_float",
|
| 238 |
+
"notes": "Scalar Sign must preserve NaN inputs."
|
| 239 |
},
|
| 240 |
"inputs": {
|
| 241 |
"input": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": ["NaN", -2.0, 2.0] } }
|
|
|
|
| 300 |
},
|
| 301 |
{
|
| 302 |
"name": "f32_scalar_tail_4095",
|
| 303 |
+
"provenance": { "notes": "An odd float32 element count exercises scalar Sign with exact sign semantics." },
|
|
|
|
|
|
|
| 304 |
"inputs": {
|
| 305 |
"input": {
|
| 306 |
"dtype": "float32",
|
build/webgpu/unary-scalar.wgsl.jinja
CHANGED
|
@@ -1,21 +1,22 @@
|
|
| 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 |
-
//
|
| 6 |
-
let invocation = gid.x + gid.y *
|
| 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 |
-
|
| 10 |
-
let
|
|
|
|
| 11 |
for (var i = begin; i < end; i = i + 1u) {
|
| 12 |
{%- endmacro %}
|
| 13 |
{% macro flat_tail_close() %}
|
| 14 |
}
|
| 15 |
{% endmacro %}
|
| 16 |
|
| 17 |
-
// Scalar unary
|
| 18 |
-
//
|
| 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 |
+
let begin = invocation * {{ itemsPerInvocation }}u;
|
| 11 |
+
let end = min(begin + {{ itemsPerInvocation }}u, params.count);
|
| 12 |
for (var i = begin; i < end; i = i + 1u) {
|
| 13 |
{%- endmacro %}
|
| 14 |
{% macro flat_tail_close() %}
|
| 15 |
}
|
| 16 |
{% endmacro %}
|
| 17 |
|
| 18 |
+
// Scalar unary elementwise implementation. Specialization emits only the
|
| 19 |
+
// selected operation and any numerical helper it requires.
|
| 20 |
{% if usesF16 %}
|
| 21 |
enable f16;
|
| 22 |
{% endif %}
|
build/webgpu/unary-vec4.wgsl.jinja
CHANGED
|
@@ -1,19 +1,35 @@
|
|
| 1 |
-
// Loads and stores vec4<T>
|
| 2 |
-
//
|
| 3 |
-
{% if usesF16 %}
|
| 4 |
-
enable f16;
|
| 5 |
-
{% endif %}
|
| 6 |
{{ env.wgsl.resourceDeclarations }}
|
| 7 |
|
| 8 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 9 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 10 |
-
fn main(@builtin(global_invocation_id) gid: vec3<u32>
|
| 11 |
// 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
|
| 12 |
-
//
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 14 |
if (i >= params.count) {
|
| 15 |
return;
|
| 16 |
}
|
|
|
|
|
|
|
| 17 |
let xv = x[i];
|
| 18 |
{% if scalar == "i32" %}
|
| 19 |
y[i] = select(select(vec4<i32>(0i), vec4<i32>(1i), xv > vec4<i32>(0i)), vec4<i32>(-1i), xv < vec4<i32>(0i));
|
|
@@ -28,4 +44,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 28 |
let signed = select(select(vec4<f32>(0.0), vec4<f32>(1.0), fv > vec4<f32>(0.0)), vec4<f32>(-1.0), fv < vec4<f32>(0.0));
|
| 29 |
y[i] = {{ vectorScalar }}(select(signed, fv, fv_is_nan));
|
| 30 |
{% endif %}
|
|
|
|
|
|
|
|
|
|
| 31 |
}
|
|
|
|
| 1 |
+
// Loads and stores vec4<T> while evaluating the selected unary operation per
|
| 2 |
+
// component.
|
|
|
|
|
|
|
|
|
|
| 3 |
{{ env.wgsl.resourceDeclarations }}
|
| 4 |
|
| 5 |
|
| 6 |
+
{% set vec4PerThread = vec4PerThread %}
|
| 7 |
+
{% if vec4PerThread > 1 %}
|
| 8 |
+
const ITEMS: u32 = {{ vec4PerThread }}u;
|
| 9 |
+
{% endif %}
|
| 10 |
+
|
| 11 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 12 |
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
| 13 |
// 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
|
| 14 |
+
// per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
|
| 15 |
+
{% if vec4PerThread > 1 %}
|
| 16 |
+
// Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
|
| 17 |
+
// access consecutive words on every step, while each lane can keep several
|
| 18 |
+
// independent loads in flight.
|
| 19 |
+
let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 20 |
+
let span = (params.count + ITEMS - 1u) / ITEMS;
|
| 21 |
+
for (var j = 0u; j < ITEMS; j = j + 1u) {
|
| 22 |
+
let i = tid + j * span;
|
| 23 |
+
if (i >= params.count) {
|
| 24 |
+
break;
|
| 25 |
+
}
|
| 26 |
+
{% else %}
|
| 27 |
+
let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 28 |
if (i >= params.count) {
|
| 29 |
return;
|
| 30 |
}
|
| 31 |
+
{% endif %}
|
| 32 |
+
|
| 33 |
let xv = x[i];
|
| 34 |
{% if scalar == "i32" %}
|
| 35 |
y[i] = select(select(vec4<i32>(0i), vec4<i32>(1i), xv > vec4<i32>(0i)), vec4<i32>(-1i), xv < vec4<i32>(0i));
|
|
|
|
| 44 |
let signed = select(select(vec4<f32>(0.0), vec4<f32>(1.0), fv > vec4<f32>(0.0)), vec4<f32>(-1.0), fv < vec4<f32>(0.0));
|
| 45 |
y[i] = {{ vectorScalar }}(select(signed, fv, fv_is_nan));
|
| 46 |
{% endif %}
|
| 47 |
+
{% if vec4PerThread > 1 %}
|
| 48 |
+
}
|
| 49 |
+
{% endif %}
|
| 50 |
}
|