| { |
| "domain": "ai.onnx", |
| "name": "TensorScatter", |
| "sinceVersion": 24, |
| "inputs": { |
| "past": { "onnx": "past_cache", "dtype": "T" }, |
| "update": { "dtype": "T" }, |
| "writeIndices": { |
| "onnx": "write_indices", |
| "dtype": "I", |
| "rank": 1, |
| "optional": true, |
| "storage": "uint32", |
| "narrowing": "checked" |
| } |
| }, |
| "outputs": { "present": { "onnx": "present_cache", "dtype": "T", "rank": "ranks.past", "shape": "shapes.past" } }, |
| "attributes": { "axis": { "default": -2 }, "mode": { "default": "linear" } }, |
| "attributeConstraints": { "mode": { "values": ["linear", "circular"] } }, |
| "typeConstraints": { |
| "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"], |
| "I": ["int64"] |
| }, |
| "tunables": { "WORKGROUP_SIZE": { "default": 256 } }, |
| "when": ["(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) < ranks.past", "(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) >= 1", "f16Ok(dtypes.T)", "numel(shapes.past) == numel(shapes.present)", "ranks.past == ranks.present", "ranks.past == ranks.update", "ranks.past >= 2"], |
| "bindings": { |
| "update": { "elementType": "$scalar" }, |
| "present": { "elementType": "$scalar" }, |
| "params": { "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.update)" }] }, |
| "writeIndices": { "elementType": "u32" } |
| }, |
| "variants": [ |
| { |
| "id": "no_write_indices", |
| "priority": 0, |
| "when": ["not present.writeIndices"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "copyElementType": "(\"vec4<\" ~ dtypes.T ~ \">\") if numel(shapes.present) % 4 == 0 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "copy", |
| "name": "TensorScatter.Copy", |
| "shader": "scatter-flat-copy.wgsl.jinja", |
| "derive": { |
| "count": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))", |
| "copyWidth": "16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4)", |
| "elementCount": "numel(shapes.present)", |
| "scalarTailSafe": "numel(shapes.present) % 4 != 0" |
| }, |
| "bindings": [ |
| { "arg": "past", "name": "src", "elementType": "$copyElementType" }, |
| { "arg": "present", "name": "dst", "elementType": "$copyElementType" } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| }, |
| "viewAlias": [{ "input": "src", "output": "dst" }] |
| }, |
| { |
| "id": "scatter", |
| "name": "TensorScatter.Scatter", |
| "shader": "tensor-scatter.wgsl.jinja", |
| "derive": { |
| "pastShape": "shapes.past", |
| "updateShape": "shapes.update", |
| "rank": "ranks.past", |
| "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past", |
| "circular": "attrs.mode == \"circular\"", |
| "hasIndices": "false" |
| }, |
| "bindings": ["update", "present", "params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "with_write_indices", |
| "priority": 20, |
| "when": ["present.writeIndices", "dtypes.I == \"u32\""], |
| "derive": { |
| "scalar": "dtypes.T", |
| "copyElementType": "(\"vec4<\" ~ dtypes.T ~ \">\") if numel(shapes.present) % 4 == 0 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "copy", |
| "name": "TensorScatter.Copy", |
| "shader": "scatter-flat-copy.wgsl.jinja", |
| "derive": { |
| "count": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))", |
| "copyWidth": "16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4)", |
| "elementCount": "numel(shapes.present)", |
| "scalarTailSafe": "numel(shapes.present) % 4 != 0" |
| }, |
| "bindings": [ |
| { "arg": "past", "name": "src", "elementType": "$copyElementType" }, |
| { "arg": "present", "name": "dst", "elementType": "$copyElementType" } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| }, |
| "viewAlias": [{ "input": "src", "output": "dst" }] |
| }, |
| { |
| "id": "scatter", |
| "name": "TensorScatter.Scatter", |
| "shader": "tensor-scatter.wgsl.jinja", |
| "derive": { |
| "pastShape": "shapes.past", |
| "updateShape": "shapes.update", |
| "rank": "ranks.past", |
| "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past", |
| "circular": "attrs.mode == \"circular\"", |
| "hasIndices": "true" |
| }, |
| "bindings": ["update", "writeIndices", "present", "params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
