| { |
| "domain": "com.microsoft", |
| "name": "TransposeMatMul", |
| "sinceVersion": 1, |
| "inputs": { "A": { "dtype": "T" }, "B": { "dtype": "T" } }, |
| "outputs": { |
| "Y": { |
| "dtype": "T", |
| "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)", |
| "shape": "matmulShape(logicalAShape, logicalBShape)" |
| } |
| }, |
| "attributes": { "alpha": { "default": 1 }, "transA": { "default": 0 }, "transB": { "default": 0 } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { |
| "TILED_REG_MIN_WORKGROUPS": { "default": 64 }, |
| "PLAIN_RANK2_REG_DEEP_K_TILES": { "default": 128 }, |
| "GEMV_TARGET_BLOCKS": { "default": 512 }, |
| "SUBGROUP_MATRIX_MIN_M": { "default": 2 }, |
| "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": { "default": 512 }, |
| "SUBGROUP_MATRIX_SPLITK_MIN_K": { "default": 1024 }, |
| "SUBGROUP_MATRIX_SPLITK_MAX_TILES": { "default": 128 }, |
| "BAND_VEC4_MAX_ROWS": { "default": 16 }, |
| "BAND_SPLIT_TARGET_WORKGROUPS": { "default": 256 }, |
| "BAND_SPLIT_MAX_COLUMN_GROUPS": { "default": 24 }, |
| "BAND_SPLIT_SLICES": { "default": 8 }, |
| "BAND_PREFER_MAX_ROWS": { "default": 8 }, |
| "BAND_PREFER_DEEP_K": { "default": 4096 }, |
| "BROADCAST_TRANSB_MIN_WORKGROUPS": { "default": 64 }, |
| "BROADCAST_TRANSB_MAX_PADDING_RATIO": { "default": 2 } |
| }, |
| "derive": { |
| "batchMovedAShape": "shapes.A", |
| "batchMovedBShape": "shapes.B", |
| "logicalAShape": "moveAxis(batchMovedAShape, -1, -2) if attrs.transA != 0 and ranks.A > 1 else batchMovedAShape", |
| "logicalBShape": "moveAxis(batchMovedBShape, -1, -2) if attrs.transB != 0 and ranks.B > 1 else batchMovedBShape", |
| "gemvN": "dim(shapes.B, 1)", |
| "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", |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "rank2DeepPortableTier": "variableSubgroup16To32 or (has(device.adapterInfo, \"architecture\") and (device.adapterInfo.architecture == \"pascal\" or (not device.features.has(\"subgroups\") and (device.adapterInfo.architecture == \"apple\" or device.adapterInfo.architecture == \"gen-9\"))))", |
| "broadcastTransbM": "dim(shapes.A, ranks.A - 2)", |
| "broadcastTransbN": "dim(shapes.B, ranks.B - 2)", |
| "broadcastTransbK": "dim(shapes.A, ranks.A - 1)", |
| "broadcastTransbBatches": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 2))", |
| "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", |
| "registerTile": 64, |
| "generalTile": 32, |
| "tiledRegResourcesFit": "registerTile / 4 <= device.limits.maxComputeWorkgroupSizeX and registerTile / 4 <= device.limits.maxComputeWorkgroupSizeY and registerTile * registerTile / 16 <= device.limits.maxComputeInvocationsPerWorkgroup and 32 * registerTile * dtypeBytes(dtypes.T) <= device.limits.maxComputeWorkgroupStorageSize", |
| "plainRank2RegDeepPreferredTier": "rank2DeepPortableTier and dim(shapes.A, 1) >= tunables.PLAIN_RANK2_REG_DEEP_K_TILES * 16 and dim(shapes.A, 1) % 16 == 0", |
| "subgroupMatrixResourcesFit": "128 <= deviceWorkgroupCap and ((32 * 32 + 64 * 32) * dtypeBytes(dtypes.T) + 4 * 4 * 64 * 4) <= device.limits.maxComputeWorkgroupStorageSize", |
| "sgmatSplitKDepth": "dim(shapes.A, ranks.A - 1)", |
| "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0", |
| "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0", |
| "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0", |
| "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0", |
| "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0", |
| "bandSplitWant": "pow2ceil(ceilDiv(tunables.BAND_SPLIT_TARGET_WORKGROUPS, max(1, gemvWorkgroups)))", |
| "bandSplitK": "16 if (bandSplitWant >= 16 and dim(shapes.A, ranks.A - 1) >= 4096) else (8 if (bandSplitWant >= 8 and dim(shapes.A, ranks.A - 1) >= 2048) else (4 if (bandSplitWant >= 4 and dim(shapes.A, ranks.A - 1) >= 1024) else (2 if (bandSplitWant >= 2 and dim(shapes.A, ranks.A - 1) >= 512) else 1)))", |
| "scalar": "dtypes.T", |
| "alpha": "attrs.alpha", |
| "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"", |
| "aRank": "ranks.A", |
| "bRank": "ranks.B", |
| "fusedSgmatRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and dim(shapes.A, 1) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.B, 1)", |
| "sgmatOutTiles": "ceilDiv(dim(shapes.A, 0), 32) * ceilDiv(dim(shapes.B, 1), 64) if fusedSgmatRank2Ok else 1", |
| "sgmatSplitKWant": "ceilDiv(tunables.SUBGROUP_MATRIX_SPLITK_TARGET_WGS, sgmatOutTiles)", |
| "sgmatSplitK": "32 if (sgmatSplitKWant > 16 and sgmatSplitK32Ok) else (16 if (sgmatSplitKWant > 8 and sgmatSplitK16Ok) else (8 if (sgmatSplitKWant > 4 and sgmatSplitK8Ok) else (4 if (sgmatSplitKWant > 2 and sgmatSplitK4Ok) else (2 if sgmatSplitK2Ok else 1))))", |
| "bandRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and dim(shapes.A, 1) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.B, 1)", |
| "broadcastTransbContract": "(f16Ok(dtypes.T)) and (attrs.transA == 0 and attrs.transB != 0) and (ranks.A > ranks.B and ranks.B >= 2) and (ranks.Y == ranks.A) and (sameShape(shapes.Y, matmulShape(logicalAShape, logicalBShape))) and (dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 1)) and (dim(shapes.A, ranks.A - 2) >= 64) and (dim(shapes.A, ranks.A - 1) >= 32) and (dim(shapes.B, ranks.B - 2) >= 64) and (numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 2)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| }, |
| "bindings": { |
| "a": { "arg": "A", "elementType": "$scalar" }, |
| "b": { "arg": "B", "elementType": "$vectorScalar" }, |
| "partials": { "buffer": "read-only-storage", "elementType": "f32" }, |
| "y": { "arg": "Y", "elementType": "$scalar" }, |
| "params": { "struct": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] }, |
| "b_scalar": { "arg": "B", "name": "b", "elementType": "$scalar" }, |
| "params_rows": { "name": "params", "struct": [{ "name": "M", "type": "u32", "value": "rowCount" }] } |
| }, |
| "variants": [ |
| { |
| "id": "broadcast_transb_tiled_reg", |
| "priority": 6, |
| "when": ["broadcastTransbContract", "tiledRegResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { |
| "bShape": "logicalBShape", |
| "bTransposed": true, |
| "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM", |
| "N": "broadcastTransbN", |
| "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1", |
| "transBatchA": false, |
| "regSequentialK": "dtypes.T == \"f16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.BroadcastTransBTiledReg", |
| "shader": "matmul-tiled-general-reg.wgsl.jinja", |
| "derive": { |
| "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []", |
| "aRank": "ranks.A if ranks.B > 2 else 2", |
| "K": "dim(shapes.A, ranks.A - 1)" |
| }, |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { "x": "ceilDiv(N, registerTile)", "y": "ceilDiv(rowCount, registerTile)", "z": "batchCount" } |
| } |
| ], |
| "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM, registerTile) * ceilDiv(broadcastTransbN, registerTile) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM, registerTile) * registerTile * ceilDiv(broadcastTransbN, registerTile) * registerTile * ceilDiv(broadcastTransbK,16) * 16 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"] |
| }, |
| { |
| "id": "broadcast_transb_subgroup_matrix_f16", |
| "priority": 11, |
| "when": ["broadcastTransbContract", "dtypes.T == \"f16\"", "wave32Effective", "subgroupMatrixResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2),32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2),64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "bShape": "logicalBShape", |
| "bTransposed": true, |
| "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM", |
| "K": "broadcastTransbK", |
| "N": "broadcastTransbN", |
| "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1", |
| "hasBias": false, |
| "fScalar": "dtypes.T", |
| "outScalar": "dtypes.T", |
| "generalAddressing": true, |
| "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 2) % 64 != 0", |
| "outputBuffer": "\"y\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.BroadcastTransBSubgroupMatrix", |
| "shader": "matmul-subgroup-matrix-ext.wgsl.jinja", |
| "derive": { |
| "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []", |
| "aRank": "ranks.A if ranks.B > 2 else 2" |
| }, |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { "x": "ceilDiv(N,64)", "y": "ceilDiv(rowCount,32)", "z": "batchCount" } |
| } |
| ], |
| "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM,32) * ceilDiv(broadcastTransbN,64) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM,32) * 32 * ceilDiv(broadcastTransbN,64) * 64 * ceilDiv(broadcastTransbK,32) * 32 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"] |
| }, |
| { |
| "id": "broadcast_transb_subgroup_matrix_f32", |
| "priority": 11, |
| "when": ["broadcastTransbContract", "dtypes.T == \"f32\"", "wave32Effective", "subgroupMatrixResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2),32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2),64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "bShape": "logicalBShape", |
| "bTransposed": true, |
| "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM", |
| "K": "broadcastTransbK", |
| "N": "broadcastTransbN", |
| "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1", |
| "hasBias": false, |
| "fScalar": "dtypes.T", |
| "outScalar": "dtypes.T", |
| "generalAddressing": true, |
| "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 2) % 64 != 0", |
| "outputBuffer": "\"y\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.BroadcastTransBSubgroupMatrix", |
| "shader": "matmul-subgroup-matrix-ext.wgsl.jinja", |
| "derive": { |
| "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []", |
| "aRank": "ranks.A if ranks.B > 2 else 2" |
| }, |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { "x": "ceilDiv(N,64)", "y": "ceilDiv(rowCount,32)", "z": "batchCount" } |
| } |
| ], |
| "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM,32) * ceilDiv(broadcastTransbN,64) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM,32) * 32 * ceilDiv(broadcastTransbN,64) * 64 * ceilDiv(broadcastTransbK,32) * 32 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"] |
| }, |
| { |
| "id": "m1_gemv_vec4", |
| "priority": 30, |
| "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 0) == 1", "dim(shapes.Y, 0) == 1", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "f16Ok(dtypes.T)"], |
| "derive": { |
| "unrollK2": "dtypes.T == \"f16\"", |
| "gemvScalar": "dtypes.T", |
| "gemvVector": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "alphaScale": "attrs.alpha" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.M1GemvVec4", |
| "shader": "matmul-vector-matrix-vec4.wgsl.jinja", |
| "bindings": [ |
| { "arg": "A", "name": "a", "elementType": "$gemvScalar" }, |
| { "arg": "B", "name": "b", "elementType": "$gemvVector" }, |
| { "arg": "Y", "name": "c", "elementType": "$gemvVector" }, |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" }, |
| { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "gemvWorkgroups" } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_band_vec4_splitk", |
| "priority": 11, |
| "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvLanes <= device.limits.maxComputeWorkgroupSizeX", "gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS", "bandSplitK >= 2", "bandSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "bandSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeWorkgroupSizeY", "gemvLanes * tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeInvocationsPerWorkgroup"], |
| "derive": { |
| "batched": false, |
| "outputBuffer": "\"y\"", |
| "M": "dim(shapes.A, 0)", |
| "K": "dim(shapes.A, 1)", |
| "N": "dim(shapes.B, 1)", |
| "gemvSlices": "tunables.BAND_SPLIT_SLICES", |
| "kSplits": "bandSplitK", |
| "split": "bandSplitK", |
| "workgroupSize": 256 |
| }, |
| "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[bandSplitK * numel(shapes.Y)]" }], |
| "passes": [ |
| { |
| "id": "partial", |
| "name": "TransposeMatMul.Rank2BandVec4SplitK", |
| "shader": "matmul-band-vec4.wgsl.jinja", |
| "bindings": ["a", "b", { "scratch": "partials", "name": "y", "elementType": "vec4<f32>" }], |
| "dispatch": { "x": "gemvWorkgroups", "y": "bandSplitK" } |
| }, |
| { |
| "id": "combine", |
| "name": "TransposeMatMul.Rank2BandVec4SplitKCombine", |
| "shader": "reduce-axis0-splitk-combine.wgsl.jinja", |
| "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" }, |
| "bindings": ["partials", "y", "params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (256)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_band_vec4", |
| "priority": 11, |
| "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"], |
| "derive": { |
| "batched": false, |
| "outputBuffer": "\"y\"", |
| "M": "dim(shapes.A, 0)", |
| "K": "dim(shapes.A, 1)", |
| "N": "dim(shapes.B, 1)" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.Rank2BandVec4", |
| "shader": "matmul-band-vec4.wgsl.jinja", |
| "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }], |
| "dispatch": { "x": "gemvWorkgroups" } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_band_vec4_f32_preferred", |
| "priority": 13, |
| "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"], |
| "demoteWhen": ["dtypes.T != \"f32\" or (dim(shapes.A, 0) > tunables.BAND_PREFER_MAX_ROWS and dim(shapes.A, 1) >= tunables.BAND_PREFER_DEEP_K)"], |
| "derive": { |
| "batched": false, |
| "outputBuffer": "\"y\"", |
| "M": "dim(shapes.A, 0)", |
| "K": "dim(shapes.A, 1)", |
| "N": "dim(shapes.B, 1)" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.Rank2BandVec4", |
| "shader": "matmul-band-vec4.wgsl.jinja", |
| "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }], |
| "dispatch": { "x": "gemvWorkgroups" } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_matrix_splitk", |
| "priority": 12, |
| "when": ["(dtypes.T == \"f16\" or dtypes.T == \"f32\") and f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_SPLITK_MIN_K", "dim(shapes.B, 1) % 64 == 0", "sgmatSplitK >= 2", "sgmatOutTiles < tunables.SUBGROUP_MATRIX_SPLITK_MAX_TILES", "sgmatSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "sgmatSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.Y, 1), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.Y, 0), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "subgroupMatrixResourcesFit", "wave32Effective"], |
| "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 } |
| ] |
| }, |
| "derive": { |
| "hasBias": false, |
| "generalAddressing": true, |
| "tailSafe": false, |
| "outputBuffer": "\"partials\"", |
| "outScalar": "\"f32\"", |
| "rowCount": "dim(shapes.A, 0)", |
| "K": "dim(shapes.A, 1)", |
| "N": "dim(shapes.B, 1)", |
| "batchCount": 1, |
| "splitK": "sgmatSplitK", |
| "kPerSplit": "dim(shapes.A, 1) / sgmatSplitK", |
| "split": "sgmatSplitK", |
| "workgroupSize": 256 |
| }, |
| "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[sgmatSplitK * numel(shapes.Y)]" }], |
| "passes": [ |
| { |
| "id": "partial", |
| "name": "TransposeMatMul.SubgroupMatrixSplitK", |
| "shader": "matmul-subgroup-matrix-ext.wgsl.jinja", |
| "derive": { "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"], "aRank": 2, "bRank": 2 }, |
| "bindings": ["a", "b_scalar", { "name": "partials", "elementType": "f32" }, "params_rows"], |
| "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" } |
| }, |
| { |
| "id": "combine", |
| "name": "TransposeMatMul.SubgroupMatrixSplitKCombine", |
| "shader": "reduce-axis0-splitk-combine.wgsl.jinja", |
| "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" }, |
| "bindings": ["partials", "y", "params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (256)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_matrix_tail_broadcast", |
| "priority": 11, |
| "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "(((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0)))) or (ranks.A == 4 and ranks.B == 2 and ranks.Y == 4 and sameShape(prefix(shapes.Y, 2), prefix(shapes.A, 2))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "hasBias": false, |
| "fScalar": "\"f16\"", |
| "outScalar": "\"f16\"", |
| "generalAddressing": true, |
| "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0", |
| "outputBuffer": "\"y\"", |
| "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else dim(shapes.A, ranks.A - 2)", |
| "K": "dim(shapes.A, ranks.A - 1)", |
| "N": "dim(shapes.B, ranks.B - 1)", |
| "batchCount": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) if ranks.B > 2 else 1" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.SubgroupMatrixTailBroadcast", |
| "shader": "matmul-subgroup-matrix-ext.wgsl.jinja", |
| "derive": { |
| "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []", |
| "aRank": "ranks.A if ranks.B > 2 else 2", |
| "bShape": "shapes.B" |
| }, |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "batchCount" } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_matrix", |
| "priority": 10, |
| "when": ["f16Ok(dtypes.T)", "ranks.A >= 2", "ranks.B == ranks.A", "ranks.Y == ranks.A", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) == (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2))", "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) >= tunables.SUBGROUP_MATRIX_MIN_M", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32Effective"], |
| "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 } |
| ] |
| }, |
| "derive": { |
| "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"", |
| "transA": "attrs.transA != 0", |
| "transB": "attrs.transB != 0", |
| "transBatchA": "false", |
| "rowCount": "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))", |
| "K": "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1))", |
| "N": "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", |
| "batchCount": "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)))" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.SubgroupMatrix", |
| "shader": "fused-matmul-subgroup-matrix.wgsl.jinja", |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "numel(shapes.Y) / (rowCount * N)" } |
| } |
| ] |
| }, |
| { |
| "id": "broadcast_rank4_tiled_reg", |
| "priority": 6, |
| "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 4", "(ranks.B == 2 or ranks.B == 3)", "ranks.Y == 4", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "(ranks.B == 2 or dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1)", "dim(shapes.Y, 1) == (dim(shapes.A, 1) if ranks.B == 2 else max(dim(shapes.A, 1), dim(shapes.B, 0)))", "dim(shapes.A, 3) == dim(shapes.B, ranks.B - 2)", "dim(shapes.Y, 2) == dim(shapes.A, 2)", "dim(shapes.Y, 3) == dim(shapes.B, ranks.B - 1)", "dim(shapes.A, 2) >= 64", "dim(shapes.A, 3) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "ceil(dim(shapes.B, ranks.B - 1) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 2) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.BroadcastRank4TiledReg", |
| "shader": "matmul-tiled-general-reg.wgsl.jinja", |
| "derive": { |
| "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []", |
| "aRank": "ranks.A if ranks.B > 2 else 2", |
| "K": "dim(shapes.A, ranks.A - 1)", |
| "bShape": "shapes.B", |
| "transBatchA": "false" |
| }, |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.B, ranks.B - 1) / registerTile)", |
| "y": "ceil(rowCount / registerTile)", |
| "z": "batchCount" |
| } |
| } |
| ], |
| "derive": { |
| "rowCount": "outer(shapes.A, 3) if ranks.B == 2 else dim(shapes.A, 2)", |
| "batchCount": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) if ranks.B > 2 else 1" |
| } |
| }, |
| { |
| "id": "plain_rank2_tiled_reg", |
| "priority": 4, |
| "demoteWhen": ["plainRank2RegDeepPreferredTier"], |
| "when": ["f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / registerTile) * ceil(dim(shapes.B, 1) / registerTile) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 0) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.PlainRank2TiledReg", |
| "shader": "matmul-tiled-general-reg.wgsl.jinja", |
| "derive": { |
| "rowCount": "dim(shapes.A, 0)", |
| "K": "dim(shapes.A, 1)", |
| "bShape": "shapes.B", |
| "transBatchA": "false" |
| }, |
| "bindings": ["a", "b_scalar", "y", "params_rows"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.B, 1) / registerTile)", |
| "y": "ceil(dim(shapes.A, 0) / registerTile)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "tiled", |
| "priority": 0, |
| "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else dim(shapes.A, ranks.A - 2))) == (dim(shapes.B, 0) if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2)))", "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) / 16) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2))) / 16) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2)))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "TransposeMatMul.Tiled", |
| "shader": "matmul-tiled-general.wgsl.jinja", |
| "derive": { |
| "aShape": "shapes.A", |
| "bShape": "shapes.B", |
| "transA": "attrs.transA != 0", |
| "transB": "attrs.transB != 0", |
| "transBatchA": "false", |
| "transBatchB": "false" |
| }, |
| "bindings": ["a", "b_scalar", "y"], |
| "dispatch": { |
| "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2))) / generalTile)", |
| "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) / generalTile)", |
| "z": "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2))))" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
