Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.GatedAdd

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
d27ea5a
·
verified ·
1 Parent(s): b5d19c4

sync c928d21e6cc1

Browse files
Files changed (3) hide show
  1. build/webgpu/gated-add.wgsl.jinja +30 -2
  2. build/webgpu/manifest.json +2 -1
  3. build/webgpu/metadata.json +4 -4
build/webgpu/gated-add.wgsl.jinja CHANGED
@@ -1,3 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
@@ -48,16 +73,19 @@ enable f16;
48
  // make this op more accurate than the graph it replaces.
49
  const HIDDEN: u32 = {{ hidden }}u;
50
 
 
51
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
52
  fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
53
  {{ flat_index_2d() }}
54
- {% if vec4 %}
55
  // A vec4 group is four consecutive channels of one row: HIDDEN % 4 == 0 stops
56
  // it from ever straddling two rows, so the whole group shares one gate value.
57
  let g = vec4<{{ scalar }}>(gate[i * 4u / HIDDEN]);
58
  output[i] = x[i] + fma(y[i], g, vec4<{{ scalar }}>(0.0));
 
59
  {% else %}
 
60
  let g = gate[i / HIDDEN];
61
  output[i] = x[i] + fma(y[i], g, {{ scalar }}(0.0));
62
- {% endif %}
63
  }
 
 
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
+ {% if source.itemsPerInvocation is defined %}
8
+ // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
9
+ // the logical tensor length (or its storage binding) to be vec4 aligned.
10
+ let begin = invocation * {{ source.itemsPerInvocation }}u;
11
+ let end = min(begin + {{ source.itemsPerInvocation }}u, params.count);
12
+ for (var i = begin; i < end; i = i + 1u) {
13
+ {%- else %}
14
+ let i = invocation;
15
+ if (i >= params.count) {
16
+ return;
17
+ }
18
+ {%- endif %}
19
+ {% endmacro %}
20
+ {% macro flat_tail_close() %}
21
+ {% if source.itemsPerInvocation is defined %}
22
+ }
23
+ {% endif %}
24
+ {% endmacro %}
25
+
26
  {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
27
  {% if note == "dispatch-limit" %}
28
  // 2D-folded flat index: gid.y carries the high bits past the
 
73
  // make this op more accurate than the graph it replaces.
74
  const HIDDEN: u32 = {{ hidden }}u;
75
 
76
+ {% if vec4 %}
77
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
78
  fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
79
  {{ flat_index_2d() }}
 
80
  // A vec4 group is four consecutive channels of one row: HIDDEN % 4 == 0 stops
81
  // it from ever straddling two rows, so the whole group shares one gate value.
82
  let g = vec4<{{ scalar }}>(gate[i * 4u / HIDDEN]);
83
  output[i] = x[i] + fma(y[i], g, vec4<{{ scalar }}>(0.0));
84
+ }
85
  {% else %}
86
+ {{ flat_tail_open() }}
87
  let g = gate[i / HIDDEN];
88
  output[i] = x[i] + fma(y[i], g, {{ scalar }}(0.0));
89
+ {{ flat_tail_close() -}}
90
  }
91
+ {% endif %}
build/webgpu/manifest.json CHANGED
@@ -103,6 +103,7 @@
103
  "id": "main",
104
  "name": "GatedAdd.scalar",
105
  "shader": "gated-add.wgsl.jinja",
 
106
  "bindings": [
107
  {
108
  "name": "x",
@@ -139,7 +140,7 @@
139
  "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
140
  }
141
  ],
142
- "dispatch": { "threads": "numel(shapes.X)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
143
  }
144
  ]
145
  }
 
103
  "id": "main",
104
  "name": "GatedAdd.scalar",
105
  "shader": "gated-add.wgsl.jinja",
106
+ "source": { "inputs": { "itemsPerInvocation": 4 } },
107
  "bindings": [
108
  {
109
  "name": "x",
 
140
  "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
141
  }
142
  ],
143
+ "dispatch": { "threads": "ceilDiv(numel(shapes.X), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
144
  }
145
  ]
146
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "com.microsoft.GatedAdd",
3
- "id": "_com_microsoft_gatedadd_webgpu_c9e5b9d",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,11 +8,11 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "Yo5szY3Ccb8NOxk0YLKXGApxHMFCENocZt9z2/GaJew=",
11
- "gated-add.wgsl.jinja": "GiWIIFkb2GXB3/MDWr4Buyc/kKQrH5xr2tcee1RCOog=",
12
- "manifest.json": "eGSFQTnmYL4/Zqwfz8zkKFJgFT1Fn5IjP+slAJPCuTU=",
13
  "test.json": "5Gsy9SVd5u5t8bVgKOLTcNMfO4wqI7rFSjHIhPpLEgk="
14
  }
15
  },
16
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
17
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GatedAdd" }
18
  }
 
1
  {
2
  "name": "com.microsoft.GatedAdd",
3
+ "id": "_com_microsoft_gatedadd_webgpu_bd62a87",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "Yo5szY3Ccb8NOxk0YLKXGApxHMFCENocZt9z2/GaJew=",
11
+ "gated-add.wgsl.jinja": "1jbbfo5jhq8VJQHFSPM20gaFGAman4LWBgrssi9fQ5E=",
12
+ "manifest.json": "t9s9nEQJMkL8JEPMCsIRRZhZlohyMFsAkyIW7GIwTRw=",
13
  "test.json": "5Gsy9SVd5u5t8bVgKOLTcNMfO4wqI7rFSjHIhPpLEgk="
14
  }
15
  },
16
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
17
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GatedAdd" }
18
  }