{ "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))" } } ] } ] }