sync 91d990483a17

Browse files

Files changed (6) hide show

README.md +20 -17
build/webgpu/bench.json +1 -3
build/webgpu/manifest.json +113 -199
build/webgpu/metadata.json +15 -7
build/webgpu/split-to-sequence.wgsl.jinja +14 -17
build/webgpu/test.json +5 -6

README.md CHANGED Viewed

@@ -12,27 +12,27 @@ tags:
 ## Description
-Support status: the standard ONNX `SplitToSequence` operator is not implemented because this kernel ABI exposes tensors rather than an ONNX sequence value. This internal lowering copies an input into two, three, four, or six caller-shaped tensor outputs and therefore does not implement the standard `output_sequence` contract; it must not be treated as ONNX `SplitToSequence`.
 See the [standard ONNX `SplitToSequence` spec](https://onnx.ai/onnx/operators/onnx__SplitToSequence.html) for the contract this internal lowering does not implement.
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `input` | `input` | `T` | — | — | The tensor to split. | required |
-| `split` | `split` | `S` | — | — | Length of each output slice: a scalar for uniform chunks or a 1-D tensor of per-output lengths. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y0` | `y0` | `T` | derived | — | First output tensor slice. | required |
-| `Y1` | `y1` | `T` | derived | — | Second output tensor slice. | required |
-| `Y2` | `y2` | `T` | derived | — | Third output tensor slice. | optional |
-| `Y3` | `y3` | `T` | derived | — | Fourth output tensor slice. | optional |
-| `Y4` | `y4` | `T` | derived | — | Fifth output tensor slice. | optional |
-| `Y5` | `y5` | `T` | derived | — | Sixth output tensor slice. | optional |
 ## Attributes
@@ -52,7 +52,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
@@ -60,16 +60,19 @@ 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:
 - `y0`
 - `y1`
-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.

 ## Description
+Internal fixed-output split lowering that returns two, three, four, or six tensors with invocation-supplied shapes. It does not construct an ONNX sequence value.
 See the [standard ONNX `SplitToSequence` spec](https://onnx.ai/onnx/operators/onnx__SplitToSequence.html) for the contract this internal lowering does not implement.
 ## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `input` | `T` | — | — | The tensor to split. | required |
+| `split` | `S` | — | — | Length of each output slice: a scalar for uniform chunks or a 1-D tensor of per-output lengths. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y0` | `Y0` | `T` | derived | — | First output tensor slice. | required |
+| `y1` | `Y1` | `T` | derived | — | Second output tensor slice. | required |
+| `y2` | `Y2` | `T` | derived | — | Third output tensor slice. | optional |
+| `y3` | `Y3` | `T` | derived | — | Fourth output tensor slice. | optional |
+| `y4` | `Y4` | `T` | derived | — | Fifth output tensor slice. | optional |
+| `y5` | `Y5` | `T` | derived | — | Sixth output tensor slice. | optional |
 ## 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:
 - `y0`
 - `y1`
 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.SplitToSequence",
   "cases": [
     {
       "name": "split3-1536x256-f32",
@@ -20,8 +19,7 @@
       "preset": "stress",
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
-        "input": { "dtype": "float32", "shape": [8, 6, 4, 4, 16, 64], "dist": "normal", "seed": 1234, "scale": 2 },
-        "split": { "dtype": "uint32", "shape": [], "dist": "constant", "value": 1 }
       },
       "outputs": {
         "y0": { "dtype": "float32", "shape": [8, 4, 4, 16, 64], "dist": "empty" },

 {
   "cases": [
     {
       "name": "split3-1536x256-f32",
       "preset": "stress",
       "attrs": { "axis": 1, "keepdims": 0 },
       "inputs": {
+        "input": { "dtype": "float32", "shape": [8, 6, 4, 4, 16, 64], "dist": "normal", "seed": 1234, "scale": 2 }
       },
       "outputs": {
         "y0": { "dtype": "float32", "shape": [8, 4, 4, 16, 64], "dist": "empty" },

build/webgpu/manifest.json CHANGED Viewed

@@ -3,85 +3,47 @@
   "name": "SplitToSequence",
   "conformance": "internal-lowering",
   "sinceVersion": 11,
-  "description": "Support status: the standard ONNX `SplitToSequence` operator is not implemented because this kernel ABI exposes tensors rather than an ONNX sequence value. This internal lowering copies an input into two, three, four, or six caller-shaped tensor outputs and therefore does not implement the standard `output_sequence` contract; it must not be treated as ONNX `SplitToSequence`.",
-  "inputs": [
-    { "role": "input", "dtype": "T", "description": "The tensor to split." },
-    {
-      "role": "split",
-      "dtype": "S",
-      "optional": true,
-      "description": "Length of each output slice: a scalar for uniform chunks or a 1-D tensor of per-output lengths."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "Y0",
       "dtype": "T",
-      "description": "First output tensor slice.",
-      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1"
     },
-    {
-      "role": "Y1",
       "dtype": "T",
-      "description": "Second output tensor slice.",
-      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1"
     },
-    {
-      "role": "Y2",
       "dtype": "T",
-      "optional": true,
-      "description": "Third output tensor slice.",
-      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1"
     },
-    {
-      "role": "Y3",
       "dtype": "T",
-      "optional": true,
-      "description": "Fourth output tensor slice.",
-      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1"
-    },
-    {
-      "role": "Y4",
-      "dtype": "T",
-      "optional": true,
-      "description": "Fifth output tensor slice.",
-      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1"
-    },
-    {
-      "role": "Y5",
-      "dtype": "T",
-      "optional": true,
-      "description": "Sixth output tensor slice.",
-      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1"
     }
-  ],
-  "attributes": { "axis": 0, "keepdims": 1 },
-  "attributeDescriptions": {
-    "axis": "Axis along which to split; negative values count from the back. Accepted range is `[-rank, rank-1]`.",
-    "keepdims": "Whether to keep the split dimension in the output (default `1`). Ignored when `split` is provided."
   },
   "attributeConstraints": { "keepdims": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "bool"], "S": ["uint32"] },
-  "args": {
-    "input": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "split": { "kind": "tensor", "semantic": "split", "role": "split", "required": false },
-    "y0": { "kind": "tensor", "semantic": "Y0", "role": "output0" },
-    "y1": { "kind": "tensor", "semantic": "Y1", "role": "output1" },
-    "y2": { "kind": "tensor", "semantic": "Y2", "role": "output2", "required": false },
-    "y3": { "kind": "tensor", "semantic": "Y3", "role": "output3", "required": false },
-    "y4": { "kind": "tensor", "semantic": "Y4", "role": "output4", "required": false },
-    "y5": { "kind": "tensor", "semantic": "Y5", "role": "output5", "required": false }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256 },
   "derive": {
     "baseContract": "ranks.input >= 1 and (attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input) >= 0 and (attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input) < ranks.input and f16Ok(dtypes.T)",
-    "outputRankContract": "ranks.Y0 == ranks.input if present.split else (ranks.Y0 == ranks.input if attrs.keepdims == 1 else ranks.Y0 == ranks.input - 1)",
-    "twoOutputContract": "baseContract and not present.y2 and ranks.Y0 == ranks.Y1 and outputRankContract",
-    "threeOutputContract": "baseContract and present.y2 and not present.y3 and ranks.Y0 == ranks.Y1 and ranks.Y0 == ranks.Y2 and outputRankContract",
-    "fourOutputContract": "baseContract and present.y2 and present.y3 and not present.y4 and ranks.Y0 == ranks.Y1 and ranks.Y0 == ranks.Y2 and ranks.Y0 == ranks.Y3 and outputRankContract",
-    "sixOutputContract": "baseContract and present.y2 and present.y3 and present.y4 and present.y5 and ranks.Y0 == ranks.Y1 and ranks.Y0 == ranks.Y2 and ranks.Y0 == ranks.Y3 and ranks.Y0 == ranks.Y4 and ranks.Y0 == ranks.Y5 and outputRankContract"
   },
-  "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
   "variants": [
     {
       "id": "two_outputs",
@@ -90,43 +52,31 @@
         {
           "id": "main",
           "name": "SplitToSequence2",
-          "source": {
-            "shader": "split-to-sequence.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "rank": "ranks.input",
-              "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
-              "squeeze": "ranks.Y0 < ranks.input",
-              "numOutputs": "2",
-              "outShapes": ["shapes.Y0", "shapes.Y1"]
-            }
           },
           "bindings": [
-            {
-              "name": "input",
-              "arg": "input",
-              "semantic": "input",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$scalar" },
             {
               "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [
-                  { "name": "y0Count", "type": "u32", "value": "numel(shapes.Y0)" },
-                  { "name": "y1Count", "type": "u32", "value": "numel(shapes.Y1)" }
-                ]
-              }
             }
           ],
           "dispatch": {
-            "threads": "max(numel(shapes.Y0), numel(shapes.Y1))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -139,45 +89,33 @@
         {
           "id": "main",
           "name": "SplitToSequence3",
-          "source": {
-            "shader": "split-to-sequence.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "rank": "ranks.input",
-              "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
-              "squeeze": "ranks.Y0 < ranks.input",
-              "numOutputs": "3",
-              "outShapes": ["shapes.Y0", "shapes.Y1", "shapes.Y2"]
-            }
           },
           "bindings": [
-            {
-              "name": "input",
-              "arg": "input",
-              "semantic": "input",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y2", "arg": "y2", "semantic": "Y2", "buffer": { "type": "storage" }, "elementType": "$scalar" },
             {
               "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [
-                  { "name": "y0Count", "type": "u32", "value": "numel(shapes.Y0)" },
-                  { "name": "y1Count", "type": "u32", "value": "numel(shapes.Y1)" },
-                  { "name": "y2Count", "type": "u32", "value": "numel(shapes.Y2)" }
-                ]
-              }
             }
           ],
           "dispatch": {
-            "threads": "max(max(numel(shapes.Y0), numel(shapes.Y1)), numel(shapes.Y2))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -190,47 +128,35 @@
         {
           "id": "main",
           "name": "SplitToSequence4",
-          "source": {
-            "shader": "split-to-sequence.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "rank": "ranks.input",
-              "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
-              "squeeze": "ranks.Y0 < ranks.input",
-              "numOutputs": "4",
-              "outShapes": ["shapes.Y0", "shapes.Y1", "shapes.Y2", "shapes.Y3"]
-            }
           },
           "bindings": [
-            {
-              "name": "input",
-              "arg": "input",
-              "semantic": "input",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y2", "arg": "y2", "semantic": "Y2", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y3", "arg": "y3", "semantic": "Y3", "buffer": { "type": "storage" }, "elementType": "$scalar" },
             {
               "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [
-                  { "name": "y0Count", "type": "u32", "value": "numel(shapes.Y0)" },
-                  { "name": "y1Count", "type": "u32", "value": "numel(shapes.Y1)" },
-                  { "name": "y2Count", "type": "u32", "value": "numel(shapes.Y2)" },
-                  { "name": "y3Count", "type": "u32", "value": "numel(shapes.Y3)" }
-                ]
-              }
             }
           ],
           "dispatch": {
-            "threads": "max(max(max(numel(shapes.Y0), numel(shapes.Y1)), numel(shapes.Y2)), numel(shapes.Y3))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -243,51 +169,39 @@
         {
           "id": "main",
           "name": "SplitToSequence6",
-          "source": {
-            "shader": "split-to-sequence.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "rank": "ranks.input",
-              "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
-              "squeeze": "ranks.Y0 < ranks.input",
-              "numOutputs": "6",
-              "outShapes": ["shapes.Y0", "shapes.Y1", "shapes.Y2", "shapes.Y3", "shapes.Y4", "shapes.Y5"]
-            }
           },
           "bindings": [
-            {
-              "name": "input",
-              "arg": "input",
-              "semantic": "input",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y2", "arg": "y2", "semantic": "Y2", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y3", "arg": "y3", "semantic": "Y3", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y4", "arg": "y4", "semantic": "Y4", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            { "name": "y5", "arg": "y5", "semantic": "Y5", "buffer": { "type": "storage" }, "elementType": "$scalar" },
             {
               "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [
-                  { "name": "y0Count", "type": "u32", "value": "numel(shapes.Y0)" },
-                  { "name": "y1Count", "type": "u32", "value": "numel(shapes.Y1)" },
-                  { "name": "y2Count", "type": "u32", "value": "numel(shapes.Y2)" },
-                  { "name": "y3Count", "type": "u32", "value": "numel(shapes.Y3)" },
-                  { "name": "y4Count", "type": "u32", "value": "numel(shapes.Y4)" },
-                  { "name": "y5Count", "type": "u32", "value": "numel(shapes.Y5)" }
-                ]
-              }
             }
           ],
           "dispatch": {
-            "threads": "max(max(max(max(max(numel(shapes.Y0), numel(shapes.Y1)), numel(shapes.Y2)), numel(shapes.Y3)), numel(shapes.Y4)), numel(shapes.Y5))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]

   "name": "SplitToSequence",
   "conformance": "internal-lowering",
   "sinceVersion": 11,
+  "inputs": { "input": { "dtype": "T" }, "split": { "dtype": "S", "optional": true } },
+  "outputs": {
+    "y0": { "onnx": "Y0", "dtype": "T", "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1" },
+    "y1": { "onnx": "Y1", "dtype": "T", "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1" },
+    "y2": {
+      "onnx": "Y2",
       "dtype": "T",
+      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1",
+      "optional": true
     },
+    "y3": {
+      "onnx": "Y3",
       "dtype": "T",
+      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1",
+      "optional": true
     },
+    "y4": {
+      "onnx": "Y4",
       "dtype": "T",
+      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1",
+      "optional": true
     },
+    "y5": {
+      "onnx": "Y5",
       "dtype": "T",
+      "rank": "ranks.input if attrs.keepdims == 1 else ranks.input - 1",
+      "optional": true
     }
   },
+  "attributes": { "axis": { "default": 0 }, "keepdims": { "default": 1 } },
   "attributeConstraints": { "keepdims": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "bool"], "S": ["uint32"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
     "baseContract": "ranks.input >= 1 and (attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input) >= 0 and (attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input) < ranks.input and f16Ok(dtypes.T)",
+    "outputRankContract": "ranks.y0 == ranks.input if present.split else (ranks.y0 == ranks.input if attrs.keepdims == 1 else ranks.y0 == ranks.input - 1)",
+    "twoOutputContract": "baseContract and not present.y2 and ranks.y0 == ranks.y1 and outputRankContract",
+    "threeOutputContract": "baseContract and present.y2 and not present.y3 and ranks.y0 == ranks.y1 and ranks.y0 == ranks.y2 and outputRankContract",
+    "fourOutputContract": "baseContract and present.y2 and present.y3 and not present.y4 and ranks.y0 == ranks.y1 and ranks.y0 == ranks.y2 and ranks.y0 == ranks.y3 and outputRankContract",
+    "sixOutputContract": "baseContract and present.y2 and present.y3 and present.y4 and present.y5 and ranks.y0 == ranks.y1 and ranks.y0 == ranks.y2 and ranks.y0 == ranks.y3 and ranks.y0 == ranks.y4 and ranks.y0 == ranks.y5 and outputRankContract"
   },
   "variants": [
     {
       "id": "two_outputs",
         {
           "id": "main",
           "name": "SplitToSequence2",
+          "shader": "split-to-sequence.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "rank": "ranks.input",
+            "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
+            "squeeze": "ranks.y0 < ranks.input",
+            "numOutputs": "2",
+            "outShapes": ["shapes.y0", "shapes.y1"]
           },
           "bindings": [
+            "input",
+            "y0",
+            "y1",
             {
               "name": "params",
+              "struct": [
+                { "name": "y0Count", "type": "u32", "value": "numel(shapes.y0)" },
+                { "name": "y1Count", "type": "u32", "value": "numel(shapes.y1)" }
+              ]
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(numel(shapes.y0), numel(shapes.y1))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "SplitToSequence3",
+          "shader": "split-to-sequence.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "rank": "ranks.input",
+            "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
+            "squeeze": "ranks.y0 < ranks.input",
+            "numOutputs": "3",
+            "outShapes": ["shapes.y0", "shapes.y1", "shapes.y2"]
           },
           "bindings": [
+            "input",
+            "y0",
+            "y1",
+            "y2",
             {
               "name": "params",
+              "struct": [
+                { "name": "y0Count", "type": "u32", "value": "numel(shapes.y0)" },
+                { "name": "y1Count", "type": "u32", "value": "numel(shapes.y1)" },
+                { "name": "y2Count", "type": "u32", "value": "numel(shapes.y2)" }
+              ]
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((max(max(numel(shapes.y0), numel(shapes.y1)), numel(shapes.y2))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(max(numel(shapes.y0), numel(shapes.y1)), numel(shapes.y2))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "SplitToSequence4",
+          "shader": "split-to-sequence.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "rank": "ranks.input",
+            "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
+            "squeeze": "ranks.y0 < ranks.input",
+            "numOutputs": "4",
+            "outShapes": ["shapes.y0", "shapes.y1", "shapes.y2", "shapes.y3"]
           },
           "bindings": [
+            "input",
+            "y0",
+            "y1",
+            "y2",
+            "y3",
             {
               "name": "params",
+              "struct": [
+                { "name": "y0Count", "type": "u32", "value": "numel(shapes.y0)" },
+                { "name": "y1Count", "type": "u32", "value": "numel(shapes.y1)" },
+                { "name": "y2Count", "type": "u32", "value": "numel(shapes.y2)" },
+                { "name": "y3Count", "type": "u32", "value": "numel(shapes.y3)" }
+              ]
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((max(max(max(numel(shapes.y0), numel(shapes.y1)), numel(shapes.y2)), numel(shapes.y3))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(max(max(numel(shapes.y0), numel(shapes.y1)), numel(shapes.y2)), numel(shapes.y3))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "SplitToSequence6",
+          "shader": "split-to-sequence.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "rank": "ranks.input",
+            "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
+            "squeeze": "ranks.y0 < ranks.input",
+            "numOutputs": "6",
+            "outShapes": ["shapes.y0", "shapes.y1", "shapes.y2", "shapes.y3", "shapes.y4", "shapes.y5"]
           },
           "bindings": [
+            "input",
+            "y0",
+            "y1",
+            "y2",
+            "y3",
+            "y4",
+            "y5",
             {
               "name": "params",
+              "struct": [
+                { "name": "y0Count", "type": "u32", "value": "numel(shapes.y0)" },
+                { "name": "y1Count", "type": "u32", "value": "numel(shapes.y1)" },
+                { "name": "y2Count", "type": "u32", "value": "numel(shapes.y2)" },
+                { "name": "y3Count", "type": "u32", "value": "numel(shapes.y3)" },
+                { "name": "y4Count", "type": "u32", "value": "numel(shapes.y4)" },
+                { "name": "y5Count", "type": "u32", "value": "numel(shapes.y5)" }
+              ]
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((max(max(max(max(max(numel(shapes.y0), numel(shapes.y1)), numel(shapes.y2)), numel(shapes.y3)), numel(shapes.y4)), numel(shapes.y5))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(max(max(max(max(numel(shapes.y0), numel(shapes.y1)), numel(shapes.y2)), numel(shapes.y3)), numel(shapes.y4)), numel(shapes.y5))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,18 +1,26 @@
 {
   "name": "ai.onnx.SplitToSequence",
-  "id": "_ai_onnx_splittosequence_webgpu_e21e7a0",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "EENkbgxvL5j8em9SHo+NktW1UwzMESrX+8p/FKNCoQc=",
-      "manifest.json": "bKxHYk9Csq9b7FX6sH6WaiOZyBawG9Om33PdtCe3kKs=",
-      "split-to-sequence.wgsl.jinja": "LTXzbRgOGFK2BVTVM42WPJ9DWr+mx8PcoY9jqf4KBUc=",
-      "test.json": "r6D3suRTZe+ljNzJvTQu6Le3qDRRZ6RR68+ObgTAIMY="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.SplitToSequence" }
 }

 {
   "name": "ai.onnx.SplitToSequence",
+  "id": "_ai_onnx_splittosequence_webgpu_915cb2a",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "Zs84ZKi7bfjFdmtbBi7218N/ixgvCEbOJMXn6YS5yn4=",
+      "manifest.json": "KLNBWHNVH08iKjw2HfJ33n1lHvbg3fmdzBcp6ioyvgk=",
+      "split-to-sequence.wgsl.jinja": "RfgQQlUXADMjM2fTdCz0/SSpd5x5CMfNeIjywANBVrM=",
+      "test.json": "FAy7gFEfxoFsB44fhBjHbN4cWZYYHKlVqz09kYcJMr0="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "two_outputs": ["split-to-sequence.wgsl.jinja"],
+      "three_outputs": ["split-to-sequence.wgsl.jinja"],
+      "four_outputs": ["split-to-sequence.wgsl.jinja"],
+      "six_outputs": ["split-to-sequence.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/split-to-sequence.wgsl.jinja CHANGED Viewed

@@ -1,33 +1,30 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// SplitToSequence moves contiguous slices along the split axis; the output
-// tensors form the sequence. For size-one slices with
 // keepdims=0, each output drops the axis dimension. The copy reconstructs the
 // full input coordinate by reinserting that axis position.
-{% set R = source.rank %}
-{% set ax = source.axis %}
 {% set acc = namespace(offset=0) %}
-{% for k in range(source.numOutputs) %}
-{% set oShape = source.outShapes[k] %}
 fn input_offset_y{{ k }}(out_index: u32) -> u32 {
   var rem = out_index;
   var offset = 0u;
-{% if source.squeeze %}
 {% for d in range(R - 1) %}
 {% set out_stride = namespace(value=1) %}
 {% for j in range(d + 1, R - 1) %}{% set out_stride.value = out_stride.value * oShape[j] %}{% endfor %}
 {% set in_dim = d if d < ax else d + 1 %}
 {% set in_stride = namespace(value=1) %}
-{% for j in range(in_dim + 1, R) %}{% set in_stride.value = in_stride.value * source.inputShape[j] %}{% endfor %}
   let c{{ d }} = rem / {{ out_stride.value if out_stride.value > 0 else 1 }}u;
   rem = rem % {{ out_stride.value if out_stride.value > 0 else 1 }}u;
   offset = offset + c{{ d }} * {{ in_stride.value }}u;
 {% endfor %}
 {% set axis_stride = namespace(value=1) %}
-{% for j in range(ax + 1, R) %}{% set axis_stride.value = axis_stride.value * source.inputShape[j] %}{% endfor %}
   offset = offset + {{ acc.offset }}u * {{ axis_stride.value }}u;
 {% set acc.offset = acc.offset + 1 %}
 {% else %}
@@ -35,7 +32,7 @@ fn input_offset_y{{ k }}(out_index: u32) -> u32 {
 {% set out_stride = namespace(value=1) %}
 {% for j in range(d + 1, R) %}{% set out_stride.value = out_stride.value * oShape[j] %}{% endfor %}
 {% set in_stride = namespace(value=1) %}
-{% for j in range(d + 1, R) %}{% set in_stride.value = in_stride.value * source.inputShape[j] %}{% endfor %}
   var c{{ d }} = rem / {{ out_stride.value if out_stride.value > 0 else 1 }}u;
   rem = rem % {{ out_stride.value if out_stride.value > 0 else 1 }}u;
 {% if d == ax %}
@@ -50,11 +47,11 @@ fn input_offset_y{{ k }}(out_index: u32) -> u32 {
 {% endfor %}
 @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 past the
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
-  let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
-{% for k in range(source.numOutputs) %}
   if (i < params.y{{ k }}Count) {
     y{{ k }}[i] = input[input_offset_y{{ k }}(i)];
   }

 {{ env.wgsl.resourceDeclarations }}
+// This internal lowering moves contiguous slices along the split axis into
+// a fixed set of output tensors. For size-one slices with
 // keepdims=0, each output drops the axis dimension. The copy reconstructs the
 // full input coordinate by reinserting that axis position.
+{% set R = rank %}
+{% set ax = axisSpec %}
 {% set acc = namespace(offset=0) %}
+{% for k in range(numOutputs) %}
+{% set oShape = outShapes[k] %}
 fn input_offset_y{{ k }}(out_index: u32) -> u32 {
   var rem = out_index;
   var offset = 0u;
+{% if squeeze %}
 {% for d in range(R - 1) %}
 {% set out_stride = namespace(value=1) %}
 {% for j in range(d + 1, R - 1) %}{% set out_stride.value = out_stride.value * oShape[j] %}{% endfor %}
 {% set in_dim = d if d < ax else d + 1 %}
 {% set in_stride = namespace(value=1) %}
+{% for j in range(in_dim + 1, R) %}{% set in_stride.value = in_stride.value * inputShape[j] %}{% endfor %}
   let c{{ d }} = rem / {{ out_stride.value if out_stride.value > 0 else 1 }}u;
   rem = rem % {{ out_stride.value if out_stride.value > 0 else 1 }}u;
   offset = offset + c{{ d }} * {{ in_stride.value }}u;
 {% endfor %}
 {% set axis_stride = namespace(value=1) %}
+{% for j in range(ax + 1, R) %}{% set axis_stride.value = axis_stride.value * inputShape[j] %}{% endfor %}
   offset = offset + {{ acc.offset }}u * {{ axis_stride.value }}u;
 {% set acc.offset = acc.offset + 1 %}
 {% else %}
 {% set out_stride = namespace(value=1) %}
 {% for j in range(d + 1, R) %}{% set out_stride.value = out_stride.value * oShape[j] %}{% endfor %}
 {% set in_stride = namespace(value=1) %}
+{% for j in range(d + 1, R) %}{% set in_stride.value = in_stride.value * inputShape[j] %}{% endfor %}
   var c{{ d }} = rem / {{ out_stride.value if out_stride.value > 0 else 1 }}u;
   rem = rem % {{ out_stride.value if out_stride.value > 0 else 1 }}u;
 {% if d == ax %}
 {% endfor %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width (outputs > 16.7M elements).
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
+{% for k in range(numOutputs) %}
   if (i < params.y{{ k }}Count) {
     y{{ k }}[i] = input[input_offset_y{{ k }}(i)];
   }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.SplitToSequence",
   "fixtureArrays": {
     "ort_positive_axis_scalar_split_rank3_input_input": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24],
     "onnx_backend_split_scalar_axis1_three_outputs_input_input": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
@@ -303,7 +302,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc",
         "test": "SequenceOpsTest.SplitToSequence_Axis0DefaultSplitFloatSetAxisExplicitly",
-        "notes": "Valid ONNX sequence output modelled as four fixed tensors."
       },
       "attrs": { "axis": 0 },
       "inputs": {
@@ -345,7 +344,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc",
         "test": "SequenceOpsTest.SplitToSequence_BoolSplit",
-        "notes": "Valid ONNX bool sequence output modelled as four fixed tensors."
       },
       "attrs": { "axis": 0 },
       "inputs": {
@@ -363,7 +362,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc",
         "test": "SequenceOpsTest.SplitToSequence_PositiveAxisDontKeepDims",
-        "notes": "Valid ONNX keepdims=0 split along a non-leading axis, modelled as four fixed tensors."
       },
       "attrs": { "axis": 2, "keepdims": 0 },
       "inputs": {
@@ -473,7 +472,7 @@
     {
       "name": "two_outputs_zero_dim",
       "provenance": {
-        "notes": "Two-chunk split of a tensor with a zero-width trailing axis. The two-output kernel floors each chunk stride at one so a zero-width output cannot divide by zero; the existing zero-dim case splits into three chunks, which is a different kernel."
       },
       "attrs": { "axis": 0 },
       "inputs": {
@@ -488,7 +487,7 @@
     {
       "name": "two_outputs_zero_dim_squeeze",
       "provenance": {
-        "notes": "Squeezed twin of the case above: keepdims=0 drops the split axis, which is a separate coordinate walk with its own chunk-stride floor. The zero-width trailing axis makes that floor fire."
       },
       "attrs": { "axis": 0, "keepdims": 0 },
       "inputs": { "input": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } } },

 {
   "fixtureArrays": {
     "ort_positive_axis_scalar_split_rank3_input_input": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24],
     "onnx_backend_split_scalar_axis1_three_outputs_input_input": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc",
         "test": "SequenceOpsTest.SplitToSequence_Axis0DefaultSplitFloatSetAxisExplicitly",
+        "notes": "The package's fixed-output lowering maps the four sequence elements to tensor outputs y0 through y3."
       },
       "attrs": { "axis": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc",
         "test": "SequenceOpsTest.SplitToSequence_BoolSplit",
+        "notes": "The package's fixed-output lowering maps the four boolean sequence elements to tensor outputs y0 through y3."
       },
       "attrs": { "axis": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc",
         "test": "SequenceOpsTest.SplitToSequence_PositiveAxisDontKeepDims",
+        "notes": "With keepdims=0, splitting axis 2 removes that axis from each chunk; the fixed-output lowering maps the four chunks to y0 through y3."
       },
       "attrs": { "axis": 2, "keepdims": 0 },
       "inputs": {
     {
       "name": "two_outputs_zero_dim",
       "provenance": {
+        "notes": "Splitting shape [4,0] along axis 0 produces two fixed tensor outputs of shape [2,0]. Zero-width chunks must not cause division by zero during offset calculation."
       },
       "attrs": { "axis": 0 },
       "inputs": {
     {
       "name": "two_outputs_zero_dim_squeeze",
       "provenance": {
+        "notes": "With keepdims=0, each fixed output drops the split axis and uses the squeezed coordinate walk. A zero-width trailing axis exercises the chunk-stride floor."
       },
       "attrs": { "axis": 0, "keepdims": 0 },
       "inputs": { "input": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } } },