{ "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": { "kDim": "dim(shapes.W, 0)", "nDim": "dim(shapes.W, 1)", "gemvN": "nDim", "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", "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32", "gemvLanes": 32, "vec4OutputTile": "4 * gemvLanes", "gemvWorkgroups": "ceilDiv(gemvN, vec4OutputTile)", "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / gemvLanes), floor(device.limits.maxComputeWorkgroupStorageSize / (16 * gemvLanes)))", "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))", "gemvResourcesFit": "gemvLanes <= device.limits.maxComputeWorkgroupSizeX and gemvSlices <= device.limits.maxComputeWorkgroupSizeY and gemvLanes * gemvSlices <= device.limits.maxComputeInvocationsPerWorkgroup and 16 * gemvLanes * gemvSlices <= device.limits.maxComputeWorkgroupStorageSize", "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\"", "transB": false, "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)", "sgmatOutputGroups": "ceilDiv(gfgRows, sgmatDirectRows) * ceilDiv(nDim, sgmatDirectCols)", "sgmatSmallRows": 32, "sgmatSmallCols": 16, "sgmatSmallStorageBytes": "2 * sgmatSmallCols * 32 * sgmatOperandBytes + sgmatSmallRows * sgmatSmallCols * 4", "sgmatSmallContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatSmallRows == 0 and kDim % 32 == 0 and nDim % sgmatSmallCols == 0 and deviceWorkgroupCap >= 128 and sgmatSmallStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatSmallRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatSmallCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "sgmatMediumRows": 32, "sgmatMediumCols": 32, "sgmatMediumStorageBytes": "2 * sgmatMediumCols * 32 * sgmatOperandBytes + sgmatMediumRows * sgmatMediumCols * 4", "sgmatMediumContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatMediumRows == 0 and kDim % 32 == 0 and nDim % sgmatMediumCols == 0 and deviceWorkgroupCap >= 128 and sgmatMediumStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatMediumRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatMediumCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "portableSmallGroups": "ceilDiv(gfgRows, 32) * ceilDiv(nDim, 32)", "portableLargeGroups": "ceilDiv(gfgRows, 64) * ceilDiv(nDim, 64)", "portableSmallBudget": "tunables.PORTABLE_SMALL_TILE_GROUPS_F16 if usesF16 else tunables.PORTABLE_SMALL_TILE_GROUPS_F32", "portableAvoidsPadding": "portableLargeGroups * 3 >= portableSmallGroups", "portableXe3LpgF32": "not usesF16 and has(device.adapterInfo, \"vendor\") and device.adapterInfo.vendor == \"intel\" and has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"xe-3lpg\"", "portableFixedWave32NoF16": "wave32Effective and not device.features.has(\"shader-f16\") and gfgRows >= device.adapterInfo.subgroupMaxSize and kDim >= 16 * device.adapterInfo.subgroupMaxSize and portableSmallGroups >= 8 * device.adapterInfo.subgroupMaxSize", "portableLargeNoSubgroupGrid": "not usesF16 and not device.features.has(\"subgroups\") and portableSmallGroups >= 4 * deviceWorkgroupCap", "gemmTileRows": "32 if portableXe3LpgF32 or portableFixedWave32NoF16 or portableLargeNoSubgroupGrid or portableSmallGroups <= portableSmallBudget * (2 if portableAvoidsPadding else 1) else 64", "gemmTileCols": "gemmTileRows", "gfgWideMicroPreferred": "(variableSubgroup16To32 or portableFixedWave32NoF16) and (8 if usesF16 else 4) * device.adapterInfo.subgroupMaxSize <= ceilDiv(gfgRows, gemmTileRows) * ceilDiv(nDim, gemmTileCols) and gemmTileRows * gemmTileCols / 8 <= deviceWorkgroupCap and gemmTileCols / 2 <= device.limits.maxComputeWorkgroupSizeX and gemmTileRows / 4 <= device.limits.maxComputeWorkgroupSizeY", "gemmMicroRows": "4 if gfgWideMicroPreferred else gemmTileRows / 16", "gemmMicroCols": "2 if gfgWideMicroPreferred else gemmTileCols / 16", "gemmTileK": "(16 * 16 * 4) / gemmTileRows if usesF16 else (16 if portableLargeNoSubgroupGrid and has(device.adapterInfo, \"vendor\") and device.adapterInfo.vendor == \"intel\" else 32)", "gemmTilePadding": "1 if not usesF16 and 2 * gemmTileRows * (gemmTileK + 4) * 4 <= device.limits.maxComputeWorkgroupStorageSize else 0", "gemmCompensatedTiles": "not usesF16 and kDim > gemmTileK", "gemmStaticM": "gfgRows", "gemmStaticN": "nDim", "gemmStaticK": "kDim" }, "bindings": { "x": { "arg": "X", "elementType": "$xScalar" }, "w": { "arg": "W", "elementType": "$wScalar" }, "y": { "arg": "Y", "elementType": "$yScalar" }, "params": { "struct": [{ "name": "M", "type": "u32", "value": "gfgRows" }] }, "a": { "arg": "X", "elementType": "$T" }, "b": { "arg": "W", "elementType": "$T" }, "y_t": { "arg": "Y", "name": "y", "elementType": "$T" } }, "variants": [ { "id": "vector_matrix_bias", "priority": 130, "when": ["biasContract", "gfgRows == 1", "nDim % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit"], "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputBuffer": "\"y\"", "gemvBias": true, "unrollK2": "usesF16" }, "passes": [ { "id": "main", "name": "GemmFastGelu.VectorMatrixBias", "shader": "matmul-vector-matrix-vec4.wgsl.jinja", "bindings": [ "a", { "name": "b", "arg": "W", "elementType": "$vectorScalar" }, { "arg": "bias", "elementType": "$vectorScalar" }, { "name": "y", "arg": "Y", "elementType": "$vectorScalar" }, { "name": "params", "struct": [ { "name": "K", "type": "u32", "value": "kDim" }, { "name": "N4", "type": "u32", "value": "nDim / 4" } ] } ], "dispatch": { "x": "gemvWorkgroups" } } ] }, { "id": "vector_matrix", "priority": 130, "when": ["noBiasContract", "gfgRows == 1", "nDim % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit"], "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputBuffer": "\"y\"", "gemvBias": false, "unrollK2": "usesF16" }, "passes": [ { "id": "main", "name": "GemmFastGelu.VectorMatrix", "shader": "matmul-vector-matrix-vec4.wgsl.jinja", "bindings": [ "a", { "name": "b", "arg": "W", "elementType": "$vectorScalar" }, { "name": "y", "arg": "Y", "elementType": "$vectorScalar" }, { "name": "params", "struct": [ { "name": "K", "type": "u32", "value": "kDim" }, { "name": "N4", "type": "u32", "value": "nDim / 4" } ] } ], "dispatch": { "x": "gemvWorkgroups" } } ] }, { "id": "sgmat_medium_bias", "priority": 120, "when": ["biasContract", "sgmatMediumContract", "tensorDtypes.X == \"float32\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] }, "derive": { "hasBias": true, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatMediumRows", "tileN": "sgmatMediumCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": true }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatMediumBias", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_medium", "priority": 120, "when": ["noBiasContract", "sgmatMediumContract", "tensorDtypes.X == \"float32\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] }, "derive": { "hasBias": false, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatMediumRows", "tileN": "sgmatMediumCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": true }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatMedium", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_medium_bias_f16", "priority": 120, "when": ["biasContract", "sgmatMediumContract", "tensorDtypes.X == \"float16\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasBias": true, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatMediumRows", "tileN": "sgmatMediumCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatMediumBiasF16", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_medium_f16", "priority": 120, "when": ["noBiasContract", "sgmatMediumContract", "tensorDtypes.X == \"float16\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasBias": false, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatMediumRows", "tileN": "sgmatMediumCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatMediumF16", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_small_bias", "priority": 125, "when": ["biasContract", "sgmatSmallContract", "tensorDtypes.X == \"float32\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] }, "derive": { "hasBias": true, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatSmallRows", "tileN": "sgmatSmallCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": true }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatSmallBias", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_small", "priority": 125, "when": ["noBiasContract", "sgmatSmallContract", "tensorDtypes.X == \"float32\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }] }, "derive": { "hasBias": false, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatSmallRows", "tileN": "sgmatSmallCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": true }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatSmall", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_small_bias_f16", "priority": 125, "when": ["biasContract", "sgmatSmallContract", "tensorDtypes.X == \"float16\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasBias": true, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatSmallRows", "tileN": "sgmatSmallCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatSmallBiasF16", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"] }, { "id": "sgmat_small_f16", "priority": 125, "when": ["noBiasContract", "sgmatSmallContract", "tensorDtypes.X == \"float16\""], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "derive": { "hasBias": false, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatSmallRows", "tileN": "sgmatSmallCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatSmallF16", "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)" } } ], "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"] }, { "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, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatDirectRows", "tileN": "sgmatDirectCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": true }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatDirectBias", "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatDirectRows", "tileN": "sgmatDirectCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": true }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatDirect", "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatDirectRows", "tileN": "sgmatDirectCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatDirectBiasF16", "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": "sgmatDirectRows", "tileN": "sgmatDirectCols", "directAStorage": true, "doubleBufferedB": true, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatDirectF16", "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": 32, "tileN": 64, "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatBias", "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f32\"", "wScalar": "\"f32\"", "bScalar": "\"f32\"", "yScalar": "\"f32\"", "mmaScalar": "\"f32\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": 32, "tileN": 64, "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": 32, "tileN": 64, "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatBiasF16", "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)" } } ], "demoteWhen": ["false"] }, { "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, "xScalar": "\"f16\"", "wScalar": "\"f16\"", "bScalar": "\"f16\"", "yScalar": "\"f16\"", "mmaScalar": "\"f16\"", "inFeatures": "kDim", "outFeatures": "nDim", "tileM": 32, "tileN": 64, "directAStorage": false, "doubleBufferedB": false, "balancedMatrixAccumulation": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.SgmatF16", "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)" } } ], "demoteWhen": ["false"] }, { "id": "tiled_bias", "priority": 0, "when": ["biasContract"], "derive": { "hasBias": true, "rowBias": true, "scalarBias": false, "columnBias": false }, "passes": [ { "id": "main", "name": "GemmFastGelu.Bias", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "bindings": [ "a", "b", "bias", "y_t", { "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, gemmTileCols), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(gfgRows, gemmTileRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } } ] }, { "id": "tiled", "priority": 0, "when": ["noBiasContract"], "derive": { "hasBias": false }, "passes": [ { "id": "main", "name": "GemmFastGelu", "shader": "gemm-epilogue-tiled-reg.wgsl.jinja", "bindings": [ "a", "b", "y_t", { "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, gemmTileCols), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "min(ceilDiv(gfgRows, gemmTileRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" } } ] } ], "tunables": { "SMALL_MATRIX_MAX_GROUPS": { "default": 48 }, "MEDIUM_MATRIX_MAX_GROUPS_F32": { "default": 96 }, "MEDIUM_MATRIX_MAX_GROUPS_F16": { "default": 80 }, "GEMV_TARGET_BLOCKS": { "default": 512 }, "PORTABLE_SMALL_TILE_GROUPS_F32": { "default": 64 }, "PORTABLE_SMALL_TILE_GROUPS_F16": { "default": 128 } } }