Xenova HF Staff commited on
Commit
f72eeb2
·
verified ·
1 Parent(s): 221692d

sync 91d990483a17

Browse files
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 | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
22
- | --- | --- | --- | --- | --- | --- | --- | --- |
23
- | `data` | `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` | `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 | Bind key | Logical dtype | Rank | Shape | Description | Presence |
29
- | --- | --- | --- | --- | --- | --- | --- |
30
- | `squeezed` | `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,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
- The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
 
 
51
 
52
- The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
53
 
54
- - `squeezed`
55
 
56
- Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
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 = {{ source.count }}u;
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
- "description": "Removes single-dimensional entries from the shape of a tensor. If `axes` is provided, only those dimensions are removed; if omitted, all size-1 dimensions are removed. Selecting an axis whose size is not 1 raises an error.",
6
- "inputs": [
7
- {
8
- "role": "data",
9
- "dtype": "T",
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
- "args": {
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
- "constants": { "usesF16": "dtypes.T == \"f16\"" },
46
  "variants": [
47
  {
48
  "id": "copy_vec4",
49
  "priority": 20,
50
- "when": ["axesContract", "numel(shapes.data) > 0", "numel(shapes.data) % 4 == 0", "numel(shapes.data) == numel(shapes.squeezed)", "f16Ok(dtypes.T)"],
51
- "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
52
  "passes": [
53
  {
54
  "id": "main",
55
  "name": "Squeeze.copyVec4",
56
- "source": { "shader": "datamove-flat-copy.wgsl.jinja", "inputs": { "count": "numel(shapes.squeezed) / 4" } },
 
57
  "bindings": [
58
- {
59
- "name": "src",
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": { "gridStride": "numel(shapes.squeezed) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" },
 
 
 
 
74
  "viewAlias": [{ "input": "src", "output": "dst" }]
75
  }
76
  ]
77
  },
78
  {
79
  "id": "copy",
80
- "when": ["axesContract", "numel(shapes.data) == numel(shapes.squeezed)", "f16Ok(dtypes.T)"],
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
- "name": "x",
90
- "arg": "data",
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": { "gridStride": "numel(shapes.squeezed)", "workgroupSize": "tunables.WORKGROUP_SIZE" },
 
 
 
 
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": "_ai_onnx_squeeze_webgpu_832197a",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "nv9qMEydsiA2N2U6kTkSg/0bUKGXz3olKJQONAe+ALw=",
11
- "datamove-elementwise-copy.wgsl.jinja": "J5yC2bAddPiP+odXLgVGS3TJ9jeNsfRTvedKrj/fhZg=",
12
- "datamove-flat-copy.wgsl.jinja": "rEpeYMvIZQuQnB2f3lt51jrrklKICsqUTzCJPRTknNk=",
13
- "manifest.json": "X1FYKHl8iSCmL4jG2aIzk6pzQGL3QrYp3Jca594J2vQ=",
14
- "test.json": "McQKVkPXhYGWdGWzD6vp5gR1x5Cwujt0CkQGow0TGkg="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
18
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Squeeze" }
 
 
 
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": "ONNX int64 metadata/index tensors use framework int32/uint32 slots where representable."
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": "ONNX int64 metadata/index tensors use framework int32/uint32 slots where representable."
282
  }
283
  },
284
  {
@@ -380,7 +379,7 @@
380
  {
381
  "name": "f32_vec4_copy_1024",
382
  "provenance": {
383
- "notes": "Compact vec4-copy sibling for the 1M f32 Squeeze benchmark; keeps the all-unit-axis removal pattern with a small aligned payload."
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": {