Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Concat

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
5d188ac
·
verified ·
1 Parent(s): ee861ee

sync c928d21e6cc1

Browse files
Files changed (3) hide show
  1. build/webgpu/concat-n.wgsl.jinja +17 -17
  2. build/webgpu/manifest.json +12 -12
  3. build/webgpu/metadata.json +4 -4
build/webgpu/concat-n.wgsl.jinja CHANGED
@@ -60,18 +60,17 @@ fn axis_coord({% if c_count.value > 0 %}out_index: u32{% endif %}) -> u32 {
60
 
61
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
62
  fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
63
- // 2D-folded flat index: gid.y carries the high bits past the
64
- // maxComputeWorkgroupsPerDimension limit. Named out_idx (not i) so it never
65
- // collides with an input bound to the WGSL name `i`.
66
- let out_idx = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
67
- if (out_idx >= {{ c_count.value }}u) {
68
- return;
69
- }
70
- let coord = axis_coord({% if c_count.value > 0 %}out_idx{% endif %});
71
  {% if axisBase > 0 %}
72
- if (coord < {{ axisBase }}u) {
73
- return;
74
- }
75
  {% endif %}
76
  {% set cumulative = namespace(value=axisBase) %}
77
  {% for input in source.inputs %}
@@ -79,15 +78,16 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
79
  {% set cumulative.value = cumulative.value + input.shape[source.axis] %}
80
  {% set bounded = (not loop.last) or partialTail %}
81
  {% if lower > 0 and bounded %}
82
- if (coord >= {{ lower }}u && coord < {{ cumulative.value }}u) {
83
  {% elif lower > 0 %}
84
- if (coord >= {{ lower }}u) {
85
  {% elif bounded %}
86
- if (coord < {{ cumulative.value }}u) {
87
  {% else %}
88
- {
89
  {% endif %}
90
- c[out_idx] = {{ input.name }}[{{ input.name }}_offset({% if c_count.value > 0 %}out_idx{% endif %})];
91
- }
92
  {% endfor %}
 
93
  }
 
60
 
61
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
62
  fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
63
+ // Four consecutive outputs per invocation. Named out_idx (not i) so it never
64
+ // collides with an input bound to the WGSL name `i`; gid.y carries the high
65
+ // bits of the folded dispatch past the maxComputeWorkgroupsPerDimension limit.
66
+ let begin = (gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u) * 4u;
67
+ let end = min(begin + 4u, {{ c_count.value }}u);
68
+ for (var out_idx = begin; out_idx < end; out_idx = out_idx + 1u) {
69
+ let coord = axis_coord({% if c_count.value > 0 %}out_idx{% endif %});
 
70
  {% if axisBase > 0 %}
71
+ if (coord < {{ axisBase }}u) {
72
+ continue;
73
+ }
74
  {% endif %}
75
  {% set cumulative = namespace(value=axisBase) %}
76
  {% for input in source.inputs %}
 
78
  {% set cumulative.value = cumulative.value + input.shape[source.axis] %}
79
  {% set bounded = (not loop.last) or partialTail %}
80
  {% if lower > 0 and bounded %}
81
+ if (coord >= {{ lower }}u && coord < {{ cumulative.value }}u) {
82
  {% elif lower > 0 %}
83
+ if (coord >= {{ lower }}u) {
84
  {% elif bounded %}
85
+ if (coord < {{ cumulative.value }}u) {
86
  {% else %}
87
+ {
88
  {% endif %}
89
+ c[out_idx] = {{ input.name }}[{{ input.name }}_offset({% if c_count.value > 0 %}out_idx{% endif %})];
90
+ }
91
  {% endfor %}
92
+ }
93
  }
build/webgpu/manifest.json CHANGED
@@ -317,7 +317,7 @@
317
  }
318
  },
319
  "bindings": "fiveInputs",
320
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
321
  },
322
  {
323
  "id": "tail",
@@ -333,7 +333,7 @@
333
  }
334
  },
335
  "bindings": "tailThreeInputs",
336
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
337
  }
338
  ]
339
  },
@@ -357,7 +357,7 @@
357
  }
358
  },
359
  "bindings": "fiveInputs",
360
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
361
  },
362
  {
363
  "id": "tail",
@@ -373,7 +373,7 @@
373
  }
374
  },
375
  "bindings": "tailFourInputs",
376
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
377
  }
378
  ]
379
  },
@@ -548,7 +548,7 @@
548
  }
549
  },
550
  "bindings": "twoInputs",
551
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
552
  }
553
  ]
554
  },
@@ -571,7 +571,7 @@
571
  }
572
  },
573
  "bindings": "fourInputs",
574
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
575
  }
576
  ]
577
  },
@@ -594,7 +594,7 @@
594
  }
595
  },
596
  "bindings": "fiveInputs",
597
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
598
  }
599
  ]
600
  },
@@ -617,7 +617,7 @@
617
  }
618
  },
619
  "bindings": "sixInputs",
620
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
621
  }
622
  ]
623
  },
@@ -640,7 +640,7 @@
640
  }
641
  },
642
  "bindings": "sevenInputs",
643
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
644
  }
645
  ]
646
  },
@@ -663,7 +663,7 @@
663
  }
664
  },
665
  "bindings": "eightInputs",
666
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
667
  }
668
  ]
669
  },
@@ -710,7 +710,7 @@
710
  }
711
  },
712
  "bindings": "nineInputs",
713
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
714
  }
715
  ]
716
  },
@@ -733,7 +733,7 @@
733
  }
734
  },
735
  "bindings": "threeInputs",
736
- "dispatch": { "threads": "numel(shapes.c)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
737
  }
738
  ]
739
  }
 
317
  }
318
  },
319
  "bindings": "fiveInputs",
320
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
321
  },
322
  {
323
  "id": "tail",
 
333
  }
334
  },
335
  "bindings": "tailThreeInputs",
336
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
337
  }
338
  ]
339
  },
 
357
  }
358
  },
359
  "bindings": "fiveInputs",
360
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
361
  },
362
  {
363
  "id": "tail",
 
373
  }
374
  },
375
  "bindings": "tailFourInputs",
376
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
377
  }
378
  ]
379
  },
 
548
  }
549
  },
550
  "bindings": "twoInputs",
551
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
552
  }
553
  ]
554
  },
 
571
  }
572
  },
573
  "bindings": "fourInputs",
574
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
575
  }
576
  ]
577
  },
 
594
  }
595
  },
596
  "bindings": "fiveInputs",
597
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
598
  }
599
  ]
600
  },
 
617
  }
618
  },
619
  "bindings": "sixInputs",
620
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
621
  }
622
  ]
623
  },
 
640
  }
641
  },
642
  "bindings": "sevenInputs",
643
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
644
  }
645
  ]
646
  },
 
663
  }
664
  },
665
  "bindings": "eightInputs",
666
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
667
  }
668
  ]
669
  },
 
710
  }
711
  },
712
  "bindings": "nineInputs",
713
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
714
  }
715
  ]
716
  },
 
733
  }
734
  },
735
  "bindings": "threeInputs",
736
+ "dispatch": { "threads": "ceilDiv(numel(shapes.c), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
737
  }
738
  ]
739
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Concat",
3
- "id": "_ai_onnx_concat_webgpu_64c71dd",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -9,14 +9,14 @@
9
  "files": {
10
  "bench.json": "4vi/7xLqrnj6XrQmLczlokRTUtwhNt7BTJd+4EF7Qts=",
11
  "concat-last-axis-vec4-tail.wgsl.jinja": "N4rfqKJeeGVJvZkNTRN9Iki76pk2fDo0ehpbkpNvS3Q=",
12
- "concat-n.wgsl.jinja": "jBOHWgvT2OB9RVHPT7p9ctbKdFoG46G4ggJ9kP0UHSM=",
13
  "datamove-concat-block.wgsl.jinja": "Kg3TKvB5kMzYzIp7dliwq2bAIzHRmsqUpBZ6F1VGKrc=",
14
  "datamove-concat-last-axis-cooperative.wgsl.jinja": "28h6uIprJtOXG69eZ5v27WC3BfDRweQMsXoFRprsHYg=",
15
  "datamove-flat-copy.wgsl.jinja": "rEpeYMvIZQuQnB2f3lt51jrrklKICsqUTzCJPRTknNk=",
16
- "manifest.json": "AM5zQkbwoNHCmTWhNQYD7Ny/0ba5PmEi9U4r4GVJHcM=",
17
  "test.json": "7iHOtLpxgTfHaKsafULD0XqApFBxgEqIlblmZDjpTq0="
18
  }
19
  },
20
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
21
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Concat" }
22
  }
 
1
  {
2
  "name": "ai.onnx.Concat",
3
+ "id": "_ai_onnx_concat_webgpu_e443e2a",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
9
  "files": {
10
  "bench.json": "4vi/7xLqrnj6XrQmLczlokRTUtwhNt7BTJd+4EF7Qts=",
11
  "concat-last-axis-vec4-tail.wgsl.jinja": "N4rfqKJeeGVJvZkNTRN9Iki76pk2fDo0ehpbkpNvS3Q=",
12
+ "concat-n.wgsl.jinja": "OHkvG+IVeZTKsQAjZnNUvLQZRrEedHMZ2/2KCWwKuz8=",
13
  "datamove-concat-block.wgsl.jinja": "Kg3TKvB5kMzYzIp7dliwq2bAIzHRmsqUpBZ6F1VGKrc=",
14
  "datamove-concat-last-axis-cooperative.wgsl.jinja": "28h6uIprJtOXG69eZ5v27WC3BfDRweQMsXoFRprsHYg=",
15
  "datamove-flat-copy.wgsl.jinja": "rEpeYMvIZQuQnB2f3lt51jrrklKICsqUTzCJPRTknNk=",
16
+ "manifest.json": "NLe1nI97qznOuYRz69A7PWN8NFwzVOPoVurULnJuIXw=",
17
  "test.json": "7iHOtLpxgTfHaKsafULD0XqApFBxgEqIlblmZDjpTq0="
18
  }
19
  },
20
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
21
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Concat" }
22
  }