sync 6fdf6301e2bb

Browse files

Files changed (6) hide show

README.md +2 -2
build/webgpu/lp-norm-reduce.wgsl.jinja +2 -2
build/webgpu/manifest.json +9 -18
build/webgpu/metadata.json +7 -7
build/webgpu/norm-row-stats.wgsl.jinja +5 -109
build/webgpu/test.json +1 -1

README.md CHANGED Viewed

@@ -48,7 +48,7 @@ Default values (overridable per request):
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`lp-norm-divide.wgsl.jinja`](build/webgpu/lp-norm-divide.wgsl.jinja)
 - [`lp-norm-reduce.wgsl.jinja`](build/webgpu/lp-norm-reduce.wgsl.jinja)
 - [`lp-norm-split-combine.wgsl.jinja`](build/webgpu/lp-norm-split-combine.wgsl.jinja)
@@ -57,7 +57,7 @@ Default values (overridable per request):
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`lp-norm-divide.wgsl.jinja`](build/webgpu/lp-norm-divide.wgsl.jinja)
 - [`lp-norm-reduce.wgsl.jinja`](build/webgpu/lp-norm-reduce.wgsl.jinja)
 - [`lp-norm-split-combine.wgsl.jinja`](build/webgpu/lp-norm-split-combine.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/lp-norm-reduce.wgsl.jinja CHANGED Viewed

@@ -15,9 +15,9 @@ const SPLIT: u32 = {{ split }}u;
 // tensor.
 @compute @workgroup_size(WG)
 fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>{% if segmented %},
   @builtin(workgroup_id) wid: vec3<u32>,
-{% else %},
 {% endif %}
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {

 // tensor.
 @compute @workgroup_size(WG)
 fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+{% if segmented %}
   @builtin(workgroup_id) wid: vec3<u32>,
 {% endif %}
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {

build/webgpu/manifest.json CHANGED Viewed

@@ -22,9 +22,8 @@
   },
   "when": ["ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "f16Ok(dtypes.T)"],
   "bindings": {
-    "input": { "buffer": "read-only-storage", "elementType": "$scalar" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "axisRows" },
         { "name": "dim", "type": "u32", "value": "axisDim" },
@@ -32,12 +31,11 @@
         { "name": "p", "type": "u32", "value": "attrs.p" }
       ]
     },
-    "input_2": { "name": "input", "buffer": "read-only-storage", "elementType": "$divElem" },
     "rowNorms": { "buffer": "read-only-storage", "elementType": "f32" },
-    "output": { "buffer": "storage", "elementType": "$divElem" },
-    "params_3": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "count", "type": "u32", "value": "divideCount" },
         { "name": "dim", "type": "u32", "value": "axisDim" },
@@ -80,10 +78,9 @@
           "id": "combine",
           "name": "LpNormalization.SplitCombine",
           "shader": "lp-norm-split-combine.wgsl.jinja",
-          "derive": {},
           "bindings": [
             { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
-            { "name": "rowNorms", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
@@ -102,7 +99,7 @@
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
-          "bindings": ["input_2", "rowNorms", "output", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -140,7 +137,7 @@
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
-          "bindings": ["input_2", "rowNorms", "output", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -162,8 +159,6 @@
           "derive": {
             "modeSpec": "\"lp\"",
             "vec4": true,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.input, -1)",
             "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1) / 4))",
             "p": "attrs.p",
@@ -190,8 +185,7 @@
             "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -208,8 +202,6 @@
           "derive": {
             "modeSpec": "\"lp\"",
             "vec4": false,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.input, -1)",
             "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1)))",
             "p": "attrs.p",
@@ -236,8 +228,7 @@
             "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

   },
   "when": ["ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "f16Ok(dtypes.T)"],
   "bindings": {
+    "input": { "elementType": "$scalar" },
     "params": {
       "struct": [
         { "name": "rows", "type": "u32", "value": "axisRows" },
         { "name": "dim", "type": "u32", "value": "axisDim" },
         { "name": "p", "type": "u32", "value": "attrs.p" }
       ]
     },
+    "input_axis": { "name": "input", "elementType": "$divElem" },
     "rowNorms": { "buffer": "read-only-storage", "elementType": "f32" },
+    "output": { "elementType": "$divElem" },
+    "params_axis": {
       "name": "params",
       "struct": [
         { "name": "count", "type": "u32", "value": "divideCount" },
         { "name": "dim", "type": "u32", "value": "axisDim" },
           "id": "combine",
           "name": "LpNormalization.SplitCombine",
           "shader": "lp-norm-split-combine.wgsl.jinja",
           "bindings": [
             { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "rowNorms", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
+          "bindings": ["input_axis", "rowNorms", "output", "params_axis"],
           "dispatch": {
             "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
+          "bindings": ["input_axis", "rowNorms", "output", "params_axis"],
           "dispatch": {
             "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "derive": {
             "modeSpec": "\"lp\"",
             "vec4": true,
             "hidden": "dim(shapes.input, -1)",
             "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1) / 4))",
             "p": "attrs.p",
             "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "z": 1
+          }
         }
       ]
     },
           "derive": {
             "modeSpec": "\"lp\"",
             "vec4": false,
             "hidden": "dim(shapes.input, -1)",
             "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1)))",
             "p": "attrs.p",
             "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
             "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "ai.onnx.LpNormalization",
-  "id": "_ai_onnx_lpnormalization_webgpu_7afe9b5",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -9,16 +9,16 @@
     "files": {
       "bench.json": "Zqjg8X2eQ7EWdSXd2uSW+U227BKDhM9aUvyzA2Fd75o=",
       "lp-norm-divide.wgsl.jinja": "bIMXO9qn/IyThJQr5tapg04aNYkQqYt6LK/yz9D9ZBo=",
-      "lp-norm-reduce.wgsl.jinja": "Yu7Q9BM+uAAsEfA4VC1PQR7VOxdTyn+A+mjIryfLJwc=",
       "lp-norm-split-combine.wgsl.jinja": "PzQVkT5bUQ3+l2lqn3MgA2P0MbokIMa8Q6fVT4lFq9Q=",
-      "manifest.json": "EVQo8vZnPKz7vHj4+ZebBiS7UI2rAlpagBmiNdNEp6c=",
-      "norm-row-stats.wgsl.jinja": "iwey3jfLc6FXtfY4bBqo5+YG3ALWvo0RHLn12HHeJ0Q=",
-      "test.json": "7Vll6iQ9EAVx7E3pLedW2baT6dCaQLOfgy3qgY3xzyQ="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "axis_splitk": ["lp-norm-divide.wgsl.jinja", "lp-norm-reduce.wgsl.jinja", "lp-norm-split-combine.wgsl.jinja"],
       "axis": ["lp-norm-divide.wgsl.jinja", "lp-norm-reduce.wgsl.jinja"],

 {
   "name": "ai.onnx.LpNormalization",
+  "id": "_ai_onnx_lpnormalization_webgpu_9e7a6b8",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "files": {
       "bench.json": "Zqjg8X2eQ7EWdSXd2uSW+U227BKDhM9aUvyzA2Fd75o=",
       "lp-norm-divide.wgsl.jinja": "bIMXO9qn/IyThJQr5tapg04aNYkQqYt6LK/yz9D9ZBo=",
+      "lp-norm-reduce.wgsl.jinja": "S2wrYWXN0Mbb6OXcSlaxOnMQXmQ6k0R0WgPai4nXjZA=",
       "lp-norm-split-combine.wgsl.jinja": "PzQVkT5bUQ3+l2lqn3MgA2P0MbokIMa8Q6fVT4lFq9Q=",
+      "manifest.json": "sTPEZdTEkXALooJrYq/AV8nXbfxC5uyAOG1GsngDYMM=",
+      "norm-row-stats.wgsl.jinja": "iUons1Qj/Og11tYG6yRsNsKi3sdLtXC454tF+T2b8wg=",
+      "test.json": "W80CQp2OhX4ZWC8hdozFfQcDfmcpt4ubQfRGeVk5z1I="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "axis_splitk": ["lp-norm-divide.wgsl.jinja", "lp-norm-reduce.wgsl.jinja", "lp-norm-split-combine.wgsl.jinja"],
       "axis": ["lp-norm-divide.wgsl.jinja", "lp-norm-reduce.wgsl.jinja"],

build/webgpu/norm-row-stats.wgsl.jinja CHANGED Viewed

@@ -1,15 +1,6 @@
-{% if usesF16Spec %}
-enable f16;
-{% endif %}
-{% set combineSubgroups = combineSubgroups %}
-{% set scalarIo = scalarIo if scalarIo is defined else false %}
-{% set packedBf16Embedding = packedBf16Embedding if packedBf16Embedding is defined else false %}
 {% set lpOrder = p if p is defined else 0 %}
-{% set rmsChainNorm = rmsChainNorm if rmsChainNorm is defined else false %}
-{% set hiddenPairs = hiddenPairs | default(0) %}
-{% set numRows = numRows | default(0) %}
-{% set epsilon = epsilon | default("0.0") %}
-{% set epsilon2 = epsilon2 | default("0.0") %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -34,52 +25,7 @@ const HIDDEN: u32 = {{ hidden }}u;
 {% if vec4 %}
 const HIDDEN_V: u32 = {{ hiddenVec }}u;
 {% endif %}
-{% if packedBf16Embedding %}
-const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
-const NUM_ROWS: u32 = {{ numRows }}u;
-{% endif %}
 const WG: u32 = {{ wg }}u;
-{% if rmsChainNorm %}
-const EPSILON2: f32 = {{ epsilon2 }};
-{% endif %}
-{% if packedBf16Embedding %}
-{% if vec4 %}
-fn unpack_bf16_pair(word: u32) -> vec2<f32> {
-  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
-  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
-}
-{% endif %}
-{% if not vec4 %}
-fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
-  if (source_row >= NUM_ROWS) {
-    return 0.0;
-  }
-  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
-  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
-  return bitcast<f32>(bits << 16u);
-}
-{% endif %}
-{% if vec4 %}
-fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
-  if (source_row >= NUM_ROWS) {
-    return vec4<f32>(0.0);
-  }
-  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
-  let low = unpack_bf16_pair(x[base]);
-  let high = unpack_bf16_pair(x[base + 1u]);
-  return vec4<f32>(low, high);
-}
-{% endif %}
-{% endif %}
-{% if vec4 and scalarIo %}
-fn load_vec4(index: u32) -> vec4<f32> {
-  return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
-}
-{% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
@@ -133,31 +79,16 @@ fn main(
     return;
   }
   let tid = lid.x;
-{% if packedBf16Embedding %}
-  let source_row = indices[row];
-{% if vec4 %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
-{% elif vec4 and not scalarIo %}
-  let base = row * HIDDEN_V;
-{% else %}
-  let base = row * HIDDEN;
-{% endif %}
   var acc = 0.0;
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if packedBf16Embedding %}
-    let v = embedding_vec4(source_row, i);
-    embedding_out[base + i] = v;
-{% elif scalarIo %}
-    let v = load_vec4(base + i * 4u);
-{% else %}
-    let v = vec4<f32>(x[base + i]);
-{% endif %}
 {% if lpOrder == 1 %}
     let a = abs(v);
     acc = acc + a.x + a.y + a.z + a.w;
@@ -167,12 +98,7 @@ fn main(
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
-{% if packedBf16Embedding %}
-    let v = embedding_scalar(source_row, i);
-    embedding_out[base + i] = v;
-{% else %}
     let v = f32(x[base + i]);
-{% endif %}
 {% if lpOrder == 1 %}
     acc = acc + abs(v);
 {% else %}
@@ -189,47 +115,17 @@ fn main(
   let norm = total;
 {% endif %}
-{% if rmsChainNorm %}
-  var acc2 = 0.0;
-{% endif %}
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if packedBf16Embedding %}
     let idx = base + i;
-    let v = embedding_vec4(source_row, i);
-{% elif scalarIo %}
-    let idx = base + i * 4u;
-    let v = load_vec4(idx);
-{% else %}
-    let idx = base + i;
-    let v = vec4<f32>(x[idx]);
-{% endif %}
     let normalized = select(v / vec4<f32>(norm), vec4<f32>(0.0), vec4<bool>(norm == 0.0));
     y[idx] = {{ vecType }}(normalized);
   }
-{% if rmsChainNorm %}
-  // The chained second norm reads the residual row this loop just stored. This
-  // barrier completes those stores and any preceding shared-scratch use before
-  // the next reduction reuses its scratch; each lane then re-reads only the
-  // elements it wrote itself.
-  workgroupBarrier();
-  let total2 = reduce_scalar(acc2{{ reduceThreadArguments }});
-  let inv2 = inverseSqrt(total2 / f32(HIDDEN) + EPSILON2);
-  for (var i = tid; i < HIDDEN_V; i = i + WG) {
-    let idx = base + i;
-    let hv = vec4<f32>(y[idx]);
-    normed2[idx] = {{ vecType }}(hv * inv2 * vec4<f32>(scale2[i]));
-  }
-{% endif %}
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
-{% if packedBf16Embedding %}
-    let v = embedding_scalar(source_row, i);
-{% else %}
     let v = f32(x[idx]);
-{% endif %}
     let normalized = select(v / norm, 0.0, norm == 0.0);
     y[idx] = {{ scalar }}(normalized);
   }

+{% set scalarIo = false %}
+{% set packedF32 = "vec4<f32>" %}
 {% set lpOrder = p if p is defined else 0 %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 {% if vec4 %}
 const HIDDEN_V: u32 = {{ hiddenVec }}u;
 {% endif %}
 const WG: u32 = {{ wg }}u;
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
     return;
   }
   let tid = lid.x;
+{% if vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
   var acc = 0.0;
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let v = {{ packedF32 }}(x[base + i]);
 {% if lpOrder == 1 %}
     let a = abs(v);
     acc = acc + a.x + a.y + a.z + a.w;
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
 {% if lpOrder == 1 %}
     acc = acc + abs(v);
 {% else %}
   let norm = total;
 {% endif %}
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
+    let v = {{ packedF32 }}(x[idx]);
     let normalized = select(v / vec4<f32>(norm), vec4<f32>(0.0), vec4<bool>(norm == 0.0));
     y[idx] = {{ vecType }}(normalized);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
     let normalized = select(v / norm, 0.0, norm == 0.0);
     y[idx] = {{ scalar }}(normalized);
   }

build/webgpu/test.json CHANGED Viewed

@@ -225,7 +225,7 @@
       "outputs": { "output": { "dtype": "float16", "shape": [2, 128], "tolerance": 0.00002, "relTolerance": 0.002 } }
     },
     {
-      "name": "f16_p1_last_axis_constant200_norm_scale_lock",
       "provenance": {
         "notes": "Each 128-element row contains 200, so its L1 norm is 25,600 and every normalized value is exactly 2^-7 in float16. A 1e-6 tolerance makes an incorrect divisor observable."
       },

       "outputs": { "output": { "dtype": "float16", "shape": [2, 128], "tolerance": 0.00002, "relTolerance": 0.002 } }
     },
     {
+      "name": "f16_p1_last_axis_constant200_norm_scale",
       "provenance": {
         "notes": "Each 128-element row contains 200, so its L1 norm is 25,600 and every normalized value is exactly 2^-7 in float16. A 1e-6 tolerance makes an incorrect divisor observable."
       },