Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Transpose

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
f5ec031
·
verified ·
1 Parent(s): 30e6997

sync c928d21e6cc1

Browse files
Files changed (3) hide show
  1. build/webgpu/manifest.json +3 -2
  2. build/webgpu/metadata.json +4 -4
  3. build/webgpu/transpose.wgsl.jinja +18 -8
build/webgpu/manifest.json CHANGED
@@ -228,11 +228,12 @@
228
  "inputShape": "shapes.data",
229
  "outputShape": "shapes.transposed",
230
  "rank": "ranks.data",
231
- "perm": "genericPerm"
 
232
  }
233
  },
234
  "bindings": "scalarData",
235
- "dispatch": { "threads": "numel(shapes.transposed)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
236
  }
237
  ]
238
  }
 
228
  "inputShape": "shapes.data",
229
  "outputShape": "shapes.transposed",
230
  "rank": "ranks.data",
231
+ "perm": "genericPerm",
232
+ "itemsPerInvocation": 4
233
  }
234
  },
235
  "bindings": "scalarData",
236
+ "dispatch": { "threads": "ceilDiv(numel(shapes.transposed), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
237
  }
238
  ]
239
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Transpose",
3
- "id": "_ai_onnx_transpose_webgpu_dff5d66",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -12,11 +12,11 @@
12
  "datamove-transpose-2d-tiled-scalar.wgsl.jinja": "zy7rlN8Se+mn8RB9akbBTHq1RK8ot6NBboxG3vX1Xdg=",
13
  "datamove-transpose-2d-tiled.wgsl.jinja": "AtmY7NlBNCV1unIf2epvpBO3hfnLBmgGKGB7WiWMCQk=",
14
  "datamove-transpose-vec4.wgsl.jinja": "RgwjTbu7SHDVsUW7W2f5S81gv1UJBpT93udQX/MZdJ0=",
15
- "manifest.json": "qACTV3+3jqKTescFk98we7Nm2pwajZoBwGAo1n3dL9M=",
16
  "test.json": "acNlC/DcnpoXb3ni/Z4QMbpYD7C1YR16iHXX20mNp0Y=",
17
- "transpose.wgsl.jinja": "BbRpj9r38KShHUF3RXf518/bYtZ8BHFbGDZLRbOuP9g="
18
  }
19
  },
20
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
21
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Transpose" }
22
  }
 
1
  {
2
  "name": "ai.onnx.Transpose",
3
+ "id": "_ai_onnx_transpose_webgpu_713616c",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
12
  "datamove-transpose-2d-tiled-scalar.wgsl.jinja": "zy7rlN8Se+mn8RB9akbBTHq1RK8ot6NBboxG3vX1Xdg=",
13
  "datamove-transpose-2d-tiled.wgsl.jinja": "AtmY7NlBNCV1unIf2epvpBO3hfnLBmgGKGB7WiWMCQk=",
14
  "datamove-transpose-vec4.wgsl.jinja": "RgwjTbu7SHDVsUW7W2f5S81gv1UJBpT93udQX/MZdJ0=",
15
+ "manifest.json": "Av+HRbQxG96oiNzrlouEobs2tS/VmfX0tCkNdy8rWL4=",
16
  "test.json": "acNlC/DcnpoXb3ni/Z4QMbpYD7C1YR16iHXX20mNp0Y=",
17
+ "transpose.wgsl.jinja": "lYYw/Xqp/wTaYlWr8iB/vn1ZxacE3jOes9gYViMlhlo="
18
  }
19
  },
20
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
21
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Transpose" }
22
  }
build/webgpu/transpose.wgsl.jinja CHANGED
@@ -1,3 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  {% if usesF16 %}
2
  enable f16;
3
  {% endif %}
@@ -26,13 +42,7 @@ fn input_offset({% if source.rank > 0 %}out_index: u32{% endif %}) -> u32 {
26
  return offset;
27
  }
28
 
29
- @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
30
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
31
- // 2D-folded flat index: gid.y carries the high bits past the
32
- // maxComputeWorkgroupsPerDimension limit.
33
- let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
34
- if (i >= params.count) {
35
- return;
36
- }
37
  y[i] = x[input_offset({% if source.rank > 0 %}i{% endif %})];
 
38
  }
 
1
+ {% macro flat_tail_open() %}
2
+ @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
3
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
4
+ // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
5
+ // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y).
6
+ let invocation = gid.x + gid.y * nwg.x * {{ 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
+ let begin = invocation * {{ source.itemsPerInvocation }}u;
10
+ let end = min(begin + {{ source.itemsPerInvocation }}u, params.count);
11
+ for (var i = begin; i < end; i = i + 1u) {
12
+ {%- endmacro %}
13
+ {% macro flat_tail_close() %}
14
+ }
15
+ {% endmacro %}
16
+
17
  {% if usesF16 %}
18
  enable f16;
19
  {% endif %}
 
42
  return offset;
43
  }
44
 
45
+ {{ flat_tail_open() }}
 
 
 
 
 
 
 
46
  y[i] = x[input_offset({% if source.rank > 0 %}i{% endif %})];
47
+ {{ flat_tail_close() -}}
48
  }