| { |
| "domain": "com.microsoft", |
| "name": "GemmFastGelu", |
| "sinceVersion": 1, |
| "description": "Fuses MatMul, an optional bias, and FastGelu: `Y = FastGelu(X @ W + bias)`. `X` has rank at least 2 with shape `(..., K)`, `W` has shape `(K, N)`, and `bias` has shape `(N)`. The activation runs in the float32 accumulator before the output is narrowed, avoiding an intermediate `(..., N)` tensor. Bfloat16 is not implemented.", |
| "inputs": [ |
| { |
| "role": "X", |
| "dtype": "T", |
| "description": "Left operand of rank 2 or greater with shape `(..., K)`; every leading-axis coordinate identifies a row of the product." |
| }, |
| { "role": "W", "dtype": "T", "rank": 2, "description": "Right operand with shape `(K, N)`." }, |
| { |
| "role": "bias", |
| "dtype": "T", |
| "rank": 1, |
| "optional": true, |
| "description": "Optional bias with shape `(N)`, added before the activation." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "Y", |
| "dtype": "T", |
| "rank": "ranks.X", |
| "shape": "matmulShape(shapes.X, shapes.W)", |
| "description": "`FastGelu(X @ W + bias)`, with the same rank and leading dimensions as `X` and a trailing `N`." |
| } |
| ], |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "args": { |
| "X": { "kind": "tensor", "semantic": "X", "role": "input" }, |
| "W": { "kind": "tensor", "semantic": "W", "role": "weights" }, |
| "bias": { "kind": "tensor", "semantic": "bias", "role": "weights", "required": false }, |
| "Y": { "kind": "tensor", "semantic": "Y", "role": "output" } |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32", |
| "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32", |
| "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter", |
| "wave32Effective": "wave32Adapter or pinSubgroupSize32", |
| "kDim": "dim(shapes.W, 0)", |
| "nDim": "dim(shapes.W, 1)", |
| "gfgRows": "numel(shapes.X) / kDim if kDim > 0 else 0", |
| "shapeOk": "ranks.W == 2 and ranks.X >= 2 and kDim > 0 and nDim > 0 and dim(shapes.X, ranks.X - 1) == kDim and ranks.Y == ranks.X and dim(shapes.Y, ranks.Y - 1) == nDim and numel(shapes.Y) == gfgRows * nDim", |
| "dtypeOk": "(tensorDtypes.X == \"float32\" or tensorDtypes.X == \"float16\") and tensorDtypes.W == tensorDtypes.X and tensorDtypes.Y == tensorDtypes.X and f16Ok(dtypes.T)", |
| "biasOk": "ranks.bias == 1 and dim(shapes.bias, 0) == nDim and tensorDtypes.bias == tensorDtypes.X", |
| "baseContract": "shapeOk and dtypeOk", |
| "noBiasContract": "baseContract and not present.bias", |
| "biasContract": "baseContract and present.bias and biasOk", |
| "sgmatOperandBytes": "2 if tensorDtypes.X == \"float16\" and tensorDtypes.W == \"float16\" else 4", |
| "sgmatStorageBytes": "(32 * 32 + 64 * 32) * sgmatOperandBytes + 32 * 64 * 4", |
| "sgmatResourcesFit": "128 <= deviceWorkgroupCap and sgmatStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", |
| "sgmatLayoutOk": "gfgRows > 0 and kDim % 32 == 0 and nDim % 64 == 0 and ((tensorDtypes.X == \"float16\" and device.features.has(\"shader-f16\") and gfgRows >= 2) or (tensorDtypes.X == \"float32\" and gfgRows >= 32)) and ceilDiv(gfgRows, 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(nDim, 64) <= device.limits.maxComputeWorkgroupsPerDimension", |
| "sgmatContract": "wave32Effective and sgmatLayoutOk and sgmatResourcesFit" |
| }, |
| "constants": { |
| "usesF16": "tensorDtypes.X == \"float16\"", |
| "transA": false, |
| "transB": false, |
| "gemmEpilogue": "\"fastgelu\"" |
| }, |
| "bindingSets": { |
| "noBias": [ |
| { "name": "a", "arg": "X", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { "name": "b", "arg": "W", "semantic": "W", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "M", "type": "u32", "value": "gfgRows" }, |
| { "name": "N", "type": "u32", "value": "nDim" }, |
| { "name": "K", "type": "u32", "value": "kDim" }, |
| { "name": "alpha", "type": "f32", "value": 1 } |
| ] |
| } |
| } |
| ], |
| "withBias": [ |
| { "name": "a", "arg": "X", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { "name": "b", "arg": "W", "semantic": "W", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "bias", |
| "arg": "bias", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "M", "type": "u32", "value": "gfgRows" }, |
| { "name": "N", "type": "u32", "value": "nDim" }, |
| { "name": "K", "type": "u32", "value": "kDim" }, |
| { "name": "alpha", "type": "f32", "value": 1 }, |
| { "name": "beta", "type": "f32", "value": 1 } |
| ] |
| } |
| } |
| ], |
| "sgmatNoBias": [ |
| { "name": "x", "arg": "X", "semantic": "a", "buffer": { "type": "read-only-storage" }, "elementType": "$xScalar" }, |
| { "name": "w", "arg": "W", "semantic": "b", "buffer": { "type": "read-only-storage" }, "elementType": "$wScalar" }, |
| { "name": "y", "arg": "Y", "semantic": "y", "buffer": { "type": "storage" }, "elementType": "$yScalar" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { "name": "Params", "fields": [{ "name": "M", "type": "u32", "value": "gfgRows" }] } |
| } |
| ], |
| "sgmatBias": [ |
| { "name": "x", "arg": "X", "semantic": "a", "buffer": { "type": "read-only-storage" }, "elementType": "$xScalar" }, |
| { "name": "w", "arg": "W", "semantic": "b", "buffer": { "type": "read-only-storage" }, "elementType": "$wScalar" }, |
| { |
| "name": "bias", |
| "arg": "bias", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$bScalar" |
| }, |
| { "name": "y", "arg": "Y", "semantic": "y", "buffer": { "type": "storage" }, "elementType": "$yScalar" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { "name": "Params", "fields": [{ "name": "M", "type": "u32", "value": "gfgRows" }] } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "sgmat_bias", |
| "priority": 100, |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [ |
| { "componentType": "f16", "M": 8, "N": 8, "K": 8 }, |
| { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 } |
| ] |
| }, |
| "when": ["biasContract", "sgmatContract"], |
| "constants": { |
| "hasBias": true, |
| "gemmEpilogue": "\"fastgelu\"", |
| "bTransposedStorage": true, |
| "xScalar": "\"f16\" if tensorDtypes.X == \"float16\" else \"f32\"", |
| "wScalar": "\"f16\" if tensorDtypes.W == \"float16\" else \"f32\"", |
| "bScalar": "\"f16\" if tensorDtypes.bias == \"float16\" else \"f32\"", |
| "yScalar": "\"f16\" if tensorDtypes.Y == \"float16\" else \"f32\"", |
| "mmaScalar": "\"f16\" if tensorDtypes.X == \"float16\" and tensorDtypes.W == \"float16\" else \"f32\"", |
| "usesF16": "tensorDtypes.X == \"float16\" or tensorDtypes.W == \"float16\" or tensorDtypes.Y == \"float16\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.SubgroupMatrixBias", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": "sgmatBias", |
| "dispatch": { "workgroups": "ceilDiv(gfgRows, 32)", "x": "ceilDiv(nDim, 64)" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat", |
| "priority": 100, |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [ |
| { "componentType": "f16", "M": 8, "N": 8, "K": 8 }, |
| { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 } |
| ] |
| }, |
| "when": ["noBiasContract", "sgmatContract"], |
| "constants": { |
| "hasBias": false, |
| "gemmEpilogue": "\"fastgelu\"", |
| "bTransposedStorage": true, |
| "xScalar": "\"f16\" if tensorDtypes.X == \"float16\" else \"f32\"", |
| "wScalar": "\"f16\" if tensorDtypes.W == \"float16\" else \"f32\"", |
| "yScalar": "\"f16\" if tensorDtypes.Y == \"float16\" else \"f32\"", |
| "mmaScalar": "\"f16\" if tensorDtypes.X == \"float16\" and tensorDtypes.W == \"float16\" else \"f32\"", |
| "usesF16": "tensorDtypes.X == \"float16\" or tensorDtypes.W == \"float16\" or tensorDtypes.Y == \"float16\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.SubgroupMatrix", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": "sgmatNoBias", |
| "dispatch": { "workgroups": "ceilDiv(gfgRows, 32)", "x": "ceilDiv(nDim, 64)" } |
| } |
| ] |
| }, |
| { |
| "id": "tiled_bias", |
| "priority": 0, |
| "when": ["biasContract"], |
| "constants": { "hasBias": true, "rowBias": true, "scalarBias": false, "columnBias": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.Bias", |
| "shader": "gemm-fast-gelu.wgsl.jinja", |
| "bindings": "withBias", |
| "dispatch": { |
| "x": "min(ceilDiv(nDim, 64), device.limits.maxComputeWorkgroupsPerDimension)", |
| "y": "min(ceilDiv(gfgRows, 64), device.limits.maxComputeWorkgroupsPerDimension)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "tiled", |
| "priority": 0, |
| "when": ["noBiasContract"], |
| "constants": { "hasBias": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu", |
| "shader": "gemm-fast-gelu.wgsl.jinja", |
| "bindings": "noBias", |
| "dispatch": { |
| "x": "min(ceilDiv(nDim, 64), device.limits.maxComputeWorkgroupsPerDimension)", |
| "y": "min(ceilDiv(gfgRows, 64), device.limits.maxComputeWorkgroupsPerDimension)" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
