sync 91d990483a17
Browse files- README.md +15 -12
- build/webgpu/bench.json +0 -1
- build/webgpu/datamove-elementwise-copy.wgsl.jinja +0 -3
- build/webgpu/datamove-flat-copy.wgsl.jinja +1 -4
- build/webgpu/manifest.json +27 -80
- build/webgpu/metadata.json +11 -8
- build/webgpu/test.json +3 -4
README.md
CHANGED
|
@@ -18,16 +18,16 @@ See the [ONNX `Squeeze` spec](https://onnx.ai/onnx/operators/onnx__Squeeze.html)
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
-
| Name |
|
| 22 |
-
| --- | --- | --- | --- | --- | --- | --- |
|
| 23 |
-
| `data` | `
|
| 24 |
-
| `axes` | `
|
| 25 |
|
| 26 |
## Outputs
|
| 27 |
|
| 28 |
-
| Name |
|
| 29 |
-
| --- | --- | --- | --- | --- | --- |
|
| 30 |
-
| `squeezed` | `
|
| 31 |
|
| 32 |
## Type constraints
|
| 33 |
|
|
@@ -38,7 +38,7 @@ See the [ONNX `Squeeze` spec](https://onnx.ai/onnx/operators/onnx__Squeeze.html)
|
|
| 38 |
|
| 39 |
## Files
|
| 40 |
|
| 41 |
-
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
|
| 42 |
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
|
| 43 |
- [`test.json`](build/webgpu/test.json) — correctness cases
|
| 44 |
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
|
|
@@ -47,15 +47,18 @@ See the [ONNX `Squeeze` spec](https://onnx.ai/onnx/operators/onnx__Squeeze.html)
|
|
| 47 |
|
| 48 |
## Use with `@huggingface/kernels`
|
| 49 |
|
| 50 |
-
|
|
|
|
|
|
|
| 51 |
|
| 52 |
-
|
| 53 |
|
| 54 |
-
|
| 55 |
|
| 56 |
-
|
| 57 |
|
| 58 |
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
|
|
|
|
| 59 |
|
| 60 |
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
|
| 61 |
|
|
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
+
| Name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
|
| 22 |
+
| --- | --- | --- | --- | --- | --- | --- |
|
| 23 |
+
| `data` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | Input tensor with at least as many dimensions as the largest axis index specified. | required |
|
| 24 |
+
| `axes` | `S` | `int32` | `1` | — | Logical int64 1-D tensor of axis indices to squeeze; negative values count from the back, signed axes use int32 WebGPU storage, and omission squeezes all size-1 dimensions. | optional |
|
| 25 |
|
| 26 |
## Outputs
|
| 27 |
|
| 28 |
+
| Name | Logical dtype | Rank | Shape | Description | Presence |
|
| 29 |
+
| --- | --- | --- | --- | --- | --- |
|
| 30 |
+
| `squeezed` | `T` | — | — | Reshaped tensor with the same data as the input but with the specified size-1 dimensions removed. | required |
|
| 31 |
|
| 32 |
## Type constraints
|
| 33 |
|
|
|
|
| 38 |
|
| 39 |
## Files
|
| 40 |
|
| 41 |
+
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
|
| 42 |
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
|
| 43 |
- [`test.json`](build/webgpu/test.json) — correctness cases
|
| 44 |
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
|
|
|
|
| 47 |
|
| 48 |
## Use with `@huggingface/kernels`
|
| 49 |
|
| 50 |
+
```sh
|
| 51 |
+
npm install --save-exact @huggingface/kernels@0.0.1-preview.2
|
| 52 |
+
```
|
| 53 |
|
| 54 |
+
Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
|
| 55 |
|
| 56 |
+
This example supplies explicit metadata for:
|
| 57 |
|
| 58 |
+
- `squeezed`
|
| 59 |
|
| 60 |
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
|
| 61 |
+
It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
|
| 62 |
|
| 63 |
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
|
| 64 |
|
build/webgpu/bench.json
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
{
|
| 2 |
-
"op": "ai.onnx.Squeeze",
|
| 3 |
"tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
|
| 4 |
"cases": [
|
| 5 |
{
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
|
| 3 |
"cases": [
|
| 4 |
{
|
build/webgpu/datamove-elementwise-copy.wgsl.jinja
CHANGED
|
@@ -1,6 +1,3 @@
|
|
| 1 |
-
{% if usesF16 %}
|
| 2 |
-
enable f16;
|
| 3 |
-
{% endif %}
|
| 4 |
{{ env.wgsl.resourceDeclarations }}
|
| 5 |
|
| 6 |
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
{{ env.wgsl.resourceDeclarations }}
|
| 2 |
|
| 3 |
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
|
build/webgpu/datamove-flat-copy.wgsl.jinja
CHANGED
|
@@ -1,9 +1,6 @@
|
|
| 1 |
-
{% if usesF16 %}
|
| 2 |
-
enable f16;
|
| 3 |
-
{% endif %}
|
| 4 |
{{ env.wgsl.resourceDeclarations }}
|
| 5 |
|
| 6 |
-
const COUNT: u32 = {{
|
| 7 |
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
|
| 8 |
|
| 9 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
{{ env.wgsl.resourceDeclarations }}
|
| 2 |
|
| 3 |
+
const COUNT: u32 = {{ count }}u;
|
| 4 |
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
|
| 5 |
|
| 6 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
build/webgpu/manifest.json
CHANGED
|
@@ -2,114 +2,61 @@
|
|
| 2 |
"domain": "ai.onnx",
|
| 3 |
"name": "Squeeze",
|
| 4 |
"sinceVersion": 13,
|
| 5 |
-
"
|
| 6 |
-
|
| 7 |
-
{
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
"description": "Input tensor with at least as many dimensions as the largest axis index specified."
|
| 11 |
-
},
|
| 12 |
-
{
|
| 13 |
-
"role": "axes",
|
| 14 |
-
"dtype": "S",
|
| 15 |
-
"rank": 1,
|
| 16 |
-
"optional": true,
|
| 17 |
-
"description": "Logical int64 1-D tensor of axis indices to squeeze; negative values count from the back, signed axes use int32 WebGPU storage, and omission squeezes all size-1 dimensions."
|
| 18 |
-
}
|
| 19 |
-
],
|
| 20 |
-
"outputs": [
|
| 21 |
-
{
|
| 22 |
-
"role": "squeezed",
|
| 23 |
-
"dtype": "T",
|
| 24 |
-
"description": "Reshaped tensor with the same data as the input but with the specified size-1 dimensions removed."
|
| 25 |
-
}
|
| 26 |
-
],
|
| 27 |
"typeConstraints": {
|
| 28 |
"T": ["float32", "float16", "int32", "uint32", "int16", "uint8", "int8", "bool"],
|
| 29 |
"S": ["int64"]
|
| 30 |
},
|
| 31 |
-
"
|
| 32 |
-
"data": { "kind": "tensor", "semantic": "data", "role": "input" },
|
| 33 |
-
"axes": {
|
| 34 |
-
"kind": "tensor",
|
| 35 |
-
"semantic": "axes",
|
| 36 |
-
"role": "axes",
|
| 37 |
-
"dtype": "int32",
|
| 38 |
-
"narrowing": "checked",
|
| 39 |
-
"required": false
|
| 40 |
-
},
|
| 41 |
-
"squeezed": { "kind": "tensor", "semantic": "squeezed", "role": "output" }
|
| 42 |
-
},
|
| 43 |
-
"tunables": { "WORKGROUP_SIZE": 256 },
|
| 44 |
"derive": { "axesContract": "not present.axes or ranks.axes == 1" },
|
| 45 |
-
"
|
| 46 |
"variants": [
|
| 47 |
{
|
| 48 |
"id": "copy_vec4",
|
| 49 |
"priority": 20,
|
| 50 |
-
"when": ["
|
| 51 |
-
"
|
| 52 |
"passes": [
|
| 53 |
{
|
| 54 |
"id": "main",
|
| 55 |
"name": "Squeeze.copyVec4",
|
| 56 |
-
"
|
|
|
|
| 57 |
"bindings": [
|
| 58 |
-
{
|
| 59 |
-
|
| 60 |
-
"arg": "data",
|
| 61 |
-
"semantic": "data",
|
| 62 |
-
"buffer": { "type": "read-only-storage" },
|
| 63 |
-
"elementType": "$vectorScalar"
|
| 64 |
-
},
|
| 65 |
-
{
|
| 66 |
-
"name": "dst",
|
| 67 |
-
"arg": "squeezed",
|
| 68 |
-
"semantic": "squeezed",
|
| 69 |
-
"buffer": { "type": "storage" },
|
| 70 |
-
"elementType": "$vectorScalar"
|
| 71 |
-
}
|
| 72 |
],
|
| 73 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 74 |
"viewAlias": [{ "input": "src", "output": "dst" }]
|
| 75 |
}
|
| 76 |
]
|
| 77 |
},
|
| 78 |
{
|
| 79 |
"id": "copy",
|
| 80 |
-
"
|
| 81 |
-
"constants": { "scalar": "dtypes.T" },
|
| 82 |
"passes": [
|
| 83 |
{
|
| 84 |
"id": "main",
|
| 85 |
"name": "Squeeze",
|
| 86 |
"shader": "datamove-elementwise-copy.wgsl.jinja",
|
| 87 |
"bindings": [
|
| 88 |
-
{
|
| 89 |
-
|
| 90 |
-
|
| 91 |
-
"semantic": "data",
|
| 92 |
-
"buffer": { "type": "read-only-storage" },
|
| 93 |
-
"elementType": "$scalar"
|
| 94 |
-
},
|
| 95 |
-
{
|
| 96 |
-
"name": "y",
|
| 97 |
-
"arg": "squeezed",
|
| 98 |
-
"semantic": "squeezed",
|
| 99 |
-
"buffer": { "type": "storage" },
|
| 100 |
-
"elementType": "$scalar"
|
| 101 |
-
},
|
| 102 |
-
{
|
| 103 |
-
"name": "params",
|
| 104 |
-
"semantic": "kernel.params",
|
| 105 |
-
"buffer": { "type": "uniform" },
|
| 106 |
-
"struct": {
|
| 107 |
-
"name": "Params",
|
| 108 |
-
"fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.squeezed)" }]
|
| 109 |
-
}
|
| 110 |
-
}
|
| 111 |
],
|
| 112 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 113 |
"viewAlias": [{ "input": "x", "output": "y" }]
|
| 114 |
}
|
| 115 |
]
|
|
|
|
| 2 |
"domain": "ai.onnx",
|
| 3 |
"name": "Squeeze",
|
| 4 |
"sinceVersion": 13,
|
| 5 |
+
"inputs": {
|
| 6 |
+
"data": { "dtype": "T" },
|
| 7 |
+
"axes": { "dtype": "S", "rank": 1, "optional": true, "storage": "int32", "narrowing": "checked" }
|
| 8 |
+
},
|
| 9 |
+
"outputs": { "squeezed": { "dtype": "T" } },
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 10 |
"typeConstraints": {
|
| 11 |
"T": ["float32", "float16", "int32", "uint32", "int16", "uint8", "int8", "bool"],
|
| 12 |
"S": ["int64"]
|
| 13 |
},
|
| 14 |
+
"tunables": { "WORKGROUP_SIZE": { "default": 256 } },
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 15 |
"derive": { "axesContract": "not present.axes or ranks.axes == 1" },
|
| 16 |
+
"when": ["axesContract", "numel(shapes.data) == numel(shapes.squeezed)", "f16Ok(dtypes.T)"],
|
| 17 |
"variants": [
|
| 18 |
{
|
| 19 |
"id": "copy_vec4",
|
| 20 |
"priority": 20,
|
| 21 |
+
"when": ["numel(shapes.data) > 0", "numel(shapes.data) % 4 == 0"],
|
| 22 |
+
"derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
|
| 23 |
"passes": [
|
| 24 |
{
|
| 25 |
"id": "main",
|
| 26 |
"name": "Squeeze.copyVec4",
|
| 27 |
+
"shader": "datamove-flat-copy.wgsl.jinja",
|
| 28 |
+
"derive": { "count": "numel(shapes.squeezed) / 4" },
|
| 29 |
"bindings": [
|
| 30 |
+
{ "arg": "data", "name": "src", "elementType": "$vectorScalar" },
|
| 31 |
+
{ "arg": "squeezed", "name": "dst", "elementType": "$vectorScalar" }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 32 |
],
|
| 33 |
+
"dispatch": {
|
| 34 |
+
"x": "min(ceilDiv((numel(shapes.squeezed) / 4), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
|
| 35 |
+
"y": 1,
|
| 36 |
+
"z": 1
|
| 37 |
+
},
|
| 38 |
"viewAlias": [{ "input": "src", "output": "dst" }]
|
| 39 |
}
|
| 40 |
]
|
| 41 |
},
|
| 42 |
{
|
| 43 |
"id": "copy",
|
| 44 |
+
"derive": { "scalar": "dtypes.T" },
|
|
|
|
| 45 |
"passes": [
|
| 46 |
{
|
| 47 |
"id": "main",
|
| 48 |
"name": "Squeeze",
|
| 49 |
"shader": "datamove-elementwise-copy.wgsl.jinja",
|
| 50 |
"bindings": [
|
| 51 |
+
{ "arg": "data", "name": "x", "elementType": "$scalar" },
|
| 52 |
+
{ "arg": "squeezed", "name": "y", "elementType": "$scalar" },
|
| 53 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.squeezed)" }] }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 54 |
],
|
| 55 |
+
"dispatch": {
|
| 56 |
+
"x": "min(ceilDiv((numel(shapes.squeezed)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
|
| 57 |
+
"y": 1,
|
| 58 |
+
"z": 1
|
| 59 |
+
},
|
| 60 |
"viewAlias": [{ "input": "x", "output": "y" }]
|
| 61 |
}
|
| 62 |
]
|
build/webgpu/metadata.json
CHANGED
|
@@ -1,19 +1,22 @@
|
|
| 1 |
{
|
| 2 |
"name": "ai.onnx.Squeeze",
|
| 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 |
-
"datamove-elementwise-copy.wgsl.jinja": "
|
| 12 |
-
"datamove-flat-copy.wgsl.jinja": "
|
| 13 |
-
"manifest.json": "
|
| 14 |
-
"test.json": "
|
| 15 |
}
|
| 16 |
},
|
| 17 |
-
"provenance": { "kernel": { "sha": "
|
| 18 |
-
"webgpu": {
|
|
|
|
|
|
|
|
|
|
| 19 |
}
|
|
|
|
| 1 |
{
|
| 2 |
"name": "ai.onnx.Squeeze",
|
| 3 |
+
"id": "_ai_onnx_squeeze_webgpu_d1535eb",
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
+
"bench.json": "1WC6P+yyucqPFwrdWD/aka7rpPv3lqBN6HF9diF8KhQ=",
|
| 11 |
+
"datamove-elementwise-copy.wgsl.jinja": "Q1WuZCqcDf9b6cbT4rLilpeWj6327t/bSFmKr6IVBoQ=",
|
| 12 |
+
"datamove-flat-copy.wgsl.jinja": "g9d62mer2bmHfbScqX5CIT5Zq/Pdinac0W2dH0CD+5s=",
|
| 13 |
+
"manifest.json": "ZGJbhNJTmqaJQUFcnMmznVL5t66lXFS6aQyFvFXRhT0=",
|
| 14 |
+
"test.json": "KYRUy5FOrvaucqzu/7scvf/NrMiDWTdWU8e1tyzXwak="
|
| 15 |
}
|
| 16 |
},
|
| 17 |
+
"provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
|
| 18 |
+
"webgpu": {
|
| 19 |
+
"manifestSpec": "2.0",
|
| 20 |
+
"variants": { "copy_vec4": ["datamove-flat-copy.wgsl.jinja"], "copy": ["datamove-elementwise-copy.wgsl.jinja"] }
|
| 21 |
+
}
|
| 22 |
}
|
build/webgpu/test.json
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
{
|
| 2 |
-
"op": "ai.onnx.Squeeze",
|
| 3 |
"fixtureArrays": {
|
| 4 |
"bool_axis1_rank4_to_rank3_input_data": [1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 1]
|
| 5 |
},
|
|
@@ -259,7 +258,7 @@
|
|
| 259 |
"outputs": { "squeezed": { "dtype": "float32", "shape": [3, 4, 5] } },
|
| 260 |
"provenance": {
|
| 261 |
"source": "cmake/external/onnx/onnx/backend/test/data/node/test_squeeze",
|
| 262 |
-
"notes": "
|
| 263 |
}
|
| 264 |
},
|
| 265 |
{
|
|
@@ -278,7 +277,7 @@
|
|
| 278 |
"outputs": { "squeezed": { "dtype": "float32", "shape": [1, 3, 5] } },
|
| 279 |
"provenance": {
|
| 280 |
"source": "cmake/external/onnx/onnx/backend/test/data/node/test_squeeze_negative_axes",
|
| 281 |
-
"notes": "
|
| 282 |
}
|
| 283 |
},
|
| 284 |
{
|
|
@@ -380,7 +379,7 @@
|
|
| 380 |
{
|
| 381 |
"name": "f32_vec4_copy_1024",
|
| 382 |
"provenance": {
|
| 383 |
-
"notes": "
|
| 384 |
},
|
| 385 |
"inputs": {
|
| 386 |
"data": {
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"fixtureArrays": {
|
| 3 |
"bool_axis1_rank4_to_rank3_input_data": [1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 1]
|
| 4 |
},
|
|
|
|
| 258 |
"outputs": { "squeezed": { "dtype": "float32", "shape": [3, 4, 5] } },
|
| 259 |
"provenance": {
|
| 260 |
"source": "cmake/external/onnx/onnx/backend/test/data/node/test_squeeze",
|
| 261 |
+
"notes": "The integer axes tensor exactly represents the source ONNX int64 axis values."
|
| 262 |
}
|
| 263 |
},
|
| 264 |
{
|
|
|
|
| 277 |
"outputs": { "squeezed": { "dtype": "float32", "shape": [1, 3, 5] } },
|
| 278 |
"provenance": {
|
| 279 |
"source": "cmake/external/onnx/onnx/backend/test/data/node/test_squeeze_negative_axes",
|
| 280 |
+
"notes": "The integer axes tensor exactly represents the source ONNX int64 axis values."
|
| 281 |
}
|
| 282 |
},
|
| 283 |
{
|
|
|
|
| 379 |
{
|
| 380 |
"name": "f32_vec4_copy_1024",
|
| 381 |
"provenance": {
|
| 382 |
+
"notes": "A small aligned float32 payload exercises vectorized copying while removing all unit axes."
|
| 383 |
},
|
| 384 |
"inputs": {
|
| 385 |
"data": {
|