| { |
| "domain": "com.microsoft", |
| "name": "BiasSoftmax", |
| "sinceVersion": 1, |
| "inputs": { "data": { "dtype": "T" }, "bias": { "dtype": "T" } }, |
| "outputs": { "output": { "dtype": "T", "rank": "ranks.data", "shape": "shapes.data" } }, |
| "attributes": { "axis": { "default": 1 }, "is_inner_broadcast": {} }, |
| "attributeConstraints": { "is_inner_broadcast": { "required": true } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "BLOCK_COLS": { "default": 2048 }, |
| "ONLINE_WORKGROUP_SIZE": { "default": 64 } |
| }, |
| "derive": { |
| "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data", |
| "batchCount": "outer(shapes.data, axisNorm)", |
| "blockSize": "dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm)", |
| "biasBlockCount": "numel(shapes.bias) / max(1, blockSize)", |
| "biasContract": "(numel(shapes.data) == 0 and numel(shapes.bias) == 0) or (blockSize > 0 and biasBlockCount > 0 and numel(shapes.bias) % blockSize == 0 and biasBlockCount <= batchCount and batchCount % biasBlockCount == 0)", |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "subgroupRowsWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "subgroupRowsLaneVecs": "ceilDiv(blockSize / 4, device.adapterInfo.subgroupMinSize) if (device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize > 0) else 0", |
| "biasBaseWorkgroup": "min(tunables.ONLINE_WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "biasMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "biasWorkgroupSize": "min(biasMaxWorkgroup, pow2ceil(ceilDiv(blockSize, 4)), max(biasBaseWorkgroup, pow2ceil(ceilDiv(biasBaseWorkgroup * biasMaxWorkgroup, max(1, batchCount)))))" |
| }, |
| "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "f16Ok(dtypes.T)"], |
| "bindings": { |
| "data": { "elementType": "$scalar" }, |
| "bias": { "elementType": "$scalar" }, |
| "output": { "elementType": "$scalar" } |
| }, |
| "variants": [ |
| { |
| "id": "online_subgroup_rows_vec4", |
| "priority": 20, |
| "when": ["blockSize % 4 == 0", "has(device.adapterInfo, \"subgroupMinSize\")", "device.adapterInfo.subgroupMinSize >= 16", "blockSize >= 64", "subgroupRowsLaneVecs >= 1", "subgroupRowsLaneVecs <= 8", "batchCount >= 64", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMaxSize <= subgroupRowsWorkgroup", "pow2ceil(subgroupRowsWorkgroup) == subgroupRowsWorkgroup", "ceilDiv(ceilDiv(batchCount, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "scalar": "dtypes.T", |
| "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "workgroupSize": "subgroupRowsWorkgroup", |
| "vecsPerLane": "subgroupRowsLaneVecs", |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.SubgroupRowsVec4", |
| "shader": "softmax-subgroup-rows.wgsl.jinja", |
| "derive": { "op": "\"biassoftmax\"" }, |
| "bindings": [ |
| { "arg": "data", "name": "x", "elementType": "$vectorScalar" }, |
| { "arg": "bias", "name": "bias", "elementType": "$vectorScalar" }, |
| { "arg": "output", "name": "y", "elementType": "$vectorScalar" }, |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "batchCount" }, |
| { "name": "vecCols", "type": "u32", "value": "blockSize / 4" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv(batchCount, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)", |
| "y": "ceilDiv(ceilDiv(batchCount, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "online_workgroup_vec4", |
| "priority": 15, |
| "when": ["blockSize % 4 == 0", "blockSize > 32", "blockSize < 65536", "pow2ceil(biasWorkgroupSize) == biasWorkgroupSize", "ceilDiv(batchCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "workgroupSize": "biasWorkgroupSize", |
| "combineSubgroups": false, |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))", |
| "scalarVec4": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.WorkgroupVec4", |
| "shader": "softmax-online.wgsl.jinja", |
| "derive": { "op": "\"biassoftmax\"", "useVec4": true }, |
| "bindings": [ |
| { "arg": "data", "name": "x", "elementType": "$vectorScalar" }, |
| { "arg": "bias", "name": "bias", "elementType": "$vectorScalar" }, |
| { "arg": "output", "name": "y", "elementType": "$vectorScalar" }, |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "batchCount" }, |
| { "name": "vecCols", "type": "u32", "value": "blockSize / 4" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "online_workgroup_scalar_vec4", |
| "priority": 15, |
| "when": ["blockSize % 4 != 0", "blockSize > 32", "blockSize < 65536", "pow2ceil(biasWorkgroupSize) == biasWorkgroupSize", "ceilDiv(batchCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "workgroupSize": "biasWorkgroupSize", |
| "combineSubgroups": false, |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))", |
| "scalarVec4": true |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.WorkgroupScalarVec4", |
| "shader": "softmax-online.wgsl.jinja", |
| "derive": { "op": "\"biassoftmax\"", "useVec4": true }, |
| "bindings": [ |
| { "arg": "data", "name": "x", "elementType": "$scalar" }, |
| { "arg": "bias", "name": "bias", "elementType": "$scalar" }, |
| { "arg": "output", "name": "y", "elementType": "$scalar" }, |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "batchCount" }, |
| { "name": "vecCols", "type": "u32", "value": "ceilDiv(blockSize, 4)" }, |
| { "name": "cols", "type": "u32", "value": "blockSize" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "longrow_split", |
| "priority": 40, |
| "when": ["blockSize >= 65536", "batchCount > 0", "batchCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(blockSize, tunables.BLOCK_COLS) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { "scalar": "dtypes.T", "combineSubgroups": false }, |
| "intermediates": [ |
| { |
| "id": "blockMax", |
| "dtype": "float32", |
| "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]" |
| }, |
| { |
| "id": "blockSum", |
| "dtype": "float32", |
| "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]" |
| }, |
| { "id": "rowMax", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" }, |
| { "id": "rowSum", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" } |
| ], |
| "passes": [ |
| { |
| "id": "block_stats", |
| "name": "BiasSoftmax.LongRowBlockStats", |
| "shader": "softmax-longrow-stats.wgsl.jinja", |
| "derive": { |
| "stage": "\"block\"", |
| "biasRow": true, |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| }, |
| "bindings": [ |
| "data", |
| "bias", |
| { "name": "blockMax", "elementType": "f32" }, |
| { "name": "blockSum", "elementType": "f32" }, |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" } |
| }, |
| { |
| "id": "row_stats", |
| "name": "BiasSoftmax.LongRowStats", |
| "shader": "softmax-longrow-stats.wgsl.jinja", |
| "derive": { "stage": "\"row\"" }, |
| "bindings": [ |
| { "name": "blockMax", "buffer": "read-only-storage", "elementType": "f32" }, |
| { "name": "blockSum", "buffer": "read-only-storage", "elementType": "f32" }, |
| { "name": "rowMax", "elementType": "f32" }, |
| { "name": "rowSum", "elementType": "f32" }, |
| { |
| "name": "params", |
| "struct": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }] |
| } |
| ], |
| "dispatch": { "x": "batchCount" } |
| }, |
| { |
| "id": "normalize", |
| "name": "BiasSoftmax.LongRowNormalize", |
| "shader": "bias-softmax-longrow-normalize.wgsl.jinja", |
| "derive": { |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| }, |
| "bindings": [ |
| "data", |
| "bias", |
| { "name": "rowMax", "buffer": "read-only-storage", "elementType": "f32" }, |
| { "name": "rowSum", "buffer": "read-only-storage", "elementType": "f32" }, |
| "output", |
| { "name": "params", "struct": [{ "name": "blockSize", "type": "u32", "value": "blockSize" }] } |
| ], |
| "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rows", |
| "priority": 30, |
| "when": ["blockSize > 0", "blockSize <= 8", "batchCount >= 64"], |
| "derive": { "scalar": "dtypes.T", "combineSubgroups": false, "packedRows": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.PackedRows", |
| "shader": "bias-softmax.wgsl.jinja", |
| "derive": { |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| }, |
| "bindings": [ |
| "data", |
| "bias", |
| "output", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "batchCount", "type": "u32", "value": "batchCount" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((batchCount), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": 1, |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "adaptive_row", |
| "priority": 10, |
| "when": ["numel(shapes.data) >= 0"], |
| "derive": { "packedRows": false, "scalar": "dtypes.T", "combineSubgroups": "device.features.has(\"subgroups\")" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.AdaptiveRow", |
| "shader": "bias-softmax.wgsl.jinja", |
| "derive": { |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCountSpec": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| }, |
| "bindings": [ |
| "data", |
| "bias", |
| "output", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "batchCount", "type": "u32", "value": "batchCount" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 } |
| } |
| ] |
| } |
| ] |
| } |
|
|
