sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +16 -13
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +62 -196
build/webgpu/metadata.json +16 -9
build/webgpu/one-hot-fill.wgsl.jinja +16 -19
build/webgpu/one-hot-last-axis-vec4.wgsl.jinja +22 -26
build/webgpu/one-hot-scatter.wgsl.jinja +17 -20
build/webgpu/test.json +10 -11

README.md CHANGED Viewed

@@ -18,17 +18,17 @@ See the [ONNX `OneHot` spec](https://onnx.ai/onnx/operators/onnx__OneHot.html) f
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `indices` | `indices` | `I` | — | — | Integer or float index tensor; values outside `[-depth, depth-1]` produce all-`off_value` output rows. | required |
-| `depth` | `depth` | `D` | — | — | Scalar (or length-1 rank-1) tensor specifying the number of classes and the size of the one-hot dimension. | required |
-| `values` | `values` | `T` | `1` | — | Rank-1 tensor of exactly two elements `[off_value, on_value]` giving the values written to inactive and active positions respectively. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `output` | `output` | `T` | derived | — | One-hot tensor with rank equal to `rank(indices) + 1`, same element type as `values`. | required |
 ## Attributes
@@ -48,7 +48,7 @@ Default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -58,15 +58,18 @@ Default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
-The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
-- `output`
-Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `indices` | `I` | — | — | Integer or float index tensor; values outside `[-depth, depth-1]` produce all-`off_value` output rows. | required |
+| `depth` | `D` | — | — | Scalar (or length-1 rank-1) tensor specifying the number of classes and the size of the one-hot dimension. | required |
+| `values` | `T` | `1` | — | Rank-1 tensor of exactly two elements `[off_value, on_value]` giving the values written to inactive and active positions respectively. | required |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | derived | — | One-hot tensor with rank equal to `rank(indices) + 1`, same element type as `values`. | required |
 ## Attributes
 ## Files
+- [`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
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
+This example supplies explicit metadata for:
+- `output`
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.OneHot",
   "cases": [
     {
       "name": "tokens_vocab",

 {
   "cases": [
     {
       "name": "tokens_vocab",

build/webgpu/manifest.json CHANGED Viewed

@@ -2,65 +2,32 @@
   "domain": "ai.onnx",
   "name": "OneHot",
   "sinceVersion": 11,
-  "description": "Produces a one-hot tensor from an `indices` input: positions matching each index are filled with `on_value` and all other positions with `off_value`, where both are taken from the two-element `values` tensor `[off_value, on_value]`. The output rank is one greater than `indices`, with the new dimension of size `depth` inserted at the position given by `axis`; indices outside `[-depth, depth-1]` yield all-`off_value` rows.",
-  "inputs": [
-    {
-      "role": "indices",
-      "dtype": "I",
-      "description": "Integer or float index tensor; values outside `[-depth, depth-1]` produce all-`off_value` output rows."
-    },
-    {
-      "role": "depth",
-      "dtype": "D",
-      "description": "Scalar (or length-1 rank-1) tensor specifying the number of classes and the size of the one-hot dimension."
-    },
-    {
-      "role": "values",
-      "dtype": "T",
-      "description": "Rank-1 tensor of exactly two elements `[off_value, on_value]` giving the values written to inactive and active positions respectively.",
-      "rank": 1
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "description": "One-hot tensor with rank equal to `rank(indices) + 1`, same element type as `values`.",
-      "rank": "ranks.indices + 1"
-    }
-  ],
-  "attributes": { "axis": -1 },
-  "attributeDescriptions": {
-    "axis": "Axis along which the one-hot dimension is inserted; default `-1` appends it as the last dimension. Negative values count from the back; accepted range is `[-r-1, r]` where `r = rank(indices)`."
-  },
   "typeConstraints": {
     "I": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8"],
     "D": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8"],
     "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"]
   },
-  "args": {
-    "indices": { "kind": "tensor", "semantic": "indices", "role": "indices" },
-    "depth": { "kind": "tensor", "semantic": "depth", "role": "depth" },
-    "values": { "kind": "tensor", "semantic": "values", "role": "values" },
-    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256 },
   "derive": {
     "oneHotContractOk": "(dtypes.I == \"u32\" or dtypes.I == \"i32\" or dtypes.I == \"f32\" or dtypes.I == \"f16\") and ranks.indices >= 1 and ranks.output == ranks.indices + 1 and (ranks.depth == 0 or (ranks.depth == 1 and dim(shapes.depth, 0) == 1)) and ranks.values == 1 and dim(shapes.values, 0) == 2 and numel(shapes.output) >= numel(shapes.indices) and f16Ok(dtypes.T) and f16Ok(dtypes.I)",
     "oneHotLastAxisCovered": "oneHotContractOk and (attrs.axis == ranks.output - 1 or attrs.axis == -1) and dim(shapes.output, -1) % 4 == 0",
     "oneHotLastAxisVecCount": "dim(shapes.output, -1) / 4",
     "oneHotRowWorkgroup": "min(tunables.WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, max(1, pow2ceil(oneHotLastAxisVecCount)))",
     "oneHotCooperativeMinVecs": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 32",
-    "oneHotRowCooperative": "oneHotLastAxisCovered and oneHotLastAxisVecCount >= oneHotCooperativeMinVecs"
-  },
-  "constants": {
     "scalar": "dtypes.T",
     "indexScalar": "dtypes.I",
     "usesI32": "dtypes.I == \"i32\"",
     "floatIndices": "dtypes.I == \"f32\" or dtypes.I == \"f16\"",
-    "usesF16": "dtypes.T == \"f16\" or dtypes.I == \"f16\"",
     "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
   },
   "variants": [
     {
       "id": "last_axis_row_vec4",
@@ -70,40 +37,15 @@
         {
           "id": "main",
           "name": "OneHot.last_axis_row_vec4",
-          "source": {
-            "shader": "one-hot-last-axis-vec4.wgsl.jinja",
-            "inputs": {
-              "depth": "dim(shapes.output, -1)",
-              "rows": "numel(shapes.indices)",
-              "wg": "oneHotRowWorkgroup",
-              "cooperative": true
-            }
           },
-          "bindings": [
-            {
-              "name": "indices",
-              "arg": "indices",
-              "semantic": "indices",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$indexScalar"
-            },
-            {
-              "name": "values",
-              "arg": "values",
-              "semantic": "values",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar",
-              "length": 2
-            },
-            {
-              "name": "output",
-              "arg": "output",
-              "semantic": "output",
-              "buffer": { "type": "storage" },
-              "elementType": "$vectorScalar"
-            }
-          ],
-          "dispatch": { "workgroups": "numel(shapes.indices)" }
         }
       ]
     },
@@ -115,44 +57,19 @@
         {
           "id": "main",
           "name": "OneHot.last_axis_vec4",
-          "source": {
-            "shader": "one-hot-last-axis-vec4.wgsl.jinja",
-            "inputs": { "depth": "dim(shapes.output, -1)", "cooperative": false }
-          },
           "bindings": [
-            {
-              "name": "indices",
-              "arg": "indices",
-              "semantic": "indices",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$indexScalar"
-            },
-            {
-              "name": "values",
-              "arg": "values",
-              "semantic": "values",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar",
-              "length": 2
-            },
-            {
-              "name": "output",
-              "arg": "output",
-              "semantic": "output",
-              "buffer": { "type": "storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" }]
-              }
-            }
           ],
-          "dispatch": { "threads": "numel(shapes.output) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
@@ -169,108 +86,57 @@
         {
           "id": "fill_vec4",
           "name": "OneHot.FillVec4",
-          "source": { "shader": "one-hot-fill.wgsl.jinja", "inputs": { "vectorized": true } },
           "bindings": [
-            {
-              "name": "values",
-              "arg": "values",
-              "semantic": "values",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar",
-              "length": 2
-            },
-            {
-              "name": "output",
-              "arg": "output",
-              "semantic": "output",
-              "buffer": { "type": "storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "outputVecCount" }] }
-            }
           ],
-          "dispatch": { "threads": "outputVecCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         },
         {
           "id": "fill_tail",
           "name": "OneHot.FillTail",
-          "source": { "shader": "one-hot-fill.wgsl.jinja", "inputs": { "vectorized": false } },
           "bindings": [
-            {
-              "name": "values",
-              "arg": "values",
-              "semantic": "values",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar",
-              "length": 2
-            },
-            {
-              "name": "output",
-              "arg": "output",
-              "semantic": "output",
-              "buffer": { "type": "storage" },
-              "elementType": "$scalar"
-            },
             {
               "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [
-                  { "name": "count", "type": "u32", "value": "outputTailCount" },
-                  { "name": "offset", "type": "u32", "value": "outputVecCount * 4" }
-                ]
-              }
             }
           ],
-          "dispatch": { "threads": "outputTailCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         },
         {
           "id": "scatter",
           "name": "OneHot.Scatter",
-          "source": {
-            "shader": "one-hot-scatter.wgsl.jinja",
-            "inputs": { "depth": "dim(shapes.output, depthAxis)", "inner": "depthInner" }
-          },
           "bindings": [
-            {
-              "name": "indices",
-              "arg": "indices",
-              "semantic": "indices",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$indexScalar"
-            },
-            {
-              "name": "values",
-              "arg": "values",
-              "semantic": "values",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar",
-              "length": 2
-            },
-            {
-              "name": "output",
-              "arg": "output",
-              "semantic": "output",
-              "buffer": { "type": "storage" },
-              "elementType": "$scalar"
-            },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.indices)" }]
-              }
-            }
           ],
-          "dispatch": { "threads": "numel(shapes.indices)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "OneHot",
   "sinceVersion": 11,
+  "inputs": { "indices": { "dtype": "I" }, "depth": { "dtype": "D" }, "values": { "dtype": "T", "rank": 1 } },
+  "outputs": { "output": { "dtype": "T", "rank": "ranks.indices + 1" } },
+  "attributes": { "axis": { "default": -1 } },
   "typeConstraints": {
     "I": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8"],
     "D": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8"],
     "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"]
   },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
     "oneHotContractOk": "(dtypes.I == \"u32\" or dtypes.I == \"i32\" or dtypes.I == \"f32\" or dtypes.I == \"f16\") and ranks.indices >= 1 and ranks.output == ranks.indices + 1 and (ranks.depth == 0 or (ranks.depth == 1 and dim(shapes.depth, 0) == 1)) and ranks.values == 1 and dim(shapes.values, 0) == 2 and numel(shapes.output) >= numel(shapes.indices) and f16Ok(dtypes.T) and f16Ok(dtypes.I)",
     "oneHotLastAxisCovered": "oneHotContractOk and (attrs.axis == ranks.output - 1 or attrs.axis == -1) and dim(shapes.output, -1) % 4 == 0",
     "oneHotLastAxisVecCount": "dim(shapes.output, -1) / 4",
     "oneHotRowWorkgroup": "min(tunables.WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, max(1, pow2ceil(oneHotLastAxisVecCount)))",
     "oneHotCooperativeMinVecs": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 32",
+    "oneHotRowCooperative": "oneHotLastAxisCovered and oneHotLastAxisVecCount >= oneHotCooperativeMinVecs",
     "scalar": "dtypes.T",
     "indexScalar": "dtypes.I",
     "usesI32": "dtypes.I == \"i32\"",
     "floatIndices": "dtypes.I == \"f32\" or dtypes.I == \"f16\"",
     "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
   },
+  "bindings": {
+    "values": { "buffer": "read-only-storage", "elementType": "$scalar", "length": 2 },
+    "output": { "buffer": "storage", "elementType": "$vectorScalar" }
+  },
   "variants": [
     {
       "id": "last_axis_row_vec4",
         {
           "id": "main",
           "name": "OneHot.last_axis_row_vec4",
+          "shader": "one-hot-last-axis-vec4.wgsl.jinja",
+          "derive": {
+            "depth": "dim(shapes.output, -1)",
+            "rowCount": "numel(shapes.indices)",
+            "wg": "oneHotRowWorkgroup",
+            "cooperative": true
           },
+          "bindings": ["indices", "values", "output"],
+          "dispatch": { "x": "min(numel(shapes.indices), 65535)", "y": "ceilDiv(numel(shapes.indices), 65535)", "z": 1 }
         }
       ]
     },
         {
           "id": "main",
           "name": "OneHot.last_axis_vec4",
+          "shader": "one-hot-last-axis-vec4.wgsl.jinja",
+          "derive": { "depth": "dim(shapes.output, -1)", "cooperative": false },
           "bindings": [
+            "indices",
+            "values",
+            "output",
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" }] }
           ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
         {
           "id": "fill_vec4",
           "name": "OneHot.FillVec4",
+          "shader": "one-hot-fill.wgsl.jinja",
+          "derive": { "vectorizedSpec": true },
           "bindings": [
+            "values",
+            "output",
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "outputVecCount" }] }
           ],
+          "dispatch": {
+            "x": "min(ceilDiv((outputVecCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputVecCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "fill_tail",
           "name": "OneHot.FillTail",
+          "shader": "one-hot-fill.wgsl.jinja",
+          "derive": { "vectorizedSpec": false },
           "bindings": [
+            "values",
+            { "arg": "output" },
             {
               "name": "params",
+              "struct": [
+                { "name": "count", "type": "u32", "value": "outputTailCount" },
+                { "name": "offset", "type": "u32", "value": "outputVecCount * 4" }
+              ]
             }
           ],
+          "dispatch": {
+            "x": "min(ceilDiv((outputTailCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputTailCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "scatter",
           "name": "OneHot.Scatter",
+          "shader": "one-hot-scatter.wgsl.jinja",
+          "derive": { "depth": "dim(shapes.output, depthAxis)", "innerSize": "depthInner" },
           "bindings": [
+            "indices",
+            "values",
+            { "arg": "output" },
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.indices)" }] }
           ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.indices)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.indices)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,27 @@
 {
   "name": "ai.onnx.OneHot",
-  "id": "_ai_onnx_onehot_webgpu_439ff74",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "tIK5K4R1/SOMH56dBhsH6mCf7IQ0DE5tsv4L5Q2UuEk=",
-      "manifest.json": "BfEMlX6Fvw/2pFVCJ1L9FoZ2lBpbHjLoAr2PbX2GBZc=",
-      "one-hot-fill.wgsl.jinja": "+1+3T4yVF91Ta11a5kEZHAC8D3C/uGzVKZm/CSK4Htc=",
-      "one-hot-last-axis-vec4.wgsl.jinja": "nQ+vRZPx/rWi3Q+vuIhccDJzW6BATwWLFiaDf+HNwXo=",
-      "one-hot-scatter.wgsl.jinja": "xXJjRoPmpglSeIQvxCv+9mmygMaGMw1LTKy82ViMVJ8=",
-      "test.json": "UrZtlg9pSruj4lfpgeNOargkGctAr+ZUIM0fkeBJNsg="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.OneHot" }
 }

 {
   "name": "ai.onnx.OneHot",
+  "id": "_ai_onnx_onehot_webgpu_0f0f777",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "+FUw5FwoY0/2q5f1AHqGV4NyTYHeaoGS4hGUicbgYY4=",
+      "manifest.json": "yE4QIi2N3IwQ46XwkCOTuxL7t8DEXG7TX5L/updAYPA=",
+      "one-hot-fill.wgsl.jinja": "r55cBIB0D4pWDM63laO3cZbhz8Wox3J1FHZxHRplsM8=",
+      "one-hot-last-axis-vec4.wgsl.jinja": "kE19Rf4cbMJg+vJYsg3nI+T6wRd0pmF3t3Ah3xvo2mY=",
+      "one-hot-scatter.wgsl.jinja": "NEx5QGZ06RdFLM3T+T6L9jS2C0bkujL5niaeV9+HS38=",
+      "test.json": "nzqqm3CURddTStvPg1PZTCjgK+l1CyGWyK9EQwrXccU="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "last_axis_row_vec4": ["one-hot-last-axis-vec4.wgsl.jinja"],
+      "last_axis_vec4": ["one-hot-last-axis-vec4.wgsl.jinja"],
+      "generic_axis": ["one-hot-fill.wgsl.jinja", "one-hot-scatter.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/one-hot-fill.wgsl.jinja CHANGED Viewed

@@ -1,48 +1,45 @@
 {% 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
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
 {% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at the device's per-axis workgroup
-  // limit; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "dispatch" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
 {% endif %}
 {% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Dense off-value fill for arbitrary-axis OneHot. The vector form writes four
 // elements per invocation and the scalar form handles the tail. A sparse pass
 // writes the on-values afterward.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d(note="") }}
-{% if source.vectorized %}
   output[i] = {{ vectorScalar }}(values[0]);
 {% else %}
   output[params.offset + i] = values[0];

 {% 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 dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // Dense off-value fill for arbitrary-axis OneHot. The vector form writes four
 // elements per invocation and the scalar form handles the tail. A sparse pass
 // writes the on-values afterward.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d(note="") }}
+{% if vectorizedSpec %}
   output[i] = {{ vectorScalar }}(values[0]);
 {% else %}
   output[params.offset + i] = values[0];

build/webgpu/one-hot-last-axis-vec4.wgsl.jinja CHANGED Viewed

@@ -1,45 +1,42 @@
 {% 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
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
 {% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at the device's per-axis workgroup
-  // limit; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "dispatch" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
 {% endif %}
 {% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Vec4 OneHot for a contiguous depth-last output. Wide rows are assigned one
 // cooperative workgroup so they share the index load and avoid flat division;
 // short rows use the dense flat launch to avoid underfilled workgroups.
-const DEPTH: u32 = {{ source.depth }}u;
 const VEC_DEPTH: u32 = DEPTH / 4u;
 struct NormalizedIndex {
@@ -62,20 +59,19 @@ fn normalize_index(raw: {{ indexScalar }}) -> NormalizedIndex {
 {% endif %}
 }
-{% if source.cooperative %}
-const ROWS: u32 = {{ source.rows }}u;
-const WG: u32 = {{ source.wg }}u;
 var<workgroup> row_index: {{ indexScalar }};
-@compute @workgroup_size({{ source.wg }})
 fn main(
   @builtin(local_invocation_id) lid3: vec3<u32>,
-  @builtin(workgroup_id) wid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let lane = lid3.x;
-  let row = wid.x + wid.y * nwg.x;
   if (row >= ROWS) {
     return;
   }
@@ -99,7 +95,7 @@ fn main(
 }
 {% else %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d("i4", note="") }}
   let d4 = i4 % VEC_DEPTH;
   let d0 = d4 * 4u;

 {% 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 dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // Vec4 OneHot for a contiguous depth-last output. Wide rows are assigned one
 // cooperative workgroup so they share the index load and avoid flat division;
 // short rows use the dense flat launch to avoid underfilled workgroups.
+const DEPTH: u32 = {{ depth }}u;
 const VEC_DEPTH: u32 = DEPTH / 4u;
 struct NormalizedIndex {
 {% endif %}
 }
+{% if cooperative %}
+const ROWS: u32 = {{ rowCount }}u;
+const WG: u32 = {{ wg }}u;
 var<workgroup> row_index: {{ indexScalar }};
+@compute @workgroup_size({{ wg }})
 fn main(
   @builtin(local_invocation_id) lid3: vec3<u32>,
+  @builtin(workgroup_id) wid: vec3<u32>
 ) {
   let lane = lid3.x;
+  let row = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= ROWS) {
     return;
   }
 }
 {% else %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d("i4", note="") }}
   let d4 = i4 % VEC_DEPTH;
   let d0 = d4 * 4u;

build/webgpu/one-hot-scatter.wgsl.jinja CHANGED Viewed

@@ -1,49 +1,46 @@
 {% 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
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
 {% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at the device's per-axis workgroup
-  // limit; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "dispatch" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
 {% endif %}
 {% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Sparse on-value scatter for an output laid out as [outer, depth, inner].
 // One thread owns one input index, so output locations are unique even when
 // several input elements select the same depth value.
-const DEPTH: u32 = {{ source.depth }}u;
-const INNER: u32 = {{ source.inner }}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(note="") }}
   let raw = indices[i];
 {% if floatIndices %}

 {% 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 dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // Sparse on-value scatter for an output laid out as [outer, depth, inner].
 // One thread owns one input index, so output locations are unique even when
 // several input elements select the same depth value.
+const DEPTH: u32 = {{ depth }}u;
+const INNER: u32 = {{ innerSize }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d(note="") }}
   let raw = indices[i];
 {% if floatIndices %}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.OneHot",
   "cases": [
     {
       "name": "float16_indices_depth_and_values",
@@ -63,7 +62,7 @@
     {
       "name": "float32_indices_int16_depth_float32_values",
       "provenance": {
-        "notes": "An integral float32 index route with independently typed int16 depth connects both standard type variables while the all-int16 sibling retains widened-storage coverage."
       },
       "attrs": { "axis": -1 },
       "inputs": {
@@ -76,7 +75,7 @@
     {
       "name": "float32_indices_int8_depth_float32_values",
       "provenance": {
-        "notes": "An integral float32 index route with independently typed int8 depth connects both standard type variables while the all-int8 sibling retains widened-storage coverage."
       },
       "attrs": { "axis": 0 },
       "inputs": {
@@ -89,7 +88,7 @@
     {
       "name": "float32_indices_uint8_depth_float32_values",
       "provenance": {
-        "notes": "An integral float32 index route with independently typed uint8 depth connects both standard type variables while the all-uint8 sibling retains widened-storage coverage."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -435,7 +434,7 @@
       "outputs": { "output": { "dtype": "float32", "shape": [3, 10] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_negative_indices",
-        "notes": "ONNX int64 metadata/index tensors use framework int32/uint32 slots where representable. Official backend depth is represented with framework uint32 metadata."
       }
     },
     {
@@ -449,7 +448,7 @@
       "outputs": { "output": { "dtype": "float32", "shape": [2, 10, 2] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_with_axis",
-        "notes": "Official backend fixture stores integer-looking indices/depth as floats; represented with framework integer metadata tensors. Official backend depth is represented with framework uint32 metadata."
       }
     },
     {
@@ -463,7 +462,7 @@
       "outputs": { "output": { "dtype": "float32", "shape": [2, 10, 2] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_with_negative_axis",
-        "notes": "Official backend fixture stores integer-looking indices/depth as floats; represented with framework integer metadata tensors. Official backend depth is represented with framework uint32 metadata."
       }
     },
     {
@@ -476,7 +475,7 @@
       "outputs": { "output": { "dtype": "int32", "shape": [3, 12] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_without_axis",
-        "notes": "ONNX int64 metadata/index tensors use framework int32/uint32 slots where representable. Official backend depth is represented with framework uint32 metadata."
       }
     },
     {
@@ -484,7 +483,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc",
         "test": "OneHotOpTest.DefaultAxis_int64_int32_float_NonZeroOffValue",
-        "notes": "Uses framework int32 depth instead of ORT's float depth so WebGPU can treat depth as shape metadata."
       },
       "inputs": {
         "indices": { "dtype": "int32", "shape": [4], "data": { "kind": "values", "values": [0, 2, -1, 5] } },
@@ -666,7 +665,7 @@
     {
       "name": "f16_last_axis_vec4_depth8",
       "provenance": {
-        "notes": "f16 OneHot with a depth divisible by 4 selects last_axis_vec4 with usesF16=true (a vec4<f16> 64-bit store). Every existing f16 case has depth 3 or 10 (not %4) and routes to the scalar generic_axis kernel, so the vec4<f16> store branch is entirely untested. A real f16 embedding/classification one-hot over a %4 vocab hits this path."
       },
       "attrs": { "axis": -1 },
       "inputs": {
@@ -702,7 +701,7 @@
     {
       "name": "float_indices_last_axis_vec4_depth8",
       "provenance": {
-        "notes": "Float32 indices with a depth divisible by 4 select last_axis_vec4 with floatIndices=true (i32(trunc(raw)) + negative-wrap, then a vec4 store). Every existing float-index case has depth 10 (not %4) and routes to generic_axis, so the floatIndices branch of the vec4 kernel is untested. ONNX permits numeric indices (cast to int); a negative float index exercises the truncation+wrap."
       },
       "attrs": { "axis": -1 },
       "inputs": {

 {
   "cases": [
     {
       "name": "float16_indices_depth_and_values",
     {
       "name": "float32_indices_int16_depth_float32_values",
       "provenance": {
+        "notes": "Integral float32 indices with an independently typed int16 depth exercise both standard type variables and widened depth storage."
       },
       "attrs": { "axis": -1 },
       "inputs": {
     {
       "name": "float32_indices_int8_depth_float32_values",
       "provenance": {
+        "notes": "Integral float32 indices with an independently typed int8 depth exercise both standard type variables and widened depth storage."
       },
       "attrs": { "axis": 0 },
       "inputs": {
     {
       "name": "float32_indices_uint8_depth_float32_values",
       "provenance": {
+        "notes": "Integral float32 indices with an independently typed uint8 depth exercise both standard type variables and widened depth storage."
       },
       "attrs": { "axis": 1 },
       "inputs": {
       "outputs": { "output": { "dtype": "float32", "shape": [3, 10] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_negative_indices",
+        "notes": "The integer metadata tensors exactly represent the source ONNX int64 values, with `depth` stored as uint32."
       }
     },
     {
       "outputs": { "output": { "dtype": "float32", "shape": [2, 10, 2] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_with_axis",
+        "notes": "Integer-valued indices and depth are represented with integer metadata tensors, using uint32 for `depth`."
       }
     },
     {
       "outputs": { "output": { "dtype": "float32", "shape": [2, 10, 2] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_with_negative_axis",
+        "notes": "Integer-valued indices and depth are represented with integer metadata tensors, using uint32 for `depth`."
       }
     },
     {
       "outputs": { "output": { "dtype": "int32", "shape": [3, 12] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_onehot_without_axis",
+        "notes": "The integer metadata tensors exactly represent the source ONNX int64 values, with `depth` stored as uint32."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc",
         "test": "OneHotOpTest.DefaultAxis_int64_int32_float_NonZeroOffValue",
+        "notes": "Represents the integral `depth` value as int32 shape metadata."
       },
       "inputs": {
         "indices": { "dtype": "int32", "shape": [4], "data": { "kind": "values", "values": [0, 2, -1, 5] } },
     {
       "name": "f16_last_axis_vec4_depth8",
       "provenance": {
+        "notes": "A float16 OneHot with depth 8 selects the vectorized last-axis path and stores four float16 values per output group. Indices target the first, interior, and final classes."
       },
       "attrs": { "axis": -1 },
       "inputs": {
     {
       "name": "float_indices_last_axis_vec4_depth8",
       "provenance": {
+        "notes": "Float32 indices with depth 8 select the vectorized last-axis path. Fractional indices are truncated before indexing, negative indices wrap from the end, and out-of-range indices leave an all-off row."
       },
       "attrs": { "axis": -1 },
       "inputs": {