sync c928d21e6cc1

Browse files

Files changed (3) hide show

build/webgpu/bias-add.wgsl.jinja +30 -2
build/webgpu/manifest.json +2 -1
build/webgpu/metadata.json +4 -4

build/webgpu/bias-add.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,28 @@
 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
   // 2D-folded flat index: gid.y carries the high bits past the
@@ -48,10 +73,10 @@ enable f16;
 // for one buffer would invalidate the command buffer.
 const HIDDEN: u32 = {{ hidden }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d() }}
-{% if vec4 %}
   let base = i * 4u;
   let xv = vec4<f32>(f32(x[base]), f32(x[base + 1u]), f32(x[base + 2u]), f32(x[base + 3u]));
   let sv = vec4<f32>(f32(skip[base]), f32(skip[base + 1u]), f32(skip[base + 2u]), f32(skip[base + 3u]));
@@ -59,10 +84,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   let bv = vec4<f32>(f32(bias[bcol]), f32(bias[bcol + 1u]), f32(bias[bcol + 2u]), f32(bias[bcol + 3u]));
   let v = xv + bv + sv;
   y[i] = vec4<{{ scalar }}>(v);
 {% else %}
   let xv = f32(x[i]);
   let sv = f32(skip[i]);
   let v = xv + f32(bias[i % HIDDEN]) + sv;
   y[i] = {{ scalar }}(v);
-{% endif %}
 }

+{% macro flat_tail_open() %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
+  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y).
+  let invocation = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+{% if source.itemsPerInvocation is defined %}
+  // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
+  // the logical tensor length (or its storage binding) to be vec4 aligned.
+  let begin = invocation * {{ source.itemsPerInvocation }}u;
+  let end = min(begin + {{ source.itemsPerInvocation }}u, params.count);
+  for (var i = begin; i < end; i = i + 1u) {
+{%- else %}
+  let i = invocation;
+  if (i >= params.count) {
+    return;
+  }
+{%- endif %}
+{% endmacro %}
+{% macro flat_tail_close() %}
+{% if source.itemsPerInvocation is defined %}
+  }
+{% endif %}
+{% endmacro %}
 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
   // 2D-folded flat index: gid.y carries the high bits past the
 // for one buffer would invalidate the command buffer.
 const HIDDEN: u32 = {{ hidden }}u;
+{% if vec4 %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d() }}
   let base = i * 4u;
   let xv = vec4<f32>(f32(x[base]), f32(x[base + 1u]), f32(x[base + 2u]), f32(x[base + 3u]));
   let sv = vec4<f32>(f32(skip[base]), f32(skip[base + 1u]), f32(skip[base + 2u]), f32(skip[base + 3u]));
   let bv = vec4<f32>(f32(bias[bcol]), f32(bias[bcol + 1u]), f32(bias[bcol + 2u]), f32(bias[bcol + 3u]));
   let v = xv + bv + sv;
   y[i] = vec4<{{ scalar }}>(v);
+}
 {% else %}
+{{ flat_tail_open() }}
   let xv = f32(x[i]);
   let sv = f32(skip[i]);
   let v = xv + f32(bias[i % HIDDEN]) + sv;
   y[i] = {{ scalar }}(v);
+{{ flat_tail_close() -}}
 }
+{% endif %}

build/webgpu/manifest.json CHANGED Viewed

@@ -110,6 +110,7 @@
           "id": "main",
           "name": "BiasAdd.scalar",
           "shader": "bias-add.wgsl.jinja",
           "bindings": [
             {
               "name": "x",
@@ -140,7 +141,7 @@
               "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
             }
           ],
-          "dispatch": { "threads": "numel(shapes.X)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

           "id": "main",
           "name": "BiasAdd.scalar",
           "shader": "bias-add.wgsl.jinja",
+          "source": { "inputs": { "itemsPerInvocation": 4 } },
           "bindings": [
             {
               "name": "x",
               "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
             }
           ],
+          "dispatch": { "threads": "ceilDiv(numel(shapes.X), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.BiasAdd",
-  "id": "_com_microsoft_biasadd_webgpu_e97b808",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,11 +8,11 @@
     "algorithm": "sha256",
     "files": {
       "bench.json": "W2cQSIIm7hI8MbHD6Cro9DGB/ByB035ELwNigWwGHYw=",
-      "bias-add.wgsl.jinja": "VJnmyJ3GrZrzADYN5uP73XMs4i+YXvC39ECRQ9j40Ow=",
-      "manifest.json": "pAy9IP3Uvl8QAonrSEw9nzg8kJxnbXeRZXoQql+sxe4=",
       "test.json": "PLZQbRu0fhmrRSakHOzPjTiZL0LYBLFxmBdw7gd31m4="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.BiasAdd" }
 }

 {
   "name": "com.microsoft.BiasAdd",
+  "id": "_com_microsoft_biasadd_webgpu_40b1ba1",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "W2cQSIIm7hI8MbHD6Cro9DGB/ByB035ELwNigWwGHYw=",
+      "bias-add.wgsl.jinja": "bZPMYyEnjQ+7kuft3Vjj/ZfXZq2hV5QcMREEyyiGt8w=",
+      "manifest.json": "0KqWQgIXxxbypIfz2fQy91LPsxiEoov4j8V52QlixsE=",
       "test.json": "PLZQbRu0fhmrRSakHOzPjTiZL0LYBLFxmBdw7gd31m4="
     }
   },
+  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.BiasAdd" }
 }