Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.DepthToSpace

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
a7ec472
·
verified ·
1 Parent(s): 2fda425

sync 6fdf6301e2bb

Browse files
Files changed (5) hide show
  1. README.md +2 -2
  2. build/webgpu/depth-to-space-nchw-coarsened4.wgsl.jinja +7 -2
  3. build/webgpu/manifest.json +12 -22
  4. build/webgpu/metadata.json +6 -6
  5. build/webgpu/space-depth-permute.wgsl.jinja +10 -32
README.md CHANGED
@@ -48,14 +48,14 @@ Attributes and default values (overridable per request):
48
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
49
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
50
  - [`test.json`](build/webgpu/test.json) — correctness cases
51
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
52
  - [`depth-to-space-nchw-coarsened4.wgsl.jinja`](build/webgpu/depth-to-space-nchw-coarsened4.wgsl.jinja)
53
  - [`space-depth-permute.wgsl.jinja`](build/webgpu/space-depth-permute.wgsl.jinja)
54
 
55
  ## Use with `@huggingface/kernels`
56
 
57
  ```sh
58
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
59
  ```
60
 
61
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
48
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
49
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
50
  - [`test.json`](build/webgpu/test.json) — correctness cases
51
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
52
  - [`depth-to-space-nchw-coarsened4.wgsl.jinja`](build/webgpu/depth-to-space-nchw-coarsened4.wgsl.jinja)
53
  - [`space-depth-permute.wgsl.jinja`](build/webgpu/space-depth-permute.wgsl.jinja)
54
 
55
  ## Use with `@huggingface/kernels`
56
 
57
  ```sh
58
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
59
  ```
60
 
61
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/depth-to-space-nchw-coarsened4.wgsl.jinja CHANGED
@@ -1,3 +1,9 @@
 
 
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
 
3
  // Four adjacent output columns share N/C/H decoding and the block-row lookup.
@@ -17,8 +23,7 @@ const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
17
 
18
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
19
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
20
- let q = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
21
- if (q >= GROUPS) { return; }
22
 
23
  let owBase = (q % OUT_W_GROUPS) * 4u;
24
  var t = q / OUT_W_GROUPS;
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
6
+ if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
7
  {{ env.wgsl.resourceDeclarations }}
8
 
9
  // Four adjacent output columns share N/C/H decoding and the block-row lookup.
 
23
 
24
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
25
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
26
+ {{ flat_index_2d("WG", "q", "GROUPS", guardInline=true) }}
 
27
 
28
  let owBase = (q % OUT_W_GROUPS) * 4u;
29
  var t = q / OUT_W_GROUPS;
build/webgpu/manifest.json CHANGED
@@ -15,16 +15,23 @@
15
  "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"] },
16
  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
17
  "derive": {
 
18
  "shapeContract": "attrs.blocksize > 0 and ranks.input == 4 and ranks.output == 4 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) * attrs.blocksize * attrs.blocksize == dim(shapes.input, 1) and dim(shapes.output, 2) == dim(shapes.input, 2) * attrs.blocksize and dim(shapes.output, 3) == dim(shapes.input, 3) * attrs.blocksize and f16Ok(dtypes.T)",
19
- "scalar": "dtypes.T"
 
 
 
 
 
20
  },
21
  "when": ["shapeContract"],
22
- "bindings": { "output": { "buffer": "storage", "elementType": "$outputElement" } },
23
  "variants": [
24
  {
25
  "id": "nchw_vec4",
26
  "priority": 20,
27
  "when": ["dim(shapes.output, 3) % 4 == 0"],
 
28
  "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
29
  "passes": [
30
  {
@@ -32,17 +39,11 @@
32
  "name": "DepthToSpace",
33
  "shader": "space-depth-permute.wgsl.jinja",
34
  "derive": {
35
- "direction": "\"depthToSpace\"",
36
- "modeSpec": "attrs.mode",
37
  "vectorizedSpec": true,
38
- "blocksize": "attrs.blocksize",
39
  "count": "numel(shapes.output) / 4",
40
  "outWUnits": "dim(shapes.output, 3) / 4",
41
  "outH": "dim(shapes.output, 2)",
42
- "outC": "dim(shapes.output, 1)",
43
- "inC": "dim(shapes.input, 1)",
44
- "inH": "dim(shapes.input, 2)",
45
- "inW": "dim(shapes.input, 3)"
46
  },
47
  "bindings": ["input", "output"],
48
  "dispatch": {
@@ -63,18 +64,12 @@
63
  "name": "DepthToSpace",
64
  "shader": "space-depth-permute.wgsl.jinja",
65
  "derive": {
66
- "direction": "\"depthToSpace\"",
67
- "modeSpec": "attrs.mode",
68
  "vectorizedSpec": false,
69
- "blocksize": "attrs.blocksize",
70
  "count": "numel(shapes.output)",
71
  "outWUnits": "dim(shapes.output, 3)",
72
  "outC": "dim(shapes.output, 1)",
73
  "outH": "dim(shapes.output, 2)",
74
- "outW": "dim(shapes.output, 3)",
75
- "inC": "dim(shapes.input, 1)",
76
- "inH": "dim(shapes.input, 2)",
77
- "inW": "dim(shapes.input, 3)"
78
  },
79
  "bindings": ["input", "output"],
80
  "dispatch": {
@@ -96,16 +91,11 @@
96
  "name": "DepthToSpace.Coarsened4",
97
  "shader": "depth-to-space-nchw-coarsened4.wgsl.jinja",
98
  "derive": {
99
- "modeSpec": "attrs.mode",
100
- "blocksize": "attrs.blocksize",
101
  "groups": "dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.output, 2) * ceil(dim(shapes.output, 3) / 4)",
102
  "outWGroups": "ceil(dim(shapes.output, 3) / 4)",
103
  "outC": "dim(shapes.output, 1)",
104
  "outH": "dim(shapes.output, 2)",
105
- "outW": "dim(shapes.output, 3)",
106
- "inC": "dim(shapes.input, 1)",
107
- "inH": "dim(shapes.input, 2)",
108
- "inW": "dim(shapes.input, 3)"
109
  },
110
  "bindings": ["input", "output"],
111
  "dispatch": {
 
15
  "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"] },
16
  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
17
  "derive": {
18
+ "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
19
  "shapeContract": "attrs.blocksize > 0 and ranks.input == 4 and ranks.output == 4 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) * attrs.blocksize * attrs.blocksize == dim(shapes.input, 1) and dim(shapes.output, 2) == dim(shapes.input, 2) * attrs.blocksize and dim(shapes.output, 3) == dim(shapes.input, 3) * attrs.blocksize and f16Ok(dtypes.T)",
20
+ "scalar": "dtypes.T",
21
+ "modeSpec": "attrs.mode",
22
+ "blocksize": "attrs.blocksize",
23
+ "inC": "dim(shapes.input, 1)",
24
+ "inH": "dim(shapes.input, 2)",
25
+ "inW": "dim(shapes.input, 3)"
26
  },
27
  "when": ["shapeContract"],
28
+ "bindings": { "output": { "elementType": "$outputElement" } },
29
  "variants": [
30
  {
31
  "id": "nchw_vec4",
32
  "priority": 20,
33
  "when": ["dim(shapes.output, 3) % 4 == 0"],
34
+ "demoteWhen": ["variableSubgroup16To32 and dtypes.T == \"f16\""],
35
  "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
36
  "passes": [
37
  {
 
39
  "name": "DepthToSpace",
40
  "shader": "space-depth-permute.wgsl.jinja",
41
  "derive": {
 
 
42
  "vectorizedSpec": true,
 
43
  "count": "numel(shapes.output) / 4",
44
  "outWUnits": "dim(shapes.output, 3) / 4",
45
  "outH": "dim(shapes.output, 2)",
46
+ "outC": "dim(shapes.output, 1)"
 
 
 
47
  },
48
  "bindings": ["input", "output"],
49
  "dispatch": {
 
64
  "name": "DepthToSpace",
65
  "shader": "space-depth-permute.wgsl.jinja",
66
  "derive": {
 
 
67
  "vectorizedSpec": false,
 
68
  "count": "numel(shapes.output)",
69
  "outWUnits": "dim(shapes.output, 3)",
70
  "outC": "dim(shapes.output, 1)",
71
  "outH": "dim(shapes.output, 2)",
72
+ "outW": "dim(shapes.output, 3)"
 
 
 
73
  },
74
  "bindings": ["input", "output"],
75
  "dispatch": {
 
91
  "name": "DepthToSpace.Coarsened4",
92
  "shader": "depth-to-space-nchw-coarsened4.wgsl.jinja",
93
  "derive": {
 
 
94
  "groups": "dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.output, 2) * ceil(dim(shapes.output, 3) / 4)",
95
  "outWGroups": "ceil(dim(shapes.output, 3) / 4)",
96
  "outC": "dim(shapes.output, 1)",
97
  "outH": "dim(shapes.output, 2)",
98
+ "outW": "dim(shapes.output, 3)"
 
 
 
99
  },
100
  "bindings": ["input", "output"],
101
  "dispatch": {
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.DepthToSpace",
3
- "id": "_ai_onnx_depthtospace_webgpu_af18dda",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,15 +8,15 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "JU/Zd5udt4ll47K2EnVjXxjt0sEh9KmpBf5fZnZrXFY=",
11
- "depth-to-space-nchw-coarsened4.wgsl.jinja": "4xGeu/kV7B8e2wY47sCv1CFLAhW7NbhUkNFAO8N01jw=",
12
- "manifest.json": "4syDCQ7+CYWnOP8y3nQeAL0uF3R0PuXa2WGfqCUwI5I=",
13
- "space-depth-permute.wgsl.jinja": "AdqL6I5mdEY/Q+Pd8AMR7O4/IasCj9BhRd+8QDccF5A=",
14
  "test.json": "hs7y38XH3BuXDsaABYqenTgiC4pdBPPan6x4uUc0EU0="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": {
21
  "nchw_vec4": ["space-depth-permute.wgsl.jinja"],
22
  "nchw_const": ["space-depth-permute.wgsl.jinja"],
 
1
  {
2
  "name": "ai.onnx.DepthToSpace",
3
+ "id": "_ai_onnx_depthtospace_webgpu_6dcb40b",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "JU/Zd5udt4ll47K2EnVjXxjt0sEh9KmpBf5fZnZrXFY=",
11
+ "depth-to-space-nchw-coarsened4.wgsl.jinja": "W8vNJZ/U3oL0d11TQBVoSxOkJmvXNxolOI0xFXlG1bU=",
12
+ "manifest.json": "/2NFeZFb4YB+4NEuZeBjEgD6NuRQmZoAXjMRXWW4TjQ=",
13
+ "space-depth-permute.wgsl.jinja": "0fv5363hhXl9ojT6/Mz0lMAqdqHgoiZUSrnZTspEO30=",
14
  "test.json": "hs7y38XH3BuXDsaABYqenTgiC4pdBPPan6x4uUc0EU0="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
18
  "webgpu": {
19
+ "manifestSpec": "2.1",
20
  "variants": {
21
  "nchw_vec4": ["space-depth-permute.wgsl.jinja"],
22
  "nchw_const": ["space-depth-permute.wgsl.jinja"],
build/webgpu/space-depth-permute.wgsl.jinja CHANGED
@@ -1,36 +1,12 @@
1
- {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
2
- {% if note == "dispatch-limit" %}
3
- // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
- // per-axis workgroup fold width (outputs > 16.7M elements).
5
- {% elif note == "limit" %}
6
- // 2D-folded flat index: gid.y carries the high bits past the dispatch's
7
- // per-axis workgroup fold width.
8
- {% elif note == "device-axis" %}
9
- // The flat dispatch is folded across x/y at a fixed per-axis workgroup
10
- // width; gid.y carries the high portion of the output index.
11
- {% elif note == "vec4-limit" %}
12
- // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
13
- // per-axis workgroup fold width (the dispatch caps x and spills into y).
14
- {% elif note == "element-limit" %}
15
- // 2D-folded flat element index: gid.y carries the high bits past the
16
- // dispatch's per-axis workgroup fold width.
17
- {% elif note == "dispatch" %}
18
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
19
  // per-axis workgroup fold width.
20
- {% endif %}
21
- {% if bound == "" %}
22
- let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
23
- {%- elif guardInline %}
24
- let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
25
- if ({{ name }} >= {{ bound }}) { return; }
26
- {%- else %}
27
- let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
28
  if ({{ name }} >= {{ bound }}) {
29
  return;
30
- }
31
- {%- endif %}
32
- {% endmacro %}
33
-
34
  {{ env.wgsl.resourceDeclarations }}
35
 
36
  // Specializing all geometry lets division and modulo lower to constant
@@ -47,10 +23,11 @@ const BLOCK_SIZE: u32 = {{ blocksize }}u;
47
 
48
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
49
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
50
- {{ flat_index_2d("out_unit", "COUNT", note="") }}
 
51
 
52
- let ow_unit = out_unit % OUT_W_UNITS;
53
- var remaining = out_unit / OUT_W_UNITS;
54
  let oh = remaining % OUT_H;
55
  remaining /= OUT_H;
56
  let oc = remaining % OUT_C;
@@ -75,4 +52,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
75
  {% else %}
76
  output[out_unit] = value0;
77
  {% endif %}
 
78
  }
 
1
+ {% set flatVector = false %}
2
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
3
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
5
  // per-axis workgroup fold width.
6
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
 
 
 
 
 
 
 
7
  if ({{ name }} >= {{ bound }}) {
8
  return;
9
+ }{% endmacro %}
 
 
 
10
  {{ env.wgsl.resourceDeclarations }}
11
 
12
  // Specializing all geometry lets division and modulo lower to constant
 
23
 
24
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
25
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
26
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "out_unit", "COUNT") }}
27
+ {% for flatLane in range(4 if flatVector else 1) %}
28
 
29
+ let ow_unit = {{ "element" if flatVector else "out_unit" }} % OUT_W_UNITS;
30
+ var remaining = {{ "element" if flatVector else "out_unit" }} / OUT_W_UNITS;
31
  let oh = remaining % OUT_H;
32
  remaining /= OUT_H;
33
  let oc = remaining % OUT_C;
 
52
  {% else %}
53
  output[out_unit] = value0;
54
  {% endif %}
55
+ {% endfor %}
56
  }