sync 91d990483a17
Browse files- README.md +11 -7
- build/webgpu/bench.json +0 -1
- build/webgpu/manifest.json +92 -50
- build/webgpu/metadata.json +15 -8
- build/webgpu/test.json +2 -3
- build/webgpu/unary-scalar.wgsl.jinja +18 -7
- build/webgpu/unary-vec4.wgsl.jinja +27 -8
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 |
|
| 22 |
| --- | --- | --- | --- | --- | --- | --- |
|
| 23 |
-
| `
|
| 24 |
|
| 25 |
## Outputs
|
| 26 |
|
| 27 |
-
| Name |
|
| 28 |
| --- | --- | --- | --- | --- | --- | --- |
|
| 29 |
-
| `
|
| 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 |
-
|
| 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 | 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 |
-
"
|
| 6 |
-
"
|
| 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 |
-
"
|
| 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 |
-
"
|
| 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 |
-
"
|
| 38 |
-
"
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 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 |
-
"
|
| 60 |
"passes": [
|
| 61 |
{
|
| 62 |
"id": "main",
|
| 63 |
"name": "Log",
|
| 64 |
-
"
|
| 65 |
-
"
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 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": "
|
| 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.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": "
|
| 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": "
|
| 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>
|
| 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
{% 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>
|
| 2 |
-
//
|
| 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>
|
| 38 |
// 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
|
| 39 |
-
//
|
| 40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
}
|