| { |
| "domain": "com.microsoft", |
| "name": "GemmFastGelu", |
| "sinceVersion": 1, |
| "inputs": { |
| "X": { "dtype": "T" }, |
| "W": { "dtype": "T", "rank": 2 }, |
| "bias": { "dtype": "T", "rank": 1, "optional": true } |
| }, |
| "outputs": { "Y": { "dtype": "T", "rank": "ranks.X", "shape": "matmulShape(shapes.X, shapes.W)" } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "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) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "sgmatContract": "wave32Effective and sgmatLayoutOk and sgmatResourcesFit", |
| "usesF16": "tensorDtypes.X == \"float16\"", |
| "transA": false, |
| "transB": false, |
| "gemmEpilogue": "\"fastgelu\"", |
| "sgmatDirectRows": 64, |
| "sgmatDirectCols": 32, |
| "sgmatDirectStorageBytes": "2 * sgmatDirectCols * 32 * sgmatOperandBytes + sgmatDirectRows * sgmatDirectCols * 4", |
| "sgmatDirectContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatDirectRows == 0 and kDim % 32 == 0 and nDim % sgmatDirectCols == 0 and deviceWorkgroupCap >= 128 and sgmatDirectStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatDirectRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatDirectCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)" |
| }, |
| "bindings": { |
| "x": { "arg": "X", "buffer": "read-only-storage", "elementType": "$xScalar" }, |
| "w": { "arg": "W", "buffer": "read-only-storage", "elementType": "$wScalar" }, |
| "y": { "arg": "Y", "buffer": "storage", "elementType": "$yScalar" }, |
| "params": { "buffer": "uniform", "struct": [{ "name": "M", "type": "u32", "value": "gfgRows" }] }, |
| "a": { "arg": "X", "buffer": "read-only-storage", "elementType": "$T" }, |
| "b": { "arg": "W", "buffer": "read-only-storage", "elementType": "$T" }, |
| "y_2": { "arg": "Y", "name": "y", "buffer": "storage", "elementType": "$T" } |
| }, |
| "variants": [ |
| { |
| "id": "sgmat_direct_bias", |
| "priority": 110, |
| "when": ["biasContract", "sgmatDirectContract", "tensorDtypes.X == \"float32\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] |
| }, |
| "derive": { |
| "hasBias": true, |
| "bTransposedStorage": true, |
| "xScalar": "\"f32\"", |
| "wScalar": "\"f32\"", |
| "bScalar": "\"f32\"", |
| "yScalar": "\"f32\"", |
| "mmaScalar": "\"f32\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": "sgmatDirectRows", |
| "tileN": "sgmatDirectCols", |
| "tileK": 32, |
| "directAStorage": true, |
| "doubleBufferedB": true, |
| "balancedMatrixAccumulation": true |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_direct_bias", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_direct", |
| "priority": 110, |
| "when": ["noBiasContract", "sgmatDirectContract", "tensorDtypes.X == \"float32\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] |
| }, |
| "derive": { |
| "hasBias": false, |
| "bTransposedStorage": true, |
| "xScalar": "\"f32\"", |
| "wScalar": "\"f32\"", |
| "bScalar": "\"f32\"", |
| "yScalar": "\"f32\"", |
| "mmaScalar": "\"f32\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": "sgmatDirectRows", |
| "tileN": "sgmatDirectCols", |
| "tileK": 32, |
| "directAStorage": true, |
| "doubleBufferedB": true, |
| "balancedMatrixAccumulation": true |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_direct", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_direct_bias_f16", |
| "priority": 110, |
| "when": ["biasContract", "sgmatDirectContract", "tensorDtypes.X == \"float16\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "hasBias": true, |
| "bTransposedStorage": true, |
| "xScalar": "\"f16\"", |
| "wScalar": "\"f16\"", |
| "bScalar": "\"f16\"", |
| "yScalar": "\"f16\"", |
| "mmaScalar": "\"f16\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": "sgmatDirectRows", |
| "tileN": "sgmatDirectCols", |
| "tileK": 32, |
| "directAStorage": true, |
| "doubleBufferedB": true, |
| "balancedMatrixAccumulation": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_direct_bias_f16", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_direct_f16", |
| "priority": 110, |
| "when": ["noBiasContract", "sgmatDirectContract", "tensorDtypes.X == \"float16\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "hasBias": false, |
| "bTransposedStorage": true, |
| "xScalar": "\"f16\"", |
| "wScalar": "\"f16\"", |
| "bScalar": "\"f16\"", |
| "yScalar": "\"f16\"", |
| "mmaScalar": "\"f16\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": "sgmatDirectRows", |
| "tileN": "sgmatDirectCols", |
| "tileK": 32, |
| "directAStorage": true, |
| "doubleBufferedB": true, |
| "balancedMatrixAccumulation": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_direct_f16", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_bias", |
| "priority": 100, |
| "when": ["biasContract", "sgmatContract", "tensorDtypes.X == \"float32\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] |
| }, |
| "derive": { |
| "hasBias": true, |
| "bTransposedStorage": true, |
| "xScalar": "\"f32\"", |
| "wScalar": "\"f32\"", |
| "bScalar": "\"f32\"", |
| "yScalar": "\"f32\"", |
| "mmaScalar": "\"f32\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": 32, |
| "tileN": 64, |
| "tileK": 32, |
| "directAStorage": false, |
| "doubleBufferedB": false, |
| "balancedMatrixAccumulation": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_bias", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat", |
| "priority": 100, |
| "when": ["noBiasContract", "sgmatContract", "tensorDtypes.X == \"float32\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] |
| }, |
| "derive": { |
| "hasBias": false, |
| "bTransposedStorage": true, |
| "xScalar": "\"f32\"", |
| "wScalar": "\"f32\"", |
| "bScalar": "\"f32\"", |
| "yScalar": "\"f32\"", |
| "mmaScalar": "\"f32\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": 32, |
| "tileN": 64, |
| "tileK": 32, |
| "directAStorage": false, |
| "doubleBufferedB": false, |
| "balancedMatrixAccumulation": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_bias_f16", |
| "priority": 100, |
| "when": ["biasContract", "sgmatContract", "tensorDtypes.X == \"float16\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "hasBias": true, |
| "bTransposedStorage": true, |
| "xScalar": "\"f16\"", |
| "wScalar": "\"f16\"", |
| "bScalar": "\"f16\"", |
| "yScalar": "\"f16\"", |
| "mmaScalar": "\"f16\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": 32, |
| "tileN": 64, |
| "tileK": 32, |
| "directAStorage": false, |
| "doubleBufferedB": false, |
| "balancedMatrixAccumulation": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_bias_f16", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_f16", |
| "priority": 100, |
| "when": ["noBiasContract", "sgmatContract", "tensorDtypes.X == \"float16\""], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "hasBias": false, |
| "bTransposedStorage": true, |
| "xScalar": "\"f16\"", |
| "wScalar": "\"f16\"", |
| "bScalar": "\"f16\"", |
| "yScalar": "\"f16\"", |
| "mmaScalar": "\"f16\"", |
| "inFeatures": "kDim", |
| "outFeatures": "nDim", |
| "tileM": 32, |
| "tileN": 64, |
| "tileK": 32, |
| "directAStorage": false, |
| "doubleBufferedB": false, |
| "balancedMatrixAccumulation": false |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.sgmat_f16", |
| "shader": "gemm-subgroup-matrix.wgsl.jinja", |
| "bindings": ["x", "w", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(nDim, tileN)", |
| "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)", |
| "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "tiled_bias", |
| "priority": 0, |
| "when": ["biasContract"], |
| "derive": { "hasBias": true, "rowBias": true, "scalarBias": false, "columnBias": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu.Bias", |
| "shader": "gemm-fast-gelu.wgsl.jinja", |
| "bindings": [ |
| "a", |
| "b", |
| "bias", |
| "y_2", |
| { |
| "name": "params", |
| "struct": [ |
| { "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 } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv(nDim, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "min(ceilDiv(gfgRows, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "tiled", |
| "priority": 0, |
| "when": ["noBiasContract"], |
| "derive": { "hasBias": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmFastGelu", |
| "shader": "gemm-fast-gelu.wgsl.jinja", |
| "bindings": [ |
| "a", |
| "b", |
| "y_2", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "gfgRows" }, |
| { "name": "N", "type": "u32", "value": "nDim" }, |
| { "name": "K", "type": "u32", "value": "kDim" }, |
| { "name": "alpha", "type": "f32", "value": 1 } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv(nDim, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "min(ceilDiv(gfgRows, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
