| { |
| "domain": "com.microsoft", |
| "name": "GatherBlockQuantized", |
| "sinceVersion": 1, |
| "inputs": { |
| "dataT": { "onnx": "data", "dtype": "T1", "rank": 2 }, |
| "indicesT": { "onnx": "indices", "dtype": "Tind", "rank": 1, "storage": "uint32", "narrowing": "checked" }, |
| "scalesT": { "onnx": "scales", "dtype": "T2", "rank": 2 }, |
| "zeroPointsT": { "onnx": "zero_points", "dtype": "T1", "rank": 2, "optional": true } |
| }, |
| "outputs": { |
| "outputT": { |
| "onnx": "output", |
| "dtype": "T2", |
| "rank": 2, |
| "shape": "[dim(shapes.indicesT, 0), dim(shapes.dataT, 1) * (8 / attrs.bits)]" |
| } |
| }, |
| "attributes": { |
| "bits": { "default": 4 }, |
| "block_size": { "default": 128 }, |
| "gather_axis": { "default": 0 }, |
| "quantize_axis": { "default": 1 } |
| }, |
| "attributeConstraints": { |
| "bits": { "values": [4, 8] }, |
| "gather_axis": { "values": [0] }, |
| "quantize_axis": { "values": [1] } |
| }, |
| "typeConstraints": { "T1": ["uint8"], "T2": ["float32"], "Tind": ["int64"] }, |
| "tunables": { "WORKGROUP_SIZE": { "default": 64 } }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "bits": "attrs.bits", |
| "blockSize": "attrs.block_size", |
| "blockSizeOk": "blockSize >= 16 and pow2ceil(blockSize) == blockSize", |
| "outBlocks": "ceilDiv(dim(shapes.outputT, 1), blockSize)", |
| "zeroPointCols": "ceilDiv(outBlocks, 2) if bits == 4 else outBlocks", |
| "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "commonShapeValid": "blockSizeOk and ranks.dataT == 2 and ranks.indicesT == 1 and ranks.scalesT == 2 and ranks.outputT == 2 and tensorDtypes.indicesT == \"uint32\" and tensorDtypes.scalesT == \"float32\" and tensorDtypes.outputT == \"float32\" and dim(shapes.outputT, 0) == dim(shapes.indicesT, 0) and dim(shapes.scalesT, 0) == dim(shapes.dataT, 0) and dim(shapes.scalesT, 1) == outBlocks", |
| "q4ShapeValid": "commonShapeValid and tensorDtypes.dataT == \"uint8\" and dim(shapes.outputT, 1) == dim(shapes.dataT, 1) * 2", |
| "q8ShapeValid": "commonShapeValid and tensorDtypes.dataT == \"uint8\" and dim(shapes.outputT, 1) == dim(shapes.dataT, 1)", |
| "zeroPointsValid": "present.zeroPointsT and ranks.zeroPointsT == 2 and tensorDtypes.zeroPointsT == \"uint8\" and dim(shapes.zeroPointsT, 0) == dim(shapes.dataT, 0) and dim(shapes.zeroPointsT, 1) == zeroPointCols", |
| "noZeroMode": "not present.zeroPointsT", |
| "zeroMode": "zeroPointsValid", |
| "workgroupFits": "workgroupSize > 0", |
| "foldedDispatchFits": "ceil(ceil(numel(shapes.outputT) / min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) / workgroupSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)" |
| }, |
| "when": ["foldedDispatchFits", "workgroupFits"], |
| "bindings": { |
| "data": { "arg": "dataT", "buffer": "read-only-storage", "elementType": "$dataElement" }, |
| "indices": { "arg": "indicesT", "buffer": "read-only-storage", "elementType": "$indexScalar" }, |
| "scales": { "arg": "scalesT", "buffer": "read-only-storage", "elementType": "$scaleScalar" }, |
| "output": { "arg": "outputT", "buffer": "storage", "elementType": "$outputElement" }, |
| "params": { |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indicesT, 0)" }, |
| { "name": "cols", "type": "u32", "value": "dim(shapes.outputT, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scalesT, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.dataT, 0)" } |
| ] |
| }, |
| "params_2": { |
| "name": "params", |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indicesT, 0)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.dataT, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scalesT, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.dataT, 0)" } |
| ] |
| }, |
| "zero_points": { "arg": "zeroPointsT", "buffer": "read-only-storage", "elementType": "$zeroPointElement" }, |
| "params_3": { |
| "name": "params", |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indicesT, 0)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.dataT, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scalesT, 1)" }, |
| { "name": "zeroPointCols", "type": "u32", "value": "zeroPointCols" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.dataT, 0)" } |
| ] |
| }, |
| "params_4": { |
| "name": "params", |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "indexCount", "type": "u32", "value": "dim(shapes.indicesT, 0)" }, |
| { "name": "cols", "type": "u32", "value": "dim(shapes.outputT, 1)" }, |
| { "name": "packedCols", "type": "u32", "value": "dim(shapes.dataT, 1)" }, |
| { "name": "blocks", "type": "u32", "value": "dim(shapes.scalesT, 1)" }, |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "rows", "type": "u32", "value": "dim(shapes.dataT, 0)" } |
| ] |
| } |
| }, |
| "variants": [ |
| { |
| "id": "q8_no_zero_vec4", |
| "priority": 10, |
| "when": ["q8ShapeValid", "noZeroMode", "bits == 8", "blockSize % 4 == 0", "dim(shapes.outputT, 1) % 4 == 0"], |
| "derive": { |
| "hasZero": false, |
| "scalarTail": false, |
| "dataElement": "\"vec4<u32>\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec4<f32>\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": ["data", "indices", "scales", "output", "params"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.indicesT, 0) * (dim(shapes.outputT, 1) / 4)), (workgroupSize)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.indicesT, 0) * (dim(shapes.outputT, 1) / 4)), (workgroupSize)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q4_no_zero_pair", |
| "priority": 10, |
| "when": ["q4ShapeValid", "noZeroMode", "bits == 4"], |
| "derive": { |
| "hasZero": false, |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec2<f32>\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q4-pair.wgsl.jinja", |
| "bindings": ["data", "indices", "scales", "output", "params_2"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.indicesT, 0) * dim(shapes.dataT, 1)), (workgroupSize)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.indicesT, 0) * dim(shapes.dataT, 1)), (workgroupSize)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q4_zero_pair", |
| "priority": 10, |
| "when": ["q4ShapeValid", "zeroMode", "bits == 4"], |
| "derive": { |
| "hasZero": true, |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec2<f32>\"", |
| "zeroPointElement": "\"u32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q4-pair.wgsl.jinja", |
| "bindings": ["data", "indices", "scales", "zero_points", "output", "params_3"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.indicesT, 0) * dim(shapes.dataT, 1)), (workgroupSize)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.indicesT, 0) * dim(shapes.dataT, 1)), (workgroupSize)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q8_zero_vec4", |
| "priority": 10, |
| "when": ["q8ShapeValid", "zeroMode", "bits == 8", "blockSize % 4 == 0", "dim(shapes.outputT, 1) % 4 == 0"], |
| "derive": { |
| "hasZero": true, |
| "scalarTail": false, |
| "dataElement": "\"vec4<u32>\"", |
| "zeroPointElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"vec4<f32>\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": ["data", "indices", "scales", "zero_points", "output", "params"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.indicesT, 0) * (dim(shapes.outputT, 1) / 4)), (workgroupSize)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.indicesT, 0) * (dim(shapes.outputT, 1) / 4)), (workgroupSize)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q8_no_zero_tail4", |
| "priority": 5, |
| "when": ["q8ShapeValid", "noZeroMode", "bits == 8"], |
| "derive": { |
| "hasZero": false, |
| "scalarTail": true, |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"f32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": ["data", "indices", "scales", "output", "params_4"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.indicesT, 0) * ceilDiv(dim(shapes.outputT, 1), 4)), (workgroupSize)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.indicesT, 0) * ceilDiv(dim(shapes.outputT, 1), 4)), (workgroupSize)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "q8_zero_tail4", |
| "priority": 5, |
| "when": ["q8ShapeValid", "zeroMode", "bits == 8"], |
| "derive": { |
| "hasZero": true, |
| "scalarTail": true, |
| "dataElement": "\"u32\"", |
| "indexScalar": "\"u32\"", |
| "scaleScalar": "\"f32\"", |
| "outputElement": "\"f32\"", |
| "zeroPointElement": "\"u32\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "shader": "gather-block-quantized-q8-vec4.wgsl.jinja", |
| "bindings": ["data", "indices", "scales", "zero_points", "output", "params_4"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.indicesT, 0) * ceilDiv(dim(shapes.outputT, 1), 4)), (workgroupSize)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.indicesT, 0) * ceilDiv(dim(shapes.outputT, 1), 4)), (workgroupSize)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
