| { |
| "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", "int64"], |
| "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", |
| "outputCount": "1 + (1 if present.y1 else 0) + (1 if present.y2 else 0) + (1 if present.y3 else 0) + (1 if present.y4 else 0)", |
| "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", |
| "inputShape": "shapes.input", |
| "rank": "ranks.input", |
| "axisSpec": "axis" |
| }, |
| "bindings": { |
| "input": { "elementType": "$vectorScalar" }, |
| "y0": { "elementType": "$vectorScalar" }, |
| "y1": { "elementType": "$vectorScalar" }, |
| "input_main": { "name": "input", "elementType": "$ioElement" }, |
| "y0_main": { "name": "y0", "elementType": "$ioElement" }, |
| "y1_main": { "name": "y1", "elementType": "$ioElement" }, |
| "y2": { "elementType": "$ioElement" }, |
| "input_scalar": { "name": "input", "elementType": "$scalar" }, |
| "y0_scalar": { "name": "y0", "elementType": "$scalar" }, |
| "y1_scalar": { "name": "y1", "elementType": "$scalar" }, |
| "params": { |
| "struct": [ |
| { "name": "y0Count", "type": "u32", "value": "numel(shapes.y0)" }, |
| { "name": "y1Count", "type": "u32", "value": "numel(shapes.y1)" } |
| ] |
| }, |
| "y2_main": { "name": "y2", "elementType": "$scalar" }, |
| "y3_main": { "name": "y3", "elementType": "$scalar" }, |
| "y4": { "elementType": "$scalar" }, |
| "params_main": { |
| "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)" } |
| ] |
| }, |
| "params_split_n": { |
| "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)" } |
| ] |
| }, |
| "params__uniform": { |
| "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)" } |
| ] |
| } |
| }, |
| "variants": [ |
| { |
| "id": "one_output_copy", |
| "priority": 30, |
| "when": ["oneOutputContract"], |
| "passes": [ |
| { |
| "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": ["dtypes.T != \"vec2<u32>\"", "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": { "y0Shape": "shapes.y0", "y1Shape": "shapes.y1", "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": ["dtypes.T != \"vec2<u32>\"", "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": { |
| "y0Shape": "shapes.y0", |
| "y1Shape": "shapes.y1", |
| "y2Shape": "shapes.y2", |
| "outputCountSpec": 3, |
| "scalarBoundX4": true |
| }, |
| "bindings": ["input_main", "y0_main", "y1_main", "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": ["dtypes.T != \"vec2<u32>\"", "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 ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "Split3.BlockVec4", |
| "shader": "datamove-split-block.wgsl.jinja", |
| "derive": { "y0Shape": "shapes.y0", "y1Shape": "shapes.y1", "y2Shape": "shapes.y2", "outputCountSpec": 3 }, |
| "bindings": ["input_main", "y0_main", "y1_main", "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": ["dtypes.T != \"vec2<u32>\"", "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": { |
| "y0Shape": "shapes.y0", |
| "y1Shape": "shapes.y1", |
| "y2Shape": "shapes.y2", |
| "y3Shape": "shapes.y3", |
| "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": "two_outputs", |
| "when": ["twoOutputContract"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "Split", |
| "shader": "split-n.wgsl.jinja", |
| "derive": { |
| "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1} ]" |
| }, |
| "bindings": ["input_scalar", "y0_scalar", "y1_scalar", "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": "five_outputs", |
| "priority": 21, |
| "when": ["fiveOutputContract"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SplitN", |
| "shader": "split-n.wgsl.jinja", |
| "derive": { |
| "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} ]" |
| }, |
| "bindings": ["input_scalar", "y0_scalar", "y1_scalar", "y2_main", "y3_main", "y4", "params_main"], |
| "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": "four_outputs", |
| "priority": 20, |
| "when": ["fourOutputContract"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "Split4", |
| "shader": "split-n.wgsl.jinja", |
| "derive": { |
| "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1}, {\"name\": \"y2\", \"shape\": shapes.y2}, {\"name\": \"y3\", \"shape\": shapes.y3} ]" |
| }, |
| "bindings": ["input_scalar", "y0_scalar", "y1_scalar", "y2_main", "y3_main", "params_split_n"], |
| "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": "three_outputs", |
| "priority": 10, |
| "when": ["threeOutputContract"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "Split3", |
| "shader": "split-n.wgsl.jinja", |
| "derive": { |
| "outputs": "[ {\"name\": \"y0\", \"shape\": shapes.y0}, {\"name\": \"y1\", \"shape\": shapes.y1}, {\"name\": \"y2\", \"shape\": shapes.y2} ]" |
| }, |
| "bindings": ["input_scalar", "y0_scalar", "y1_scalar", "y2_main", "params__uniform"], |
| "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 |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
