sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +18 -15
build/webgpu/bench.json +0 -1
build/webgpu/datamove-flat-copy.wgsl.jinja +1 -4
build/webgpu/datamove-split-block.wgsl.jinja +26 -29
build/webgpu/manifest.json +194 -333
build/webgpu/metadata.json +22 -9
build/webgpu/split-n.wgsl.jinja +12 -15
build/webgpu/test.json +10 -11

README.md CHANGED Viewed

@@ -18,20 +18,20 @@ See the [ONNX `Split` spec](https://onnx.ai/onnx/operators/onnx__Split.html) for
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `input` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | The tensor to split. | required |
-| `split` | `split` | `S` | `uint32` | `1` | — | Optional logical int64 1-D tensor specifying the size of each output along the split axis; values must be non-negative, sum to the axis dimension, and use uint32 WebGPU storage. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y0` | `y0` | `T` | same as `input` | — | First output slice after splitting. | required |
-| `Y1` | `y1` | `T` | same as `input` | — | Second output slice after splitting (optional). | optional |
-| `Y2` | `y2` | `T` | same as `input` | — | Third output slice after splitting. | optional |
-| `Y3` | `y3` | `T` | same as `input` | — | Fourth output slice after splitting. | optional |
-| `Y4` | `y4` | `T` | same as `input` | — | Fifth output slice after splitting. | optional |
 ## Attributes
@@ -51,7 +51,7 @@ Attributes and 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
@@ -61,17 +61,20 @@ Attributes and 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`
 - `y2`
-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 | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `input` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | The tensor to split. | required |
+| `split` | `S` | `uint32` | `1` | — | Optional logical int64 1-D tensor specifying the size of each output along the split axis; values must be non-negative, sum to the axis dimension, and use uint32 WebGPU storage. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y0` | `Y0` | `T` | same as `input` | — | First output slice after splitting. | required |
+| `y1` | `Y1` | `T` | same as `input` | — | Second output slice after splitting (optional). | optional |
+| `y2` | `Y2` | `T` | same as `input` | — | Third output slice after splitting. | optional |
+| `y3` | `Y3` | `T` | same as `input` | — | Fourth output slice after splitting. | optional |
+| `y4` | `Y4` | `T` | same as `input` | — | Fifth output slice after splitting. | 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`
 - `y2`
 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.Split",
   "cases": [
     {
       "name": "qkv_like",

 {
   "cases": [
     {
       "name": "qkv_like",

build/webgpu/datamove-flat-copy.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,6 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const COUNT: u32 = {{ source.count }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})

 {{ env.wgsl.resourceDeclarations }}
+const COUNT: u32 = {{ count }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})

build/webgpu/datamove-split-block.wgsl.jinja CHANGED Viewed

@@ -3,35 +3,32 @@
 // divisible by four, all run boundaries are vec4-aligned and the input and
 // every output bind as vec4<T>. One division and one remainder per element
 // replace the per-axis coordinate decomposition of the scalar fallback.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set inner = namespace(value=1) %}
-{% for j in range(source.axis + 1, source.rank) %}
-{% set inner.value = inner.value * source.inputShape[j] %}
 {% endfor %}
-{% set y2Shape = source.y2Shape | default([]) %}
-{% if source.scalarBoundX4 %}
 // Scalar-bound x4 mode keeps arbitrary split boundaries legal. Each invocation
 // owns one four-scalar group within one outer slice, so a group never crosses
 // into the next slice; only the final group of a run needs scalar tail guards.
-{% set inner_in_scalar = source.inputShape[source.axis] * inner.value %}
-{% set run0_scalar = source.y0Shape[source.axis] * inner.value %}
-{% set run1_scalar = source.y1Shape[source.axis] * inner.value %}
 {% set groups0 = (((run0_scalar + 3) / 4) | int) %}
 {% set groups1 = (((run1_scalar + 3) / 4) | int) %}
 {% set outer_count = namespace(value=1) %}
-{% for j in range(source.axis) %}
-{% set outer_count.value = outer_count.value * source.inputShape[j] %}
 {% endfor %}
 {% set count0_groups = outer_count.value * groups0 %}
 {% set count1_groups = outer_count.value * groups1 %}
 {% set max_groups = namespace(value=count0_groups) %}
 {% if count1_groups > max_groups.value %}{% set max_groups.value = count1_groups %}{% endif %}
 {% set start1_scalar = run0_scalar %}
-{% set run2_scalar = y2Shape[source.axis] * inner.value %}
 {% set groups2 = (((run2_scalar + 3) / 4) | int) %}
 {% set count2_groups = outer_count.value * groups2 %}
 {% set start2_scalar = start1_scalar + run1_scalar %}
@@ -69,36 +66,36 @@ fn main(
   }
 }
 {% else %}
-{% set inner_in = (((source.inputShape[source.axis] * inner.value) / 4) | int) %}
-{% set run0 = (((source.y0Shape[source.axis] * inner.value) / 4) | int) %}
-{% set run1 = (((source.y1Shape[source.axis] * inner.value) / 4) | int) %}
 {% set has_nonempty_run = namespace(value=(run0 > 0 or run1 > 0)) %}
 {% set start1 = run0 %}
 {% set count0 = namespace(value=run0) %}
 {% set count1 = namespace(value=run1) %}
-{% for j in range(source.axis) %}
-{% set count0.value = count0.value * source.inputShape[j] %}
-{% set count1.value = count1.value * source.inputShape[j] %}
 {% endfor %}
 {% set max_count = namespace(value=count0.value) %}
 {% if count1.value > max_count.value %}{% set max_count.value = count1.value %}{% endif %}
-{% if source.outputCount >= 3 %}
-{% set run2 = (((y2Shape[source.axis] * inner.value) / 4) | int) %}
 {% if run2 > 0 %}{% set has_nonempty_run.value = true %}{% endif %}
 {% set start2 = start1 + run1 %}
 {% set count2 = namespace(value=run2) %}
-{% for j in range(source.axis) %}
-{% set count2.value = count2.value * source.inputShape[j] %}
 {% endfor %}
 {% if count2.value > max_count.value %}{% set max_count.value = count2.value %}{% endif %}
 {% endif %}
-{% if source.outputCount >= 4 %}
-{% set run3 = (((source.y3Shape[source.axis] * inner.value) / 4) | int) %}
 {% if run3 > 0 %}{% set has_nonempty_run.value = true %}{% endif %}
 {% set start3 = start2 + run2 %}
 {% set count3 = namespace(value=run3) %}
-{% for j in range(source.axis) %}
-{% set count3.value = count3.value * source.inputShape[j] %}
 {% endfor %}
 {% if count3.value > max_count.value %}{% set max_count.value = count3.value %}{% endif %}
 {% endif %}
@@ -125,12 +122,12 @@ fn main(
       y1[i] = input[(i / {{ run1 }}u) * INNER_IN + {{ start1 }}u + (i % {{ run1 }}u)];
     }
 {% endif %}
-{% if source.outputCount >= 3 and run2 > 0 %}
     if (i < {{ count2.value }}u) {
       y2[i] = input[(i / {{ run2 }}u) * INNER_IN + {{ start2 }}u + (i % {{ run2 }}u)];
     }
 {% endif %}
-{% if source.outputCount >= 4 and run3 > 0 %}
     if (i < {{ count3.value }}u) {
       y3[i] = input[(i / {{ run3 }}u) * INNER_IN + {{ start3 }}u + (i % {{ run3 }}u)];
     }

 // divisible by four, all run boundaries are vec4-aligned and the input and
 // every output bind as vec4<T>. One division and one remainder per element
 // replace the per-axis coordinate decomposition of the scalar fallback.
 {{ env.wgsl.resourceDeclarations }}
 {% set inner = namespace(value=1) %}
+{% for j in range(axisSpec + 1, rank) %}
+{% set inner.value = inner.value * inputShape[j] %}
 {% endfor %}
+{% set y2Shape = y2Shape | default([]) %}
+{% if scalarBoundX4 %}
 // Scalar-bound x4 mode keeps arbitrary split boundaries legal. Each invocation
 // owns one four-scalar group within one outer slice, so a group never crosses
 // into the next slice; only the final group of a run needs scalar tail guards.
+{% set inner_in_scalar = inputShape[axisSpec] * inner.value %}
+{% set run0_scalar = y0Shape[axisSpec] * inner.value %}
+{% set run1_scalar = y1Shape[axisSpec] * inner.value %}
 {% set groups0 = (((run0_scalar + 3) / 4) | int) %}
 {% set groups1 = (((run1_scalar + 3) / 4) | int) %}
 {% set outer_count = namespace(value=1) %}
+{% for j in range(axisSpec) %}
+{% set outer_count.value = outer_count.value * inputShape[j] %}
 {% endfor %}
 {% set count0_groups = outer_count.value * groups0 %}
 {% set count1_groups = outer_count.value * groups1 %}
 {% set max_groups = namespace(value=count0_groups) %}
 {% if count1_groups > max_groups.value %}{% set max_groups.value = count1_groups %}{% endif %}
 {% set start1_scalar = run0_scalar %}
+{% set run2_scalar = y2Shape[axisSpec] * inner.value %}
 {% set groups2 = (((run2_scalar + 3) / 4) | int) %}
 {% set count2_groups = outer_count.value * groups2 %}
 {% set start2_scalar = start1_scalar + run1_scalar %}
   }
 }
 {% else %}
+{% set inner_in = (((inputShape[axisSpec] * inner.value) / 4) | int) %}
+{% set run0 = (((y0Shape[axisSpec] * inner.value) / 4) | int) %}
+{% set run1 = (((y1Shape[axisSpec] * inner.value) / 4) | int) %}
 {% set has_nonempty_run = namespace(value=(run0 > 0 or run1 > 0)) %}
 {% set start1 = run0 %}
 {% set count0 = namespace(value=run0) %}
 {% set count1 = namespace(value=run1) %}
+{% for j in range(axisSpec) %}
+{% set count0.value = count0.value * inputShape[j] %}
+{% set count1.value = count1.value * inputShape[j] %}
 {% endfor %}
 {% set max_count = namespace(value=count0.value) %}
 {% if count1.value > max_count.value %}{% set max_count.value = count1.value %}{% endif %}
+{% if outputCountSpec >= 3 %}
+{% set run2 = (((y2Shape[axisSpec] * inner.value) / 4) | int) %}
 {% if run2 > 0 %}{% set has_nonempty_run.value = true %}{% endif %}
 {% set start2 = start1 + run1 %}
 {% set count2 = namespace(value=run2) %}
+{% for j in range(axisSpec) %}
+{% set count2.value = count2.value * inputShape[j] %}
 {% endfor %}
 {% if count2.value > max_count.value %}{% set max_count.value = count2.value %}{% endif %}
 {% endif %}
+{% if outputCountSpec >= 4 %}
+{% set run3 = (((y3Shape[axisSpec] * inner.value) / 4) | int) %}
 {% if run3 > 0 %}{% set has_nonempty_run.value = true %}{% endif %}
 {% set start3 = start2 + run2 %}
 {% set count3 = namespace(value=run3) %}
+{% for j in range(axisSpec) %}
+{% set count3.value = count3.value * inputShape[j] %}
 {% endfor %}
 {% if count3.value > max_count.value %}{% set max_count.value = count3.value %}{% endif %}
 {% endif %}
       y1[i] = input[(i / {{ run1 }}u) * INNER_IN + {{ start1 }}u + (i % {{ run1 }}u)];
     }
 {% endif %}
+{% if outputCountSpec >= 3 and run2 > 0 %}
     if (i < {{ count2.value }}u) {
       y2[i] = input[(i / {{ run2 }}u) * INNER_IN + {{ start2 }}u + (i % {{ run2 }}u)];
     }
 {% endif %}
+{% if outputCountSpec >= 4 and run3 > 0 %}
     if (i < {{ count3.value }}u) {
       y3[i] = input[(i / {{ run3 }}u) * INNER_IN + {{ start3 }}u + (i % {{ run3 }}u)];
     }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,74 +2,23 @@
   "domain": "ai.onnx",
   "name": "Split",
   "sinceVersion": 18,
-  "description": "Splits a tensor into a list of tensors along the specified `axis`. The sizes of each output slice are given by the optional `split` input; if omitted, the tensor is divided into equal parts (the last chunk may be smaller if the axis dimension is not evenly divisible).",
-  "inputs": [
-    { "role": "input", "dtype": "T", "description": "The tensor to split." },
-    {
-      "role": "split",
-      "dtype": "S",
-      "rank": 1,
-      "optional": true,
-      "description": "Optional logical int64 1-D tensor specifying the size of each output along the split axis; values must be non-negative, sum to the axis dimension, and use uint32 WebGPU storage."
-    }
-  ],
-  "outputs": [
-    { "role": "Y0", "dtype": "T", "description": "First output slice after splitting.", "rank": "ranks.input" },
-    {
-      "role": "Y1",
-      "dtype": "T",
-      "optional": true,
-      "description": "Second output slice after splitting (optional).",
-      "rank": "ranks.input"
-    },
-    {
-      "role": "Y2",
-      "dtype": "T",
-      "optional": true,
-      "description": "Third output slice after splitting.",
-      "rank": "ranks.input"
-    },
-    {
-      "role": "Y3",
-      "dtype": "T",
-      "optional": true,
-      "description": "Fourth output slice after splitting.",
-      "rank": "ranks.input"
-    },
-    {
-      "role": "Y4",
-      "dtype": "T",
-      "optional": true,
-      "description": "Fifth output slice after splitting.",
-      "rank": "ranks.input"
-    }
-  ],
-  "attributes": { "axis": 0 },
-  "attributeDescriptions": {
-    "axis": "The axis along which to split. Negative values count from the end; accepted range is `[-rank, rank-1]`.",
-    "num_outputs": "Optional number of outputs when the `split` input is omitted. The final output may be smaller when the axis dimension is not evenly divisible."
   },
   "typeConstraints": {
     "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"],
     "S": ["int64"]
   },
-  "args": {
-    "input": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "split": {
-      "kind": "tensor",
-      "semantic": "split",
-      "role": "split",
-      "dtype": "uint32",
-      "narrowing": "checked",
-      "required": false
-    },
-    "y0": { "kind": "tensor", "semantic": "Y0", "role": "output" },
-    "y1": { "kind": "tensor", "semantic": "Y1", "role": "output", "required": false },
-    "y2": { "kind": "tensor", "semantic": "Y2", "role": "output", "required": false },
-    "y3": { "kind": "tensor", "semantic": "Y3", "role": "output", "required": false },
-    "y4": { "kind": "tensor", "semantic": "Y4", "role": "output", "required": false }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256 },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
@@ -77,167 +26,67 @@
     "numOutputsContract": "not has(attrs, \"num_outputs\") or attrs.num_outputs == outputCount",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "baseContract": "ranks.input >= 1 and axis >= 0 and axis < ranks.input and numOutputsContract and f16Ok(dtypes.T)",
-    "oneOutputContract": "baseContract and not present.y1 and not present.y2 and not present.y3 and not present.y4 and sameShape(shapes.input, shapes.Y0)",
-    "twoOutputContract": "baseContract and present.y1 and not present.y2 and ranks.input == ranks.Y0 and ranks.input == ranks.Y1 and dim(shapes.input, attrs.axis) == dim(shapes.Y0, attrs.axis) + dim(shapes.Y1, attrs.axis)",
-    "threeOutputContract": "baseContract and present.y1 and present.y2 and not present.y3 and ranks.input == ranks.Y0 and ranks.input == ranks.Y1 and ranks.input == ranks.Y2 and dim(shapes.input, attrs.axis) == dim(shapes.Y0, attrs.axis) + dim(shapes.Y1, attrs.axis) + dim(shapes.Y2, attrs.axis)",
-    "fourOutputContract": "baseContract and present.y1 and present.y2 and present.y3 and not present.y4 and ranks.input == ranks.Y0 and ranks.input == ranks.Y1 and ranks.input == ranks.Y2 and ranks.input == ranks.Y3 and dim(shapes.input, attrs.axis) == dim(shapes.Y0, attrs.axis) + dim(shapes.Y1, attrs.axis) + dim(shapes.Y2, attrs.axis) + dim(shapes.Y3, attrs.axis)",
-    "fiveOutputContract": "baseContract and present.y1 and present.y2 and present.y3 and present.y4 and ranks.input == ranks.Y0 and ranks.input == ranks.Y1 and ranks.input == ranks.Y2 and ranks.input == ranks.Y3 and ranks.input == ranks.Y4 and dim(shapes.input, attrs.axis) == dim(shapes.Y0, attrs.axis) + dim(shapes.Y1, attrs.axis) + dim(shapes.Y2, attrs.axis) + dim(shapes.Y3, attrs.axis) + dim(shapes.Y4, attrs.axis)",
-    "twoBlockContract": "twoOutputContract and outer(shapes.Y0, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y0, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.Y1, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y1, attrs.axis) == inner(shapes.input, attrs.axis)",
-    "threeBlockContract": "threeOutputContract and outer(shapes.Y0, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y0, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.Y1, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y1, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.Y2, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y2, attrs.axis) == inner(shapes.input, attrs.axis)",
-    "fourBlockContract": "fourOutputContract and outer(shapes.Y0, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y0, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.Y1, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y1, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.Y2, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y2, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.Y3, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.Y3, attrs.axis) == inner(shapes.input, attrs.axis)"
   },
-  "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
-  "bindingSets": {
-    "oneOutput": [
-      {
-        "name": "src",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "dst", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$scalar" }
-    ],
-    "twoBlockVec4": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" }
-    ],
-    "threeBlock": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
-      { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
-      { "name": "y2", "arg": "y2", "semantic": "Y2", "buffer": { "type": "storage" }, "elementType": "$ioElement" }
-    ],
-    "fourBlockVec4": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      { "name": "y0", "arg": "y0", "semantic": "Y0", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      { "name": "y1", "arg": "y1", "semantic": "Y1", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      { "name": "y2", "arg": "y2", "semantic": "Y2", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
-      { "name": "y3", "arg": "y3", "semantic": "Y3", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" }
-    ],
-    "twoScalar": [
-      {
-        "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)" }
-          ]
-        }
-      }
-    ],
-    "fiveScalar": [
-      {
-        "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": "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)" }
-          ]
-        }
-      }
-    ],
-    "fourScalar": [
-      {
-        "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)" }
-          ]
-        }
-      }
-    ],
-    "threeScalar": [
-      {
-        "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)" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
@@ -248,9 +97,17 @@
         {
           "id": "main",
           "name": "Split.copy",
-          "source": { "shader": "datamove-flat-copy.wgsl.jinja", "inputs": { "count": "numel(shapes.Y0)" } },
-          "bindings": "oneOutput",
-          "dispatch": { "gridStride": "numel(shapes.Y0)", "workgroupSize": "workgroupSize" },
           "viewAlias": [{ "input": "src", "output": "dst" }]
         }
       ]
@@ -258,54 +115,55 @@
     {
       "id": "two_outputs_block_vec4",
       "priority": 15,
-      "when": ["twoBlockContract", "(dim(shapes.Y0, attrs.axis) * inner(shapes.Y0, attrs.axis)) % 4 == 0", "(dim(shapes.Y1, attrs.axis) * inner(shapes.Y1, attrs.axis)) % 4 == 0"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Split.blockVec4",
-          "source": {
-            "shader": "datamove-split-block.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "y0Shape": "shapes.Y0",
-              "y1Shape": "shapes.Y1",
-              "rank": "ranks.input",
-              "axis": "axis",
-              "outputCount": 2
-            }
           },
-          "bindings": "twoBlockVec4",
-          "dispatch": { "gridStride": "max(numel(shapes.Y0), numel(shapes.Y1)) / 4", "workgroupSize": "workgroupSize" }
         }
       ]
     },
     {
       "id": "three_outputs_block_scalar_x4",
       "priority": 14,
-      "when": ["threeBlockContract", "dim(shapes.Y0, attrs.axis) > 0", "dim(shapes.Y1, attrs.axis) > 0", "dim(shapes.Y2, attrs.axis) > 0", "max(numel(shapes.Y0), numel(shapes.Y1), numel(shapes.Y2)) >= 16"],
-      "constants": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Split3.blockScalarX4",
-          "source": {
-            "shader": "datamove-split-block.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "y0Shape": "shapes.Y0",
-              "y1Shape": "shapes.Y1",
-              "y2Shape": "shapes.Y2",
-              "rank": "ranks.input",
-              "axis": "axis",
-              "outputCount": 3,
-              "scalarBoundX4": true
-            }
           },
-          "bindings": "threeBlock",
           "dispatch": {
-            "gridStride": "max(outer(shapes.Y0, attrs.axis) * ceil(dim(shapes.Y0, attrs.axis) * inner(shapes.Y0, attrs.axis) / 4), outer(shapes.Y1, attrs.axis) * ceil(dim(shapes.Y1, attrs.axis) * inner(shapes.Y1, attrs.axis) / 4), outer(shapes.Y2, attrs.axis) * ceil(dim(shapes.Y2, attrs.axis) * inner(shapes.Y2, attrs.axis) / 4))",
-            "workgroupSize": "workgroupSize"
           }
         }
       ]
@@ -313,28 +171,27 @@
     {
       "id": "three_outputs_block_vec4",
       "priority": 15,
-      "when": ["threeBlockContract", "(dim(shapes.Y0, attrs.axis) * inner(shapes.Y0, attrs.axis)) % 4 == 0", "(dim(shapes.Y1, attrs.axis) * inner(shapes.Y1, attrs.axis)) % 4 == 0", "(dim(shapes.Y2, attrs.axis) * inner(shapes.Y2, attrs.axis)) % 4 == 0"],
-      "constants": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Split3.blockVec4",
-          "source": {
-            "shader": "datamove-split-block.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "y0Shape": "shapes.Y0",
-              "y1Shape": "shapes.Y1",
-              "y2Shape": "shapes.Y2",
-              "rank": "ranks.input",
-              "axis": "axis",
-              "outputCount": 3
-            }
           },
-          "bindings": "threeBlock",
           "dispatch": {
-            "gridStride": "max(numel(shapes.Y0), numel(shapes.Y1), numel(shapes.Y2)) / 4",
-            "workgroupSize": "workgroupSize"
           }
         }
       ]
@@ -342,29 +199,34 @@
     {
       "id": "four_outputs_block_vec4",
       "priority": 25,
-      "when": ["fourBlockContract", "(dim(shapes.Y0, attrs.axis) * inner(shapes.Y0, attrs.axis)) % 4 == 0", "(dim(shapes.Y1, attrs.axis) * inner(shapes.Y1, attrs.axis)) % 4 == 0", "(dim(shapes.Y2, attrs.axis) * inner(shapes.Y2, attrs.axis)) % 4 == 0", "(dim(shapes.Y3, attrs.axis) * inner(shapes.Y3, attrs.axis)) % 4 == 0"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Split4.blockVec4",
-          "source": {
-            "shader": "datamove-split-block.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "y0Shape": "shapes.Y0",
-              "y1Shape": "shapes.Y1",
-              "y2Shape": "shapes.Y2",
-              "y3Shape": "shapes.Y3",
-              "rank": "ranks.input",
-              "axis": "axis",
-              "outputCount": 4
-            }
           },
-          "bindings": "fourBlockVec4",
           "dispatch": {
-            "gridStride": "max(numel(shapes.Y0), numel(shapes.Y1), numel(shapes.Y2), numel(shapes.Y3)) / 4",
-            "workgroupSize": "workgroupSize"
           }
         }
       ]
@@ -376,17 +238,19 @@
         {
           "id": "main",
           "name": "Split",
-          "source": {
-            "shader": "split-n.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.Y0}, {\"name\": \"y1\", \"shape\": shapes.Y1} ]",
-              "rank": "ranks.input",
-              "axis": "axis"
-            }
           },
-          "bindings": "twoScalar",
-          "dispatch": { "threads": "max(numel(shapes.Y0), numel(shapes.Y1))", "workgroupSize": "workgroupSize" }
         }
       ]
     },
@@ -398,19 +262,18 @@
         {
           "id": "main",
           "name": "SplitN",
-          "source": {
-            "shader": "split-n.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.Y0}, {\"name\": \"y1\", \"shape\": shapes.Y1}, {\"name\": \"y2\", \"shape\": shapes.Y2}, {\"name\": \"y3\", \"shape\": shapes.Y3}, {\"name\": \"y4\", \"shape\": shapes.Y4} ]",
-              "rank": "ranks.input",
-              "axis": "axis"
-            }
           },
-          "bindings": "fiveScalar",
           "dispatch": {
-            "threads": "max(numel(shapes.Y0), numel(shapes.Y1), numel(shapes.Y2), numel(shapes.Y3), numel(shapes.Y4))",
-            "workgroupSize": "workgroupSize"
           }
         }
       ]
@@ -423,19 +286,18 @@
         {
           "id": "main",
           "name": "Split4",
-          "source": {
-            "shader": "split-n.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.Y0}, {\"name\": \"y1\", \"shape\": shapes.Y1}, {\"name\": \"y2\", \"shape\": shapes.Y2}, {\"name\": \"y3\", \"shape\": shapes.Y3} ]",
-              "rank": "ranks.input",
-              "axis": "axis"
-            }
           },
-          "bindings": "fourScalar",
           "dispatch": {
-            "threads": "max(numel(shapes.Y0), numel(shapes.Y1), numel(shapes.Y2), numel(shapes.Y3))",
-            "workgroupSize": "workgroupSize"
           }
         }
       ]
@@ -448,19 +310,18 @@
         {
           "id": "main",
           "name": "Split3",
-          "source": {
-            "shader": "split-n.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.input",
-              "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.Y0}, {\"name\": \"y1\", \"shape\": shapes.Y1}, {\"name\": \"y2\", \"shape\": shapes.Y2} ]",
-              "rank": "ranks.input",
-              "axis": "axis"
-            }
           },
-          "bindings": "threeScalar",
           "dispatch": {
-            "threads": "max(numel(shapes.Y0), numel(shapes.Y1), numel(shapes.Y2))",
-            "workgroupSize": "workgroupSize"
           }
         }
       ]

   "domain": "ai.onnx",
   "name": "Split",
   "sinceVersion": 18,
+  "inputs": {
+    "input": { "dtype": "T" },
+    "split": { "dtype": "S", "rank": 1, "optional": true, "storage": "uint32", "narrowing": "checked" }
   },
+  "outputs": {
+    "y0": { "onnx": "Y0", "dtype": "T", "rank": "ranks.input" },
+    "y1": { "onnx": "Y1", "dtype": "T", "rank": "ranks.input", "optional": true },
+    "y2": { "onnx": "Y2", "dtype": "T", "rank": "ranks.input", "optional": true },
+    "y3": { "onnx": "Y3", "dtype": "T", "rank": "ranks.input", "optional": true },
+    "y4": { "onnx": "Y4", "dtype": "T", "rank": "ranks.input", "optional": true }
+  },
+  "attributes": { "axis": { "default": 0 }, "num_outputs": {} },
   "typeConstraints": {
     "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"],
     "S": ["int64"]
   },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.input",
     "numOutputsContract": "not has(attrs, \"num_outputs\") or attrs.num_outputs == outputCount",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "baseContract": "ranks.input >= 1 and axis >= 0 and axis < ranks.input and numOutputsContract and f16Ok(dtypes.T)",
+    "oneOutputContract": "baseContract and not present.y1 and not present.y2 and not present.y3 and not present.y4 and sameShape(shapes.input, shapes.y0)",
+    "twoOutputContract": "baseContract and present.y1 and not present.y2 and ranks.input == ranks.y0 and ranks.input == ranks.y1 and dim(shapes.input, attrs.axis) == dim(shapes.y0, attrs.axis) + dim(shapes.y1, attrs.axis)",
+    "threeOutputContract": "baseContract and present.y1 and present.y2 and not present.y3 and ranks.input == ranks.y0 and ranks.input == ranks.y1 and ranks.input == ranks.y2 and dim(shapes.input, attrs.axis) == dim(shapes.y0, attrs.axis) + dim(shapes.y1, attrs.axis) + dim(shapes.y2, attrs.axis)",
+    "fourOutputContract": "baseContract and present.y1 and present.y2 and present.y3 and not present.y4 and ranks.input == ranks.y0 and ranks.input == ranks.y1 and ranks.input == ranks.y2 and ranks.input == ranks.y3 and dim(shapes.input, attrs.axis) == dim(shapes.y0, attrs.axis) + dim(shapes.y1, attrs.axis) + dim(shapes.y2, attrs.axis) + dim(shapes.y3, attrs.axis)",
+    "fiveOutputContract": "baseContract and present.y1 and present.y2 and present.y3 and present.y4 and ranks.input == ranks.y0 and ranks.input == ranks.y1 and ranks.input == ranks.y2 and ranks.input == ranks.y3 and ranks.input == ranks.y4 and dim(shapes.input, attrs.axis) == dim(shapes.y0, attrs.axis) + dim(shapes.y1, attrs.axis) + dim(shapes.y2, attrs.axis) + dim(shapes.y3, attrs.axis) + dim(shapes.y4, attrs.axis)",
+    "twoBlockContract": "twoOutputContract and outer(shapes.y0, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y0, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.y1, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y1, attrs.axis) == inner(shapes.input, attrs.axis)",
+    "threeBlockContract": "threeOutputContract and outer(shapes.y0, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y0, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.y1, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y1, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.y2, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y2, attrs.axis) == inner(shapes.input, attrs.axis)",
+    "fourBlockContract": "fourOutputContract and outer(shapes.y0, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y0, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.y1, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y1, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.y2, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y2, attrs.axis) == inner(shapes.input, attrs.axis) and outer(shapes.y3, attrs.axis) == outer(shapes.input, attrs.axis) and inner(shapes.y3, attrs.axis) == inner(shapes.input, attrs.axis)",
+    "scalar": "dtypes.T"
   },
+  "bindings": {
+    "input": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "y0": { "buffer": "storage", "elementType": "$vectorScalar" },
+    "y1": { "buffer": "storage", "elementType": "$vectorScalar" },
+    "input_2": { "name": "input", "buffer": "read-only-storage", "elementType": "$ioElement" },
+    "y0_2": { "name": "y0", "buffer": "storage", "elementType": "$ioElement" },
+    "y1_2": { "name": "y1", "buffer": "storage", "elementType": "$ioElement" },
+    "y2": { "buffer": "storage", "elementType": "$ioElement" },
+    "input_3": { "name": "input", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "y0_3": { "name": "y0", "buffer": "storage", "elementType": "$scalar" },
+    "y1_3": { "name": "y1", "buffer": "storage", "elementType": "$scalar" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "y0Count", "type": "u32", "value": "numel(shapes.y0)" },
+        { "name": "y1Count", "type": "u32", "value": "numel(shapes.y1)" }
+      ]
+    },
+    "y2_3": { "name": "y2", "buffer": "storage", "elementType": "$scalar" },
+    "y3_2": { "name": "y3", "buffer": "storage", "elementType": "$scalar" },
+    "y4": { "buffer": "storage", "elementType": "$scalar" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "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)" }
+      ]
+    },
+    "params_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "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)" }
+      ]
+    },
+    "params_4": {
+      "name": "params",
+      "buffer": "uniform",
+      "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)" }
+      ]
+    }
   },
   "variants": [
     {
         {
           "id": "main",
           "name": "Split.copy",
+          "shader": "datamove-flat-copy.wgsl.jinja",
+          "derive": { "count": "numel(shapes.y0)" },
+          "bindings": [
+            { "arg": "input", "name": "src", "elementType": "$scalar" },
+            { "arg": "y0", "name": "dst", "elementType": "$scalar" }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y0)), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          },
           "viewAlias": [{ "input": "src", "output": "dst" }]
         }
       ]
     {
       "id": "two_outputs_block_vec4",
       "priority": 15,
+      "when": ["twoBlockContract", "(dim(shapes.y0, attrs.axis) * inner(shapes.y0, attrs.axis)) % 4 == 0", "(dim(shapes.y1, attrs.axis) * inner(shapes.y1, attrs.axis)) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Split.blockVec4",
+          "shader": "datamove-split-block.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "y0Shape": "shapes.y0",
+            "y1Shape": "shapes.y1",
+            "rank": "ranks.input",
+            "axisSpec": "axis",
+            "outputCountSpec": 2
           },
+          "bindings": ["input", "y0", "y1"],
+          "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1)) / 4), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "three_outputs_block_scalar_x4",
       "priority": 14,
+      "when": ["threeBlockContract", "dim(shapes.y0, attrs.axis) > 0", "dim(shapes.y1, attrs.axis) > 0", "dim(shapes.y2, attrs.axis) > 0", "max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2)) >= 16"],
+      "derive": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Split3.blockScalarX4",
+          "shader": "datamove-split-block.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "y0Shape": "shapes.y0",
+            "y1Shape": "shapes.y1",
+            "y2Shape": "shapes.y2",
+            "rank": "ranks.input",
+            "axisSpec": "axis",
+            "outputCountSpec": 3,
+            "scalarBoundX4": true
           },
+          "bindings": ["input_2", "y0_2", "y1_2", "y2"],
           "dispatch": {
+            "x": "min(ceilDiv((max(outer(shapes.y0, attrs.axis) * ceil(dim(shapes.y0, attrs.axis) * inner(shapes.y0, attrs.axis) / 4), outer(shapes.y1, attrs.axis) * ceil(dim(shapes.y1, attrs.axis) * inner(shapes.y1, attrs.axis) / 4), outer(shapes.y2, attrs.axis) * ceil(dim(shapes.y2, attrs.axis) * inner(shapes.y2, attrs.axis) / 4))), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
           }
         }
       ]
     {
       "id": "three_outputs_block_vec4",
       "priority": 15,
+      "when": ["threeBlockContract", "(dim(shapes.y0, attrs.axis) * inner(shapes.y0, attrs.axis)) % 4 == 0", "(dim(shapes.y1, attrs.axis) * inner(shapes.y1, attrs.axis)) % 4 == 0", "(dim(shapes.y2, attrs.axis) * inner(shapes.y2, attrs.axis)) % 4 == 0"],
+      "derive": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Split3.blockVec4",
+          "shader": "datamove-split-block.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "y0Shape": "shapes.y0",
+            "y1Shape": "shapes.y1",
+            "y2Shape": "shapes.y2",
+            "rank": "ranks.input",
+            "axisSpec": "axis",
+            "outputCountSpec": 3
           },
+          "bindings": ["input_2", "y0_2", "y1_2", "y2"],
           "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2)) / 4), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
           }
         }
       ]
     {
       "id": "four_outputs_block_vec4",
       "priority": 25,
+      "when": ["fourBlockContract", "(dim(shapes.y0, attrs.axis) * inner(shapes.y0, attrs.axis)) % 4 == 0", "(dim(shapes.y1, attrs.axis) * inner(shapes.y1, attrs.axis)) % 4 == 0", "(dim(shapes.y2, attrs.axis) * inner(shapes.y2, attrs.axis)) % 4 == 0", "(dim(shapes.y3, attrs.axis) * inner(shapes.y3, attrs.axis)) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Split4.blockVec4",
+          "shader": "datamove-split-block.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "y0Shape": "shapes.y0",
+            "y1Shape": "shapes.y1",
+            "y2Shape": "shapes.y2",
+            "y3Shape": "shapes.y3",
+            "rank": "ranks.input",
+            "axisSpec": "axis",
+            "outputCountSpec": 4
           },
+          "bindings": [
+            "input",
+            "y0",
+            "y1",
+            { "arg": "y2", "elementType": "$vectorScalar" },
+            { "arg": "y3", "elementType": "$vectorScalar" }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2), numel(shapes.y3)) / 4), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "Split",
+          "shader": "split-n.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1} ]",
+            "rank": "ranks.input",
+            "axisSpec": "axis"
           },
+          "bindings": ["input_3", "y0_3", "y1_3", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1))), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(numel(shapes.y0), numel(shapes.y1))), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
         {
           "id": "main",
           "name": "SplitN",
+          "shader": "split-n.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1}, {\"name\": \"y2\", \"shape\": shapes.y2}, {\"name\": \"y3\", \"shape\": shapes.y3}, {\"name\": \"y4\", \"shape\": shapes.y4} ]",
+            "rank": "ranks.input",
+            "axisSpec": "axis"
           },
+          "bindings": ["input_3", "y0_3", "y1_3", "y2_3", "y3_2", "y4", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2), numel(shapes.y3), numel(shapes.y4))), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2), numel(shapes.y3), numel(shapes.y4))), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "Split4",
+          "shader": "split-n.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1}, {\"name\": \"y2\", \"shape\": shapes.y2}, {\"name\": \"y3\", \"shape\": shapes.y3} ]",
+            "rank": "ranks.input",
+            "axisSpec": "axis"
           },
+          "bindings": ["input_3", "y0_3", "y1_3", "y2_3", "y3_2", "params_3"],
           "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2), numel(shapes.y3))), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2), numel(shapes.y3))), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "Split3",
+          "shader": "split-n.wgsl.jinja",
+          "derive": {
+            "inputShape": "shapes.input",
+            "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1}, {\"name\": \"y2\", \"shape\": shapes.y2} ]",
+            "rank": "ranks.input",
+            "axisSpec": "axis"
           },
+          "bindings": ["input_3", "y0_3", "y1_3", "y2_3", "params_4"],
           "dispatch": {
+            "x": "min(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2))), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((max(numel(shapes.y0), numel(shapes.y1), numel(shapes.y2))), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,33 @@
 {
   "name": "ai.onnx.Split",
-  "id": "_ai_onnx_split_webgpu_fa39047",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "2lxxmh4jHqr1stKdWRuW4p+XgseMJNCWKwNsmypIqeM=",
-      "datamove-flat-copy.wgsl.jinja": "rEpeYMvIZQuQnB2f3lt51jrrklKICsqUTzCJPRTknNk=",
-      "datamove-split-block.wgsl.jinja": "BMcc7EVTp6zQzIDQBkTcgdYPqFaNCf1cUTox/l1kTxg=",
-      "manifest.json": "6wweSEZ8+U9uXmLy9VX0MtrQfDybb8JlpcRfElWwMKo=",
-      "split-n.wgsl.jinja": "GjmhWg4EmgRJccL6Y0Ll0FAvbeKn3e0JwbX2VP5zrEo=",
-      "test.json": "z07mxLLaj3ydKbjwPvFvi45EwMDYHoeTdVKoez0SnFM="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Split" }
 }

 {
   "name": "ai.onnx.Split",
+  "id": "_ai_onnx_split_webgpu_905cd69",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "k7jAXfnplRAPstbkQGbKdFPxoWX7R9DW+MK+4NiYTLU=",
+      "datamove-flat-copy.wgsl.jinja": "g9d62mer2bmHfbScqX5CIT5Zq/Pdinac0W2dH0CD+5s=",
+      "datamove-split-block.wgsl.jinja": "shWWQMUtwoSmMEgOn2yy1spChp86zzPkm/Cb7QiAnWs=",
+      "manifest.json": "qnPxFTTkmu2fktjom8/oaTjlNXRcGD2kg+oeCsJ6RQc=",
+      "split-n.wgsl.jinja": "z574+TJ8XGpxhV0wwxRfsHXSsSu7Ry4dvhXBxHcB6Cw=",
+      "test.json": "71+a2DHyXRRHacOOABnDWrpe60/NRk7A/jP9UqPwpFE="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "one_output_copy": ["datamove-flat-copy.wgsl.jinja"],
+      "two_outputs_block_vec4": ["datamove-split-block.wgsl.jinja"],
+      "three_outputs_block_scalar_x4": ["datamove-split-block.wgsl.jinja"],
+      "three_outputs_block_vec4": ["datamove-split-block.wgsl.jinja"],
+      "four_outputs_block_vec4": ["datamove-split-block.wgsl.jinja"],
+      "two_outputs": ["split-n.wgsl.jinja"],
+      "five_outputs": ["split-n.wgsl.jinja"],
+      "four_outputs": ["split-n.wgsl.jinja"],
+      "three_outputs": ["split-n.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/split-n.wgsl.jinja CHANGED Viewed

@@ -2,19 +2,16 @@
 // adding all preceding outputs' cumulative split-axis extent. One invocation
 // handles the same flat position across outputs, and each output writes only
 // when that position is within its own element count.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-{% for output in source.outputs %}
 {% set out_count = namespace(value=1) %}
 {% for d in output.shape %}
 {% set out_count.value = out_count.value * d %}
 {% endfor %}
 {% set prefix = namespace(value=0) %}
-{% for prev in source.outputs[:loop.index0] %}
-{% set prefix.value = prefix.value + prev.shape[source.axis] %}
 {% endfor %}
 fn input_offset_{{ output.name }}({% if out_count.value > 0 %}out_index: u32{% endif %}) -> u32 {
 {% if out_count.value == 0 %}
@@ -22,18 +19,18 @@ fn input_offset_{{ output.name }}({% if out_count.value > 0 %}out_index: u32{% e
 {% else %}
   var rem = out_index;
   var offset = 0u;
-{% for axis in range(source.rank) %}
 {% set out_stride = namespace(value=1) %}
-{% for j in range(axis + 1, source.rank) %}
 {% set out_stride.value = out_stride.value * output.shape[j] %}
 {% endfor %}
 {% set input_stride = namespace(value=1) %}
-{% for j in range(axis + 1, source.rank) %}
-{% set input_stride.value = input_stride.value * source.inputShape[j] %}
 {% endfor %}
   var coord{{ axis }} = rem / {{ out_stride.value }}u;
   rem = rem % {{ out_stride.value }}u;
-{% if axis == source.axis and prefix.value > 0 %}
   coord{{ axis }} = coord{{ axis }} + {{ prefix.value }}u;
 {% endif %}
   offset = offset + coord{{ axis }} * {{ input_stride.value }}u;
@@ -44,11 +41,11 @@ fn input_offset_{{ output.name }}({% if out_count.value > 0 %}out_index: u32{% e
 {% endfor %}
 @compute @workgroup_size({{ workgroupSize }})
-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 limit.
-  let i = gid.x + gid.y * nwg.x * {{ workgroupSize }}u;
-{% for output in source.outputs %}
 {% set out_count = namespace(value=1) %}
 {% for d in output.shape %}
 {% set out_count.value = out_count.value * d %}

 // adding all preceding outputs' cumulative split-axis extent. One invocation
 // handles the same flat position across outputs, and each output writes only
 // when that position is within its own element count.
 {{ env.wgsl.resourceDeclarations }}
+{% for output in outputs %}
 {% set out_count = namespace(value=1) %}
 {% for d in output.shape %}
 {% set out_count.value = out_count.value * d %}
 {% endfor %}
 {% set prefix = namespace(value=0) %}
+{% for prev in outputs[:loop.index0] %}
+{% set prefix.value = prefix.value + prev.shape[axisSpec] %}
 {% endfor %}
 fn input_offset_{{ output.name }}({% if out_count.value > 0 %}out_index: u32{% endif %}) -> u32 {
 {% if out_count.value == 0 %}
 {% else %}
   var rem = out_index;
   var offset = 0u;
+{% for axis in range(rank) %}
 {% set out_stride = namespace(value=1) %}
+{% for j in range(axis + 1, rank) %}
 {% set out_stride.value = out_stride.value * output.shape[j] %}
 {% endfor %}
 {% set input_stride = namespace(value=1) %}
+{% for j in range(axis + 1, rank) %}
+{% set input_stride.value = input_stride.value * inputShape[j] %}
 {% endfor %}
   var coord{{ axis }} = rem / {{ out_stride.value }}u;
   rem = rem % {{ out_stride.value }}u;
+{% if axis == axisSpec and prefix.value > 0 %}
   coord{{ axis }} = coord{{ axis }} + {{ prefix.value }}u;
 {% endif %}
   offset = offset + coord{{ axis }} * {{ input_stride.value }}u;
 {% endfor %}
 @compute @workgroup_size({{ workgroupSize }})
+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.
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ workgroupSize }}u;
+{% for output in outputs %}
 {% set out_count = namespace(value=1) %}
 {% for d in output.shape %}
 {% set out_count.value = out_count.value * d %}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Split",
   "fixtureArrays": {
     "ort_axis2_equal_three_outputs_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],
     "ort_axis1_middle_dimension_unequal_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, 25, 26, 27, 28, 29, 30, 31, 32]
@@ -655,7 +654,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/split_op_test.cc",
         "test": "SplitOperatorTest.Uint8NegativeAxis",
-        "notes": "Logical uint8 payload values are stored one per uint32 slot in this framework."
       },
       "attrs": { "axis": -1 },
       "inputs": {
@@ -675,7 +674,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/split_op_test.cc",
         "test": "SplitOperatorTest.Uint8Axis1SplitMiddleDimensionUnequally",
-        "notes": "Logical uint8 payload values are stored one per uint32 slot in this framework."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -848,7 +847,7 @@
       "outputs": { "y0": { "dtype": "float32", "shape": [2] }, "y1": { "dtype": "float32", "shape": [4] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_variable_parts_1d_opset18",
-        "notes": "ONNX int64 tensors use framework storage where representable. The corresponding opset-13 backend vector is execution-identical on this package surface."
       }
     },
     {
@@ -865,7 +864,7 @@
       "outputs": { "y0": { "dtype": "float32", "shape": [2, 2] }, "y1": { "dtype": "float32", "shape": [2, 4] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_variable_parts_2d_opset18",
-        "notes": "ONNX int64 tensors use framework storage where representable. The corresponding opset-13 backend vector is execution-identical on this package surface."
       }
     },
     {
@@ -881,7 +880,7 @@
       "outputs": { "y0": { "dtype": "float32", "shape": [2] }, "y1": { "dtype": "float32", "shape": [4] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_variable_parts_default_axis_opset18",
-        "notes": "ONNX int64 tensors use framework storage where representable. The corresponding opset-13 backend vector is execution-identical on this package surface."
       }
     },
     {
@@ -897,7 +896,7 @@
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_zero_size_splits_opset18",
-        "notes": "ONNX int64 tensors use framework storage where representable. The corresponding opset-13 backend vector is execution-identical on this package surface."
       }
     },
     {
@@ -944,7 +943,7 @@
     {
       "name": "rank7_middle_axis_two_outputs",
       "provenance": {
-        "notes": "Compact correctness sibling for the rank7 middle-axis split benchmark; exercises high-rank address arithmetic with an equal two-way split on axis 2."
       },
       "attrs": { "axis": 2 },
       "inputs": {
@@ -1266,7 +1265,7 @@
     {
       "name": "block_scalar_x4_three_outputs_y2_dominant_groups",
       "provenance": {
-        "notes": "Three-output block split whose runs are 2/3/6 scalars — none a multiple of 4 — so the vec4 block variant is ineligible and the scalar-bound-x4 block variant renders. Y2 needs two four-scalar groups per outer row while Y0/Y1 need one each, so the third output's group count is the strict maximum and it, not Y0/Y1, sets the loop bound. Plain ONNX Split semantics: axis=1 with split=[2,3,6] summing to the axis dim 11."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1296,7 +1295,7 @@
     {
       "name": "block_vec4_four_outputs_axis1_y3_dominant_count",
       "provenance": {
-        "notes": "Four-output vec4 block split on a middle axis with strictly growing chunks 2/2/4/6 (rows of inner 2, so 4/4/8/12 scalars, all vec4-aligned). The only four-output vec4 fixture until now split axis 0 into four equal chunks, so nothing exercised a non-zero axis (outer slices > 1 with a per-output row stride) nor a last output whose element count is the strict maximum over all four. Plain ONNX Split semantics: axis=1 with split=[2,2,4,6] summing to the axis dim 14."
       },
       "attrs": { "axis": 1 },
       "inputs": {
@@ -1316,7 +1315,7 @@
     {
       "name": "block_vec4_four_outputs_empty_y3",
       "provenance": {
-        "notes": "Four-output vec4 block split whose last chunk is empty (split=[4,4,4,0]): ONNX allows a zero-length entry in `split` as long as the entries are non-negative and sum to the axis dim, and the matching empty output must still be produced. The existing empty-chunk vec4 fixture empties Y0 in a two-output split, so no fixture ever gave the four-output block path a zero-length final run — the case where the fourth copy must be elided entirely instead of reading past the input. Values are a 13-long cycle over 12-wide rows so no repeat aligns with a row or chunk boundary."
       },
       "attrs": { "axis": 1 },
       "inputs": {

 {
   "fixtureArrays": {
     "ort_axis2_equal_three_outputs_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],
     "ort_axis1_middle_dimension_unequal_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, 25, 26, 27, 28, 29, 30, 31, 32]
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/split_op_test.cc",
         "test": "SplitOperatorTest.Uint8NegativeAxis",
+        "notes": "Logical uint8 payload values are stored one per uint32 slot."
       },
       "attrs": { "axis": -1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/split_op_test.cc",
         "test": "SplitOperatorTest.Uint8Axis1SplitMiddleDimensionUnequally",
+        "notes": "Logical uint8 payload values are stored one per uint32 slot."
       },
       "attrs": { "axis": 1 },
       "inputs": {
       "outputs": { "y0": { "dtype": "float32", "shape": [2] }, "y1": { "dtype": "float32", "shape": [4] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_variable_parts_1d_opset18",
+        "notes": "This WebGPU package stores representable ONNX int64 tensor values in int32 slots."
       }
     },
     {
       "outputs": { "y0": { "dtype": "float32", "shape": [2, 2] }, "y1": { "dtype": "float32", "shape": [2, 4] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_variable_parts_2d_opset18",
+        "notes": "This WebGPU package stores representable ONNX int64 tensor values in int32 slots."
       }
     },
     {
       "outputs": { "y0": { "dtype": "float32", "shape": [2] }, "y1": { "dtype": "float32", "shape": [4] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_variable_parts_default_axis_opset18",
+        "notes": "This WebGPU package stores representable ONNX int64 tensor values in int32 slots."
       }
     },
     {
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_split_zero_size_splits_opset18",
+        "notes": "This WebGPU package stores representable ONNX int64 tensor values in int32 slots."
       }
     },
     {
     {
       "name": "rank7_middle_axis_two_outputs",
       "provenance": {
+        "notes": "A rank-7 tensor split equally along middle axis 2 exercises high-rank address arithmetic."
       },
       "attrs": { "axis": 2 },
       "inputs": {
     {
       "name": "block_scalar_x4_three_outputs_y2_dominant_groups",
       "provenance": {
+        "notes": "Splitting axis 1 into runs of 2, 3, and 6 scalars selects grouped scalar copying because no run is four-aligned. The third output needs two four-scalar groups per row and therefore determines the loop bound."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "block_vec4_four_outputs_axis1_y3_dominant_count",
       "provenance": {
+        "notes": "A four-output vec4 block split on axis 1 uses chunks 2/2/4/6 with inner=2, giving 4/4/8/12 aligned scalars per outer slice. Multiple outer slices exercise each output row stride, and the final output has the largest element count."
       },
       "attrs": { "axis": 1 },
       "inputs": {
     {
       "name": "block_vec4_four_outputs_empty_y3",
       "provenance": {
+        "notes": "A four-output vectorized block split uses lengths [4,4,4,0]. ONNX permits the zero-length final chunk, which must be emitted without a copy or out-of-bounds input read. A 13-value cycle over 12-wide rows avoids alignment between values, rows, and chunk boundaries."
       },
       "attrs": { "axis": 1 },
       "inputs": {