| { |
| "domain": "com.microsoft", |
| "name": "FusedMatMul", |
| "sinceVersion": 1, |
| "description": "Matrix product of two N-dimensional tensors `A` and `B`, following NumPy-style matrix-multiplication broadcasting. Supports optional transposition of either operand's last two dimensions, optional batch-dimension transposition, and a scalar `alpha` multiplier. Float32 and float16 are supported; double and bfloat16 are not.", |
| "inputs": [ |
| { "role": "A", "dtype": "T", "description": "N-dimensional matrix A." }, |
| { "role": "B", "dtype": "T", "description": "N-dimensional matrix B." } |
| ], |
| "outputs": [ |
| { |
| "role": "Y", |
| "dtype": "T", |
| "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)", |
| "shape": "matmulShape(logicalAShape, logicalBShape)", |
| "description": "Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions." |
| } |
| ], |
| "attributes": { "alpha": 1, "transA": 0, "transB": 0, "transBatchA": 0, "transBatchB": 0 }, |
| "attributeDescriptions": { |
| "alpha": "Scalar multiplier applied to the product of the input tensors.", |
| "transA": "When non-zero, transposes `A` on its last two dimensions before multiplication.", |
| "transB": "When non-zero, transposes `B` on its last two dimensions before multiplication.", |
| "transBatchA": "When non-zero, transposes `A` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication.", |
| "transBatchB": "When non-zero, transposes `B` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication." |
| }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "args": { |
| "A": { "kind": "tensor", "semantic": "A", "role": "input" }, |
| "B": { "kind": "tensor", "semantic": "B", "role": "input" }, |
| "Y": { "kind": "tensor", "semantic": "Y", "role": "output" } |
| }, |
| "tunables": { |
| "TILED_REG_MIN_WORKGROUPS": 64, |
| "GEMV_TARGET_BLOCKS": 512, |
| "SUBGROUP_MATRIX_MIN_M": 2, |
| "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": 512, |
| "SUBGROUP_MATRIX_SPLITK_MIN_K": 1024, |
| "SUBGROUP_MATRIX_SPLITK_MAX_TILES": 128 |
| }, |
| "bindingSets": { |
| "matrix": [ |
| { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, |
| { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, |
| { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" } |
| ], |
| "sgmatSplitPartial": [ |
| { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, |
| { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, |
| { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "sgmatSplitCombine": [ |
| { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] } |
| } |
| ], |
| "gemvF32": [ |
| { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "b", |
| "arg": "B", |
| "semantic": "B", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { "name": "c", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" }, |
| { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" } |
| ] |
| } |
| } |
| ] |
| }, |
| "derive": { |
| "gemvWorkgroups": "ceilDiv(dim(shapes.B, 1), 128)", |
| "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / 32), floor(device.limits.maxComputeWorkgroupStorageSize / 512))", |
| "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))", |
| "batchMovedAShape": "moveAxis(shapes.A, 0, -2) if attrs.transBatchA != 0 else shapes.A", |
| "batchMovedBShape": "moveAxis(shapes.B, 0, -2) if attrs.transBatchB != 0 else 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", |
| "transBatchContract": "(attrs.transBatchA == 0 and attrs.transBatchB == 0) or (ranks.A == ranks.B and ranks.A >= 3)", |
| "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", |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "fusedSgmatRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and attrs.transBatchA == 0 and attrs.transBatchB == 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)", |
| "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)", |
| "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)", |
| "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0", |
| "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0", |
| "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0", |
| "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0", |
| "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0", |
| "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))))" |
| }, |
| "variants": [ |
| { |
| "id": "f32_m1_gemv_vec4", |
| "priority": 30, |
| "when": ["dtypes.T == \"f32\"", "attrs.alpha == 1", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 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", "ceil(dim(shapes.B, 1) / 128) <= device.limits.maxComputeWorkgroupsPerDimension"], |
| "constants": { "T": "\"f32\"", "usesF16": false, "gemvSlices": "gemvSlices" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "FusedMatMul.F32M1GemvVec4", |
| "shader": "matmul-vector-matrix-vec4.wgsl.jinja", |
| "bindings": "gemvF32", |
| "dispatch": { "x": "ceil(dim(shapes.B, 1) / 128)" } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_matrix_splitk", |
| "description": "Partitions the K reduction across workgroups while retaining subgroup-matrix tiles, then combines float32 partials that already carry alpha. It serves small-M rank-two products whose unsplit tile grids cannot occupy the device.", |
| "priority": 12, |
| "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": ["(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 <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 1), 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 0), 32) <= device.limits.maxComputeWorkgroupsPerDimension", "subgroupMatrixResourcesFit", "wave32Effective"], |
| "constants": { |
| "usesF16": "dtypes.T == \"f16\"", |
| "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"", |
| "scalar": "dtypes.T", |
| "hasBias": false, |
| "generalAddressing": true, |
| "tailSafe": false, |
| "outputBuffer": "\"partials\"", |
| "outScalar": "\"f32\"", |
| "alpha": "attrs.alpha", |
| "M": "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": "FusedMatMul.SubgroupMatrixSplitK", |
| "source": { |
| "shader": "matmul-subgroup-matrix-ext.wgsl.jinja", |
| "inputs": { |
| "aShape": ["dim(shapes.A, 0)", "dim(shapes.A, 1)"], |
| "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"], |
| "aRank": 2, |
| "bRank": 2 |
| } |
| }, |
| "bindings": "sgmatSplitPartial", |
| "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" } |
| }, |
| { |
| "id": "combine", |
| "name": "FusedMatMul.SubgroupMatrixSplitKCombine", |
| "source": { |
| "shader": "reduce-axis0-splitk-combine.wgsl.jinja", |
| "inputs": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"", "intMode": false } |
| }, |
| "bindings": "sgmatSplitCombine", |
| "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": 256 } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_matrix_tail_broadcast", |
| "priority": 11, |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 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))))", "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) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"], |
| "constants": { |
| "hasBias": false, |
| "usesF16": true, |
| "fScalar": "\"f16\"", |
| "outScalar": "\"f16\"", |
| "scalar": "dtypes.T", |
| "generalAddressing": true, |
| "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0", |
| "outputBuffer": "\"y\"", |
| "alpha": "attrs.alpha", |
| "M": "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))" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "FusedMatMul.SubgroupMatrixTailBroadcast", |
| "source": { |
| "shader": "matmul-subgroup-matrix-ext.wgsl.jinja", |
| "inputs": { "aShape": "shapes.A", "bShape": "shapes.B", "aRank": "ranks.A", "bRank": "ranks.B" } |
| }, |
| "bindings": "matrix", |
| "dispatch": { "x": "ceil(constants.N / 64)", "y": "ceil(constants.M / 32)", "z": "constants.batchCount" } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_matrix", |
| "priority": 10, |
| "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": ["f16Ok(dtypes.T)", "attrs.transBatchA == 0 or (attrs.transA == 0 and ranks.A == 3)", "attrs.transBatchB == 0", "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, 0) if attrs.transBatchA != 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, 1) if attrs.transBatchA != 0 else 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, 0) if attrs.transBatchA != 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) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 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))) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"], |
| "constants": { |
| "usesF16": "dtypes.T == \"f16\"", |
| "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"", |
| "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"", |
| "scalar": "dtypes.T", |
| "transA": "attrs.transA != 0", |
| "transB": "attrs.transB != 0", |
| "transBatchA": "attrs.transBatchA != 0", |
| "M": "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchA != 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": "FusedMatMul.SubgroupMatrix", |
| "source": { "shader": "fused-matmul-subgroup-matrix.wgsl.jinja", "inputs": { "alpha": "attrs.alpha" } }, |
| "bindings": "matrix", |
| "dispatch": { |
| "x": "ceil(constants.N / 64)", |
| "y": "ceil(constants.M / 32)", |
| "z": "numel(shapes.Y) / (constants.M * constants.N)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "broadcast_rank4_tiled_reg", |
| "priority": 6, |
| "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 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) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension"], |
| "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "FusedMatMul.BroadcastRank4TiledReg", |
| "source": { |
| "shader": "matmul-tiled-general-reg.wgsl.jinja", |
| "inputs": { |
| "aShape": "shapes.A", |
| "bShape": "shapes.B", |
| "alpha": "attrs.alpha", |
| "aRank": "ranks.A", |
| "bRank": "ranks.B", |
| "transBatchA": "false" |
| } |
| }, |
| "bindings": "matrix", |
| "dispatch": { |
| "x": "ceil(dim(shapes.B, ranks.B - 1) / 64)", |
| "y": "ceil(dim(shapes.A, 2) / 64)", |
| "z": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "plain_rank2_tiled_reg", |
| "description": "Register-blocked rank-2 Y = alpha * A @ B for the tier without subgroup matrices. It is the non-transposed rung the two specialised register-blocked variants beside it never covered, so plain rank-2 shapes fell to the scalar-indexed tile.", |
| "priority": 4, |
| "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / 64) * ceil(dim(shapes.B, 1) / 64) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension"], |
| "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "FusedMatMul.PlainRank2TiledReg", |
| "source": { |
| "shader": "matmul-tiled-general-reg.wgsl.jinja", |
| "inputs": { |
| "aShape": "shapes.A", |
| "bShape": "shapes.B", |
| "alpha": "attrs.alpha", |
| "aRank": "ranks.A", |
| "bRank": "ranks.B", |
| "transBatchA": "false" |
| } |
| }, |
| "bindings": "matrix", |
| "dispatch": { "x": "ceil(dim(shapes.B, 1) / 64)", "y": "ceil(dim(shapes.A, 0) / 64)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "transbatch_a_tiled_reg", |
| "priority": 5, |
| "when": ["f16Ok(dtypes.T)", "attrs.transBatchA != 0", "attrs.transBatchB == 0", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 3", "ranks.B == 3", "ranks.Y == 3", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.B, 0)", "dim(shapes.A, 2) == dim(shapes.B, 1)", "dim(shapes.Y, 1) == dim(shapes.A, 0)", "dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceil(dim(shapes.B, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "dim(shapes.Y, 0) <= device.limits.maxComputeWorkgroupsPerDimension"], |
| "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "FusedMatMul.TransBatchATiledReg", |
| "source": { |
| "shader": "matmul-tiled-general-reg.wgsl.jinja", |
| "inputs": { |
| "aShape": "shapes.A", |
| "bShape": "shapes.B", |
| "alpha": "attrs.alpha", |
| "aRank": "ranks.A", |
| "bRank": "ranks.B", |
| "transBatchA": "true", |
| "kTile": "4" |
| } |
| }, |
| "bindings": "matrix", |
| "dispatch": { |
| "x": "ceil(dim(shapes.B, 2) / 64)", |
| "y": "ceil(dim(shapes.A, 0) / 64)", |
| "z": "dim(shapes.Y, 0)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "tiled", |
| "priority": 0, |
| "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "transBatchContract", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else (dim(shapes.A, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchB != 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, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "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, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))) <= device.limits.maxComputeWorkgroupsPerDimension"], |
| "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "FusedMatMul.Tiled", |
| "source": { |
| "shader": "matmul-tiled-general.wgsl.jinja", |
| "inputs": { |
| "aShape": "shapes.A", |
| "bShape": "shapes.B", |
| "transA": "attrs.transA != 0", |
| "transB": "attrs.transB != 0", |
| "alpha": "attrs.alpha", |
| "aRank": "ranks.A", |
| "bRank": "ranks.B", |
| "transBatchA": "attrs.transBatchA != 0", |
| "transBatchB": "attrs.transBatchB != 0" |
| } |
| }, |
| "bindings": "matrix", |
| "dispatch": { |
| "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 32)", |
| "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 32)", |
| "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, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))))" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
