| { |
| "domain": "ai.onnx", |
| "name": "MatMulInteger", |
| "sinceVersion": 10, |
| "inputs": { |
| "a": { "onnx": "A", "dtype": "TA" }, |
| "b": { "onnx": "B", "dtype": "TB" }, |
| "a_zero_point": { "dtype": "TA", "optional": true }, |
| "b_zero_point": { "dtype": "TB", "optional": true } |
| }, |
| "outputs": { |
| "y": { |
| "onnx": "Y", |
| "dtype": "TY", |
| "rank": "max(ranks.a, ranks.b) - (1 if ranks.a == 1 or ranks.b == 1 else 0)", |
| "shape": "matmulShape(shapes.a, shapes.b)" |
| } |
| }, |
| "typeConstraints": { "TA": ["uint8", "int8"], "TB": ["uint8", "int8"], "TY": ["int32"] }, |
| "tunables": { |
| "WORKGROUP_SIZE_X": { "default": 8 }, |
| "WORKGROUP_SIZE_Y": { "default": 8 }, |
| "DP4A_TILE_M": { "default": 8 }, |
| "DP4A_TILE_N": { "default": 64 }, |
| "DP4A_GEMV_TILE_N": { "default": 32 }, |
| "DP4A_GEMV_K_UNROLL": { "default": 8 }, |
| "DP4A_K_TILE_WORDS": { "default": 32 }, |
| "DP4A_LARGE_TILE_M": { "default": 32 }, |
| "DP4A_LARGE_TILE_N": { "default": 128 }, |
| "DP4A_LARGE_K_TILE_WORDS": { "default": 16 }, |
| "DP4A_LARGE_MIN_M": { "default": 256 }, |
| "DP4A_LARGE_MIN_N": { "default": 1024 }, |
| "DP4A_LARGE_MIN_K": { "default": 1024 }, |
| "DP4A_SHORT_K_TILE_M": { "default": 32 }, |
| "DP4A_SHORT_K_TILE_N": { "default": 128 }, |
| "DP4A_SHORT_K_TILE_WORDS": { "default": 32 }, |
| "DP4A_SHORT_K_MAX": { "default": 256 }, |
| "DP4A_SHORT_K_MIN_MN": { "default": 1024 }, |
| "DP4A_MEDIUM_K_TILE_M": { "default": 32 }, |
| "DP4A_MEDIUM_K_TILE_N": { "default": 64 }, |
| "DP4A_MEDIUM_K_TILE_WORDS": { "default": 32 }, |
| "DP4A_MEDIUM_K_MAX": { "default": 512 }, |
| "DP4A_MEDIUM_K_MIN_MN": { "default": 1024 }, |
| "EXACT_MATRIX_MIN_INVOCATIONS": { "default": 16384 } |
| }, |
| "derive": { |
| "precastRows": "dim(shapes.a, 0)", |
| "precastContraction": "dim(shapes.a, 1)", |
| "precastCols": "dim(shapes.b, 1)", |
| "precastBatches": 1, |
| "sgmatTileRows": "dim(shapes.a, 0)", |
| "aUnsigned": "tensorDtypes.a == \"uint8\"", |
| "bUnsigned": "tensorDtypes.b == \"uint8\"", |
| "aZeroPerRow": false, |
| "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", |
| "sgmatTileN": "64", |
| "sgmatTileK": "32", |
| "sgmatColumnGroups": "2", |
| "sgmatStorageRowGroups": "floor((device.limits.maxComputeWorkgroupStorageSize - sgmatTileN * sgmatTileK * 4) / (16 * sgmatTileK * 4))", |
| "sgmatRowGroups": "max(1, min(4, floor(deviceWorkgroupCap / (32 * sgmatColumnGroups)), sgmatStorageRowGroups, max(floor(sgmatStorageRowGroups / 2), floor(sgmatTileRows / 16))))", |
| "sgmatTileM": "16 * sgmatRowGroups", |
| "sgmatSubgroups": "sgmatRowGroups * sgmatColumnGroups", |
| "sgmatWorkgroup": "32 * sgmatSubgroups", |
| "sgmatStagingBytes": "(sgmatTileM + sgmatTileN) * sgmatTileK * 4", |
| "sgmatScratchBytes": "sgmatSubgroups * 4 * 64 * 4", |
| "sgmatReuseStaging": "sgmatStagingBytes + sgmatScratchBytes > device.limits.maxComputeWorkgroupStorageSize", |
| "sgmatSharedBytes": "sgmatStagingBytes + (0 if sgmatReuseStaging else sgmatScratchBytes)", |
| "sgmatDeviceOk": "sgmatWorkgroup <= deviceWorkgroupCap and sgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize", |
| "precastSubgroups": "max(1, min(4, ceilDiv(precastRows, 16), floor(deviceWorkgroupCap / 32), floor(device.limits.maxComputeWorkgroupStorageSize / (8 * 64 * 4))))", |
| "precastTileM": "16 * precastSubgroups", |
| "precastTileN": "4 * 8", |
| "precastTileK": "4 * 8", |
| "precastWorkgroup": "32 * precastSubgroups", |
| "preparedK": "ceilDiv(precastContraction, precastTileK) * precastTileK", |
| "preparedN": "ceilDiv(precastCols, precastTileN) * precastTileN", |
| "preparedElements": "ceilDiv(precastRows, precastTileM) * precastTileM * preparedK", |
| "preparedBElements": "precastBatches * preparedK * preparedN", |
| "prepareWorkgroup": "min(256, deviceWorkgroupCap)", |
| "prepareTotalElements": "preparedElements + preparedBElements", |
| "prepareGroups": "ceilDiv(prepareTotalElements, 4 * prepareWorkgroup)", |
| "prepareDispatchElements": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) * min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) * prepareWorkgroup * 4", |
| "precastResourcesOk": "deviceWorkgroupCap >= precastWorkgroup and precastSubgroups * 8 * 64 * 4 <= device.limits.maxComputeWorkgroupStorageSize", |
| "precastDispatchOk": "prepareTotalElements <= 4294967295 and prepareDispatchElements <= 4294967296 and ceilDiv(precastCols, precastTileN) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(precastRows, precastTileM) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and prepareGroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "precastF32BuffersOk": "preparedElements * 4 <= device.limits.maxStorageBufferBindingSize and preparedElements * 4 <= device.limits.maxBufferSize and preparedBElements * 4 <= device.limits.maxStorageBufferBindingSize and preparedBElements * 4 <= device.limits.maxBufferSize", |
| "precastF16BuffersOk": "preparedElements * 2 <= device.limits.maxStorageBufferBindingSize and preparedElements * 2 <= device.limits.maxBufferSize and preparedBElements * 2 <= device.limits.maxStorageBufferBindingSize and preparedBElements * 2 <= device.limits.maxBufferSize", |
| "precastWorthwhile": "dim(shapes.a, 0) >= 16 and dim(shapes.b, 1) >= 32 and dim(shapes.a, 1) >= 32 and floor(dim(shapes.a, 0) / 16) * floor(dim(shapes.b, 1) / 32) * 32 >= tunables.EXACT_MATRIX_MIN_INVOCATIONS", |
| "portableWorkgroupCols": 16, |
| "portableWorkgroupRows": "pow2ceil(min(16, max(4, pow2ceil(ceilDiv(precastRows, 4))), floor(deviceWorkgroupCap / portableWorkgroupCols), device.limits.maxComputeWorkgroupSizeY) + 1) / 2", |
| "portableTileM": "4 * portableWorkgroupRows", |
| "portableTileN": "4 * portableWorkgroupCols", |
| "portableTileK": 32, |
| "portablePartialK": "floor(16777216 / (255 * 255) / portableTileK) * portableTileK", |
| "portableGeometryOk": "portableWorkgroupRows >= 4 and portableWorkgroupCols <= device.limits.maxComputeWorkgroupSizeX and portablePartialK >= portableTileK and ceilDiv(precastRows, portableTileM) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "rank2MatMulShapeOk": "ranks.a == 2 and ranks.b == 2 and ranks.y == 2 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)", |
| "rank3ByRank3ShapeOk": "ranks.a == 3 and ranks.b == 3 and ranks.y == 3 and dim(shapes.a, 2) == dim(shapes.b, 1) and dim(shapes.y, 0) == max(dim(shapes.a, 0), dim(shapes.b, 0)) and (dim(shapes.a, 0) == dim(shapes.y, 0) or dim(shapes.a, 0) == 1) and (dim(shapes.b, 0) == dim(shapes.y, 0) or dim(shapes.b, 0) == 1) and dim(shapes.y, 1) == dim(shapes.a, 1) and dim(shapes.y, 2) == dim(shapes.b, 2)", |
| "rank3ByRank2ShapeOk": "ranks.a == 3 and ranks.b == 2 and ranks.y == 3 and dim(shapes.a, 2) == dim(shapes.b, 0) and dim(shapes.y, 0) == dim(shapes.a, 0) and dim(shapes.y, 1) == dim(shapes.a, 1) and dim(shapes.y, 2) == dim(shapes.b, 1)", |
| "rank1DotShapeOk": "ranks.a == 1 and ranks.b == 1 and ranks.y == 0 and dim(shapes.a, 0) == dim(shapes.b, 0)", |
| "rank4Dp4aShapeOk": "ranks.a == 4 and ranks.b == 4 and ranks.y == 4 and dim(shapes.a, 3) == dim(shapes.b, 2) and dim(shapes.y, 0) == dim(shapes.a, 0) and dim(shapes.y, 0) == dim(shapes.b, 0) and dim(shapes.y, 1) == dim(shapes.a, 1) and dim(shapes.y, 1) == dim(shapes.b, 1) and dim(shapes.y, 2) == dim(shapes.a, 2) and dim(shapes.y, 3) == dim(shapes.b, 3)", |
| "rank4BroadcastShapeOk": "ranks.a == 4 and ranks.b == 4 and ranks.y == 4 and (dim(shapes.a, 0) == dim(shapes.y, 0) or dim(shapes.a, 0) == 1) and (dim(shapes.a, 1) == dim(shapes.y, 1) or dim(shapes.a, 1) == 1) and (dim(shapes.b, 0) == dim(shapes.y, 0) or dim(shapes.b, 0) == 1) and (dim(shapes.b, 1) == dim(shapes.y, 1) or dim(shapes.b, 1) == 1) and dim(shapes.y, 0) == max(dim(shapes.a, 0), dim(shapes.b, 0)) and dim(shapes.y, 1) == max(dim(shapes.a, 1), dim(shapes.b, 1)) and dim(shapes.a, 3) == dim(shapes.b, 2) and dim(shapes.y, 2) == dim(shapes.a, 2) and dim(shapes.y, 3) == dim(shapes.b, 3)", |
| "scalarAZeroOk": "ranks.a_zero_point == 0 or (ranks.a_zero_point == 1 and dim(shapes.a_zero_point, 0) == 1)", |
| "scalarBZeroOk": "ranks.b_zero_point == 0 or (ranks.b_zero_point == 1 and dim(shapes.b_zero_point, 0) == 1)", |
| "requiredScalarAZeroOk": "present.a_zero_point and scalarAZeroOk", |
| "requiredScalarBZeroOk": "present.b_zero_point and scalarBZeroOk", |
| "rank2PerColumnBZeroOk": "ranks.b_zero_point == 1 and dim(shapes.b_zero_point, 0) == dim(shapes.y, 1)", |
| "rank2Dp4aLargePrefill": "ranks.y == 2 and ranks.a == 2 and dim(shapes.y, 0) >= tunables.DP4A_LARGE_MIN_M and dim(shapes.y, 1) >= tunables.DP4A_LARGE_MIN_N and dim(shapes.a, 1) >= tunables.DP4A_LARGE_MIN_K", |
| "perColumnDp4aUseNarrowTile": "variableSubgroup16To32 and rank2Dp4aLargePrefill and not present.a_zero_point", |
| "rank2Dp4aWideShortK": "ranks.y == 2 and ranks.a == 2 and dim(shapes.y, 0) >= tunables.DP4A_SHORT_K_MIN_MN and dim(shapes.y, 1) >= tunables.DP4A_SHORT_K_MIN_MN and dim(shapes.a, 1) <= tunables.DP4A_SHORT_K_MAX", |
| "rank2Dp4aMediumK": "ranks.y == 2 and ranks.a == 2 and dim(shapes.y, 0) >= tunables.DP4A_MEDIUM_K_MIN_MN and dim(shapes.y, 1) >= tunables.DP4A_MEDIUM_K_MIN_MN and dim(shapes.a, 1) > tunables.DP4A_SHORT_K_MAX and dim(shapes.a, 1) <= tunables.DP4A_MEDIUM_K_MAX", |
| "rank2Dp4aTileM": "1 if ranks.y < 2 or dim(shapes.y, 0) == 1 else (tunables.DP4A_SHORT_K_TILE_M if rank2Dp4aWideShortK else (tunables.DP4A_MEDIUM_K_TILE_M if rank2Dp4aMediumK else (tunables.DP4A_LARGE_TILE_M if rank2Dp4aLargePrefill else tunables.DP4A_TILE_M)))", |
| "rank2Dp4aTileN": "tunables.DP4A_GEMV_TILE_N if ranks.y < 2 or dim(shapes.y, 0) == 1 else (tunables.DP4A_SHORT_K_TILE_N if rank2Dp4aWideShortK else (tunables.DP4A_MEDIUM_K_TILE_N if rank2Dp4aMediumK else (tunables.DP4A_LARGE_TILE_N if rank2Dp4aLargePrefill else tunables.DP4A_TILE_N)))", |
| "rank2Dp4aKtWords": "tunables.DP4A_SHORT_K_TILE_WORDS if rank2Dp4aWideShortK else (tunables.DP4A_MEDIUM_K_TILE_WORDS if rank2Dp4aMediumK else (tunables.DP4A_LARGE_K_TILE_WORDS if rank2Dp4aLargePrefill else tunables.DP4A_K_TILE_WORDS))", |
| "rank2Dp4aStorageBytes": "rank2Dp4aTileM * rank2Dp4aKtWords * dtypeBytes(\"uint32\") + rank2Dp4aTileM * dtypeBytes(\"int32\")", |
| "rank2Dp4aResourcesOk": "tunables.DP4A_GEMV_K_UNROLL >= 1 and tunables.DP4A_GEMV_K_UNROLL <= 16 and rank2Dp4aTileN <= deviceWorkgroupCap and rank2Dp4aStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", |
| "rank2Dp4aDispatchOk": "rank2Dp4aResourcesOk and ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "rank3Dp4aDispatchOk": "ceilDiv(dim(shapes.y, 2), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.y, 1), 8) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.y, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "rank4Dp4aDispatchOk": "ceilDiv(dim(shapes.y, 3), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.y, 2), 8) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.y, 0) * dim(shapes.y, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "matrixM": "dim(shapes.y, 0) if ranks.y >= 1 else 1", |
| "matrixN": "dim(shapes.y, 1) if ranks.y >= 2 else 1", |
| "matrixK": "dim(shapes.a, 1) if ranks.a >= 2 else dim(shapes.a, 0)", |
| "batchM": "matrixN", |
| "batchN": "dim(shapes.y, 2) if ranks.y >= 3 else 1", |
| "batchK": "dim(shapes.a, ranks.a - 1) if ranks.a >= 1 else 1", |
| "batchAStride": "0 if ranks.a < 3 or dim(shapes.a, 0) == 1 else dim(shapes.a, 1) * dim(shapes.a, 2)", |
| "batchBStride": "0 if ranks.b < 3 or dim(shapes.b, 0) == 1 else dim(shapes.b, 1) * dim(shapes.b, 2)", |
| "batchYStride": "dim(shapes.y, 1) * dim(shapes.y, 2) if ranks.y >= 3 else 1", |
| "batchCount": "matrixM", |
| "rank4M": "batchN", |
| "rank4N": "dim(shapes.y, 3) if ranks.y >= 4 else 1", |
| "rank4K": "dim(shapes.a, 3) if ranks.a >= 4 else 1", |
| "rank4ABatchStride": "dim(shapes.a, 2) * dim(shapes.a, 3) if ranks.a >= 4 else 1", |
| "rank4BBatchStride": "dim(shapes.b, 2) * dim(shapes.b, 3) if ranks.b >= 4 else 1", |
| "rank4YBatchStride": "dim(shapes.y, 2) * dim(shapes.y, 3) if ranks.y >= 4 else 1", |
| "aScalar": "\"u32\" if tensorDtypes.a == \"uint8\" else \"i32\"", |
| "bScalar": "\"u32\" if tensorDtypes.b == \"uint8\" else \"i32\"", |
| "sgmatM": "matrixM", |
| "sgmatN": "matrixN", |
| "sgmatK": "matrixK", |
| "sgmatDispatchOk": "ceilDiv(sgmatN, sgmatTileN) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(sgmatM, sgmatTileM) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "sgmatParallelInvocations": "floor(sgmatM / sgmatTileM) * floor(sgmatN / sgmatTileN) * sgmatWorkgroup", |
| "sgmatWorthwhile": "sgmatM >= sgmatTileM and sgmatN >= sgmatTileN / 2 and sgmatK >= sgmatTileK and sgmatParallelInvocations >= tunables.EXACT_MATRIX_MIN_INVOCATIONS", |
| "preferPackedDot": "variableSubgroup16To32 and device.wgslLanguageFeatures.has(\"packed_4x8_integer_dot_product\") and rank2MatMulShapeOk and rank2Dp4aDispatchOk and device.limits.maxComputeInvocationsPerWorkgroup >= 128 and device.limits.maxComputeWorkgroupSizeX >= 128 and device.limits.maxComputeWorkgroupStorageSize >= 8448 and matrixM >= tunables.DP4A_LARGE_MIN_M and matrixN >= tunables.DP4A_LARGE_MIN_N and matrixK > tunables.DP4A_MEDIUM_K_MAX / 2 and matrixK % 4 == 0 and matrixM * matrixN * matrixK >= 4 * tunables.DP4A_LARGE_MIN_M * tunables.DP4A_LARGE_MIN_N * tunables.DP4A_MEDIUM_K_MAX", |
| "tileM": "rank2Dp4aTileM", |
| "tileN": "rank2Dp4aTileN", |
| "ktWords": "rank2Dp4aKtWords" |
| }, |
| "bindings": { |
| "a": { "elementType": "$aStorage" }, |
| "b": { "elementType": "$bScalar" }, |
| "a_zero_point": { "elementType": "$aScalar", "length": 1 }, |
| "b_zero_point": { "elementType": "$bScalar" }, |
| "y": { "elementType": "i32" }, |
| "params": { |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "matrixM" }, |
| { "name": "N", "type": "u32", "value": "matrixN" }, |
| { "name": "K", "type": "u32", "value": "matrixK" }, |
| { "name": "aBatchStride4", "type": "u32", "value": 0 }, |
| { "name": "bBatchStride", "type": "u32", "value": 0 }, |
| { "name": "yBatchStride", "type": "u32", "value": 0 } |
| ] |
| }, |
| "a_main": { "name": "a", "elementType": "$aScalar" }, |
| "a_a_vec4": { "name": "a", "elementType": "$aVec4" }, |
| "params_main": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "matrixM" }, |
| { "name": "N", "type": "u32", "value": "matrixN" }, |
| { "name": "K", "type": "u32", "value": "matrixK" } |
| ] |
| }, |
| "params_m_n_k": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": 1 }, |
| { "name": "N", "type": "u32", "value": 1 }, |
| { "name": "K", "type": "u32", "value": "dim(shapes.a, 0)" } |
| ] |
| }, |
| "params__uniform": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "batchM" }, |
| { "name": "N", "type": "u32", "value": "batchN" }, |
| { "name": "K", "type": "u32", "value": "batchK" }, |
| { "name": "aBatchStride", "type": "u32", "value": "batchAStride" }, |
| { "name": "bBatchStride", "type": "u32", "value": "batchBStride" }, |
| { "name": "yBatchStride", "type": "u32", "value": "batchYStride" }, |
| { "name": "batchCount", "type": "u32", "value": "batchCount" } |
| ] |
| }, |
| "params_rank4_no_zero_points": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "rank4M" }, |
| { "name": "N", "type": "u32", "value": "rank4N" }, |
| { "name": "K", "type": "u32", "value": "rank4K" }, |
| { "name": "batch0", "type": "u32", "value": "dim(shapes.y, 0)" }, |
| { "name": "batch1", "type": "u32", "value": "dim(shapes.y, 1)" }, |
| { |
| "name": "aBatch0Stride", |
| "type": "u32", |
| "value": "0 if dim(shapes.a, 0) == 1 else dim(shapes.a, 1) * dim(shapes.a, 2) * dim(shapes.a, 3)" |
| }, |
| { |
| "name": "aBatch1Stride", |
| "type": "u32", |
| "value": "0 if dim(shapes.a, 1) == 1 else dim(shapes.a, 2) * dim(shapes.a, 3)" |
| }, |
| { |
| "name": "bBatch0Stride", |
| "type": "u32", |
| "value": "0 if dim(shapes.b, 0) == 1 else dim(shapes.b, 1) * dim(shapes.b, 2) * dim(shapes.b, 3)" |
| }, |
| { |
| "name": "bBatch1Stride", |
| "type": "u32", |
| "value": "0 if dim(shapes.b, 1) == 1 else dim(shapes.b, 2) * dim(shapes.b, 3)" |
| } |
| ] |
| }, |
| "params_quant_dp4a_matmul": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "batchM" }, |
| { "name": "N", "type": "u32", "value": "batchN" }, |
| { "name": "K", "type": "u32", "value": "batchK" }, |
| { "name": "aBatchStride4", "type": "u32", "value": "batchAStride / 4" }, |
| { "name": "bBatchStride", "type": "u32", "value": "batchBStride" }, |
| { "name": "yBatchStride", "type": "u32", "value": "batchYStride" } |
| ] |
| }, |
| "params_dp4a_rank4_no_zero_points": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "rank4M" }, |
| { "name": "N", "type": "u32", "value": "rank4N" }, |
| { "name": "K", "type": "u32", "value": "rank4K" }, |
| { "name": "aBatchStride4", "type": "u32", "value": "rank4ABatchStride / 4" }, |
| { "name": "bBatchStride", "type": "u32", "value": "rank4BBatchStride" }, |
| { "name": "yBatchStride", "type": "u32", "value": "rank4YBatchStride" } |
| ] |
| }, |
| "precast_f16_aFloat": { "name": "aFloat", "buffer": "storage", "elementType": "f16" }, |
| "precast_f32_aFloat": { "name": "aFloat", "buffer": "storage", "elementType": "f32" }, |
| "precast_f16_bFloat": { "name": "bFloat", "buffer": "storage", "elementType": "f16" }, |
| "precast_f32_bFloat": { "name": "bFloat", "buffer": "storage", "elementType": "f32" }, |
| "precast_prepare_params": { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "sgmatM" }, |
| { "name": "K", "type": "u32", "value": "sgmatK" }, |
| { "name": "paddedK", "type": "u32", "value": "preparedK" }, |
| { "name": "count", "type": "u32", "value": "prepareTotalElements" }, |
| { |
| "name": "groupsX", |
| "type": "u32", |
| "value": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| }, |
| { "name": "N", "type": "u32", "value": "sgmatN" }, |
| { "name": "paddedN", "type": "u32", "value": "preparedN" }, |
| { "name": "aCount", "type": "u32", "value": "preparedElements" } |
| ] |
| }, |
| "precast_f16_a": { "name": "a", "scratch": "aFloat", "buffer": "read-only-storage", "elementType": "f16" }, |
| "precast_f32_a": { "name": "a", "scratch": "aFloat", "buffer": "read-only-storage", "elementType": "f32" }, |
| "precast_f16_b": { "name": "b", "scratch": "bFloat", "buffer": "read-only-storage", "elementType": "f16" }, |
| "precast_f32_b": { "name": "b", "scratch": "bFloat", "buffer": "read-only-storage", "elementType": "f32" }, |
| "portable_f16_a": { "name": "a", "scratch": "aFloat", "buffer": "read-only-storage", "elementType": "vec4<f16>" }, |
| "portable_f16_b": { "name": "b", "scratch": "bFloat", "buffer": "read-only-storage", "elementType": "vec4<f16>" }, |
| "portable_f32_a": { "name": "a", "scratch": "aFloat", "buffer": "read-only-storage", "elementType": "vec4<f32>" }, |
| "portable_f32_b": { "name": "b", "scratch": "bFloat", "buffer": "read-only-storage", "elementType": "vec4<f32>" } |
| }, |
| "variants": [ |
| { |
| "id": "sgmat_precast_a0_bnone_f16", |
| "priority": 55, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "not present.b_zero_point", "wave32Effective", "precastResourcesOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix", "shader-f16"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f16\"", |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f16_aFloat", "b", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f16_a", "precast_f16_b", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a0_bnone_f16", |
| "priority": 39, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "not present.b_zero_point", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float16\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f16\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": ["shader-f16"] }, |
| "derive": { |
| "precastScalar": "\"f16\"", |
| "scalar": "\"f16\"", |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f16_aFloat", "b", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f16_a", "portable_f16_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a0_bnone", |
| "priority": 50, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "not present.b_zero_point", "wave32Effective", "precastResourcesOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f32\"", |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f32_aFloat", "b", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f32_a", "precast_f32_b", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a0_bnone", |
| "priority": 38, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "not present.b_zero_point", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f32\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": [] }, |
| "derive": { |
| "precastScalar": "\"f32\"", |
| "scalar": "\"f32\"", |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f32_aFloat", "b", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f32_a", "portable_f32_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a0_bscalar_f16", |
| "priority": 55, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "requiredScalarBZeroOk", "wave32Effective", "precastResourcesOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix", "shader-f16"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f16\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f16_a", "precast_f16_b", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a0_bscalar_f16", |
| "priority": 39, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "requiredScalarBZeroOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float16\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f16\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": ["shader-f16"] }, |
| "derive": { |
| "precastScalar": "\"f16\"", |
| "scalar": "\"f16\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f16_a", "portable_f16_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a0_bscalar", |
| "priority": 50, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "requiredScalarBZeroOk", "wave32Effective", "precastResourcesOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f32\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f32_a", "precast_f32_b", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a0_bscalar", |
| "priority": 38, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "requiredScalarBZeroOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f32\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": [] }, |
| "derive": { |
| "precastScalar": "\"f32\"", |
| "scalar": "\"f32\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f32_a", "portable_f32_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a0_bcolumn_f16", |
| "priority": 55, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "present.b_zero_point and rank2PerColumnBZeroOk", "wave32Effective", "precastResourcesOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix", "shader-f16"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f16\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": true |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f16_a", "precast_f16_b", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a0_bcolumn_f16", |
| "priority": 39, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "present.b_zero_point and rank2PerColumnBZeroOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float16\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f16\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": ["shader-f16"] }, |
| "derive": { |
| "precastScalar": "\"f16\"", |
| "scalar": "\"f16\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": true, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f16_a", "portable_f16_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a0_bcolumn", |
| "priority": 50, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "present.b_zero_point and rank2PerColumnBZeroOk", "wave32Effective", "precastResourcesOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f32\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": true |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f32_a", "precast_f32_b", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a0_bcolumn", |
| "priority": 38, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "present.b_zero_point and rank2PerColumnBZeroOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f32\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": [] }, |
| "derive": { |
| "precastScalar": "\"f32\"", |
| "scalar": "\"f32\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": true, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f32_a", "portable_f32_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a1_bnone_f16", |
| "priority": 55, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "not present.b_zero_point", "wave32Effective", "precastResourcesOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix", "shader-f16"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f16\"", |
| "hasAZero": true, |
| "hasBZero": false, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f16_aFloat", "b", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f16_a", "precast_f16_b", "a_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a1_bnone_f16", |
| "priority": 39, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "not present.b_zero_point", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float16\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f16\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": ["shader-f16"] }, |
| "derive": { |
| "precastScalar": "\"f16\"", |
| "scalar": "\"f16\"", |
| "hasAZero": true, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f16_aFloat", "b", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f16_a", "portable_f16_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a1_bnone", |
| "priority": 50, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "not present.b_zero_point", "wave32Effective", "precastResourcesOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f32\"", |
| "hasAZero": true, |
| "hasBZero": false, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f32_aFloat", "b", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f32_a", "precast_f32_b", "a_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a1_bnone", |
| "priority": 38, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "not present.b_zero_point", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f32\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": [] }, |
| "derive": { |
| "precastScalar": "\"f32\"", |
| "scalar": "\"f32\"", |
| "hasAZero": true, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f32_aFloat", "b", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f32_a", "portable_f32_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a1_bscalar_f16", |
| "priority": 55, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "requiredScalarBZeroOk", "wave32Effective", "precastResourcesOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix", "shader-f16"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f16\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f16_a", "precast_f16_b", "a_zero_point", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a1_bscalar_f16", |
| "priority": 39, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "requiredScalarBZeroOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float16\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f16\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": ["shader-f16"] }, |
| "derive": { |
| "precastScalar": "\"f16\"", |
| "scalar": "\"f16\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f16_a", "portable_f16_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a1_bscalar", |
| "priority": 50, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "requiredScalarBZeroOk", "wave32Effective", "precastResourcesOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f32\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f32_a", "precast_f32_b", "a_zero_point", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a1_bscalar", |
| "priority": 38, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "requiredScalarBZeroOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f32\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": [] }, |
| "derive": { |
| "precastScalar": "\"f32\"", |
| "scalar": "\"f32\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f32_a", "portable_f32_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a1_bcolumn_f16", |
| "priority": 55, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "present.b_zero_point and rank2PerColumnBZeroOk", "wave32Effective", "precastResourcesOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix", "shader-f16"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f16\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": true |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f16_a", "precast_f16_b", "a_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a1_bcolumn_f16", |
| "priority": 39, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "present.b_zero_point and rank2PerColumnBZeroOk", "precastF16BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float16\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f16\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": ["shader-f16"] }, |
| "derive": { |
| "precastScalar": "\"f16\"", |
| "scalar": "\"f16\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": true, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f16_aFloat", "b", "b_zero_point", "precast_f16_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f16_a", "portable_f16_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float16", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float16", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_precast_a1_bcolumn", |
| "priority": 50, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "present.b_zero_point and rank2PerColumnBZeroOk", "wave32Effective", "precastResourcesOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { |
| "precastOperands": true, |
| "sgmatTileM": "precastTileM", |
| "sgmatTileN": "precastTileN", |
| "sgmatTileK": "precastTileK", |
| "sgmatRowGroups": "precastSubgroups", |
| "sgmatSubgroups": "precastSubgroups", |
| "sgmatWorkgroup": "precastWorkgroup", |
| "precastScalar": "\"f32\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": true |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMultiply", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["precast_f32_a", "precast_f32_b", "a_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(sgmatN, sgmatTileN)", "y": "ceilDiv(sgmatM, sgmatTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "portable_precast_a1_bcolumn", |
| "priority": 38, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "present.b_zero_point and rank2PerColumnBZeroOk", "precastF32BuffersOk", "precastDispatchOk", "precastWorthwhile", "portableGeometryOk", "(portableTileM + portableTileN) * portableTileK * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "\"f32\" == \"f16\" or sgmatM >= portableTileN"], |
| "demoteWhen": ["preferPackedDot"], |
| "requires": { "features": [] }, |
| "derive": { |
| "precastScalar": "\"f32\"", |
| "scalar": "\"f32\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": true, |
| "M": "sgmatM", |
| "N": "sgmatN", |
| "K": "preparedK", |
| "kTile": "portableTileK", |
| "workgroupCols": "portableWorkgroupCols", |
| "workgroupRows": "portableWorkgroupRows", |
| "blockRows": "portableTileM", |
| "blockCols": "portableTileN", |
| "integerPartialK": "portablePartialK" |
| }, |
| "passes": [ |
| { |
| "id": "prepare", |
| "name": "MatMulInteger.ExactPrepareOperands", |
| "shader": "quant-exact-prepare.wgsl.jinja", |
| "bindings": ["a_main", "a_zero_point", "precast_f32_aFloat", "b", "b_zero_point", "precast_f32_bFloat", "precast_prepare_params"], |
| "dispatch": { |
| "x": "min(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "y": "ceilDiv(prepareGroups, min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| }, |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactPortableMultiply", |
| "shader": "quant-exact-portable.wgsl.jinja", |
| "bindings": ["portable_f32_a", "portable_f32_b", "y"], |
| "dispatch": { "x": "ceilDiv(sgmatN, portableTileN)", "y": "ceilDiv(sgmatM, portableTileM)" } |
| } |
| ], |
| "intermediates": [ |
| { "id": "aFloat", "dtype": "float32", "shape": "[preparedElements]" }, |
| { "id": "bFloat", "dtype": "float32", "shape": "[preparedBElements]" } |
| ] |
| }, |
| { |
| "id": "sgmat_exact_a0_bnone", |
| "priority": 40, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "not present.b_zero_point", "sgmatDeviceOk", "sgmatDispatchOk", "sgmatWorthwhile", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { "hasAZero": false, "hasBZero": false, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMatrix", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["a_main", "b", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(matrixN, sgmatTileN)", "y": "ceilDiv(matrixM, sgmatTileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_exact_a0_bscalar", |
| "priority": 40, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "requiredScalarBZeroOk", "sgmatDeviceOk", "sgmatDispatchOk", "sgmatWorthwhile", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMatrix", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(matrixN, sgmatTileN)", "y": "ceilDiv(matrixM, sgmatTileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_exact_a0_bcolumn", |
| "priority": 40, |
| "when": ["rank2MatMulShapeOk", "not present.a_zero_point", "present.b_zero_point and rank2PerColumnBZeroOk", "sgmatDeviceOk", "sgmatDispatchOk", "sgmatWorthwhile", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroPerColumn": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMatrix", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(matrixN, sgmatTileN)", "y": "ceilDiv(matrixM, sgmatTileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_exact_a1_bnone", |
| "priority": 40, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "not present.b_zero_point", "sgmatDeviceOk", "sgmatDispatchOk", "sgmatWorthwhile", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { "hasAZero": true, "hasBZero": false, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMatrix", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(matrixN, sgmatTileN)", "y": "ceilDiv(matrixM, sgmatTileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_exact_a1_bscalar", |
| "priority": 40, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "requiredScalarBZeroOk", "sgmatDeviceOk", "sgmatDispatchOk", "sgmatWorthwhile", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMatrix", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(matrixN, sgmatTileN)", "y": "ceilDiv(matrixM, sgmatTileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "sgmat_exact_a1_bcolumn", |
| "priority": 40, |
| "when": ["rank2MatMulShapeOk", "requiredScalarAZeroOk", "present.b_zero_point and rank2PerColumnBZeroOk", "sgmatDeviceOk", "sgmatDispatchOk", "sgmatWorthwhile", "wave32Effective"], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroPerColumn": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.ExactMatrix", |
| "shader": "quant-exact-matrix.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params_main"], |
| "dispatch": { "x": "ceilDiv(matrixN, sgmatTileN)", "y": "ceilDiv(matrixM, sgmatTileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_scalar_zero_points_tail", |
| "priority": 21, |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "requiredScalarAZeroOk", "requiredScalarBZeroOk", "dim(shapes.a, 1) >= 32", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 != 0", "rank2Dp4aDispatchOk"], |
| "derive": { |
| "aStorage": "\"u32\" if tensorDtypes.a == \"uint8\" else \"i32\"", |
| "aPacked": false, |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aTail", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN)", |
| "y": "ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_no_zero_points_tail", |
| "priority": 21, |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "rank2MatMulShapeOk", "dim(shapes.a, 1) >= 32", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 != 0", "rank2Dp4aDispatchOk"], |
| "derive": { |
| "aStorage": "\"u32\" if tensorDtypes.a == \"uint8\" else \"i32\"", |
| "aPacked": false, |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aTailNoZero", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a", "b", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN)", |
| "y": "ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank3_by_rank2_scalar_zero_points_tail", |
| "priority": 21, |
| "when": ["requiredScalarAZeroOk", "requiredScalarBZeroOk", "rank3ByRank2ShapeOk", "dim(shapes.a, 2) >= 32", "numel(shapes.y) >= 8192", "dim(shapes.a, 2) % 4 != 0", "rank3Dp4aDispatchOk"], |
| "derive": { |
| "aPacked": false, |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "tileM": 8, |
| "tileN": 64, |
| "ktWords": 32 |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aBatched3x2Tail", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": [ |
| "a_main", |
| "b", |
| "a_zero_point", |
| "b_zero_point", |
| "y", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "M", "type": "u32", "value": "batchM" }, |
| { "name": "N", "type": "u32", "value": "batchN" }, |
| { "name": "K", "type": "u32", "value": "batchK" }, |
| { "name": "aBatchStride4", "type": "u32", "value": "batchAStride" }, |
| { "name": "bBatchStride", "type": "u32", "value": "batchBStride" }, |
| { "name": "yBatchStride", "type": "u32", "value": "batchYStride" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "ceil(dim(shapes.y, 2) / 64)", "y": "ceil(dim(shapes.y, 1) / 8)", "z": "dim(shapes.y, 0)" } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_scalar_zero_points", |
| "priority": 20, |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "requiredScalarAZeroOk and requiredScalarBZeroOk", "dim(shapes.a, 1) >= 64", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 == 0", "rank2Dp4aDispatchOk"], |
| "derive": { |
| "aStorage": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4a", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN)", |
| "y": "ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_a_zero_point_only", |
| "priority": 20, |
| "when": ["present.a_zero_point", "not present.b_zero_point", "rank2MatMulShapeOk", "requiredScalarAZeroOk", "dim(shapes.a, 1) >= 64", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 == 0", "rank2Dp4aDispatchOk"], |
| "derive": { |
| "aVec4": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": true, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4a", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a_a_vec4", "b", "a_zero_point", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN)", |
| "y": "ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_b_zero_point_only", |
| "priority": 20, |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "requiredScalarBZeroOk", "dim(shapes.a, 1) >= 64", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 == 0", "rank2Dp4aDispatchOk"], |
| "derive": { |
| "aStorage": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4a", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a", "b", "b_zero_point", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN)", |
| "y": "ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_no_zero_points", |
| "priority": 20, |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "rank2MatMulShapeOk", "true", "dim(shapes.a, 1) >= 64", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 == 0", "rank2Dp4aDispatchOk"], |
| "derive": { |
| "aVec4": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4a", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a_a_vec4", "b", "y", "params"], |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.y, 1), rank2Dp4aTileN)", |
| "y": "ceilDiv(dim(shapes.y, 0), rank2Dp4aTileM)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_b_zero_point_per_column", |
| "priority": 20, |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "rank2PerColumnBZeroOk", "dim(shapes.a, 1) >= 64", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 == 0", "rank2Dp4aDispatchOk"], |
| "requires": { |
| "limits": { |
| "maxComputeInvocationsPerWorkgroup": 128, |
| "maxComputeWorkgroupSizeX": 128, |
| "maxComputeWorkgroupStorageSize": 8448 |
| } |
| }, |
| "tunables": { "TILE_M": { "default": 32 }, "TILE_N": { "default": 128 }, "KT_WORDS": { "default": 32 } }, |
| "derive": { |
| "aStorage": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": false, |
| "hasBZero": true, |
| "bZeroPerColumn": true, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL", |
| "tileM": "rank2Dp4aTileM if dim(shapes.y, 0) == 1 else (16 if perColumnDp4aUseNarrowTile else tunables.TILE_M)", |
| "tileN": "rank2Dp4aTileN if dim(shapes.y, 0) == 1 else (2 * tunables.TILE_N if variableSubgroup16To32 and rank2Dp4aLargePrefill and 2 * tunables.TILE_N <= deviceWorkgroupCap else tunables.TILE_N)", |
| "ktWords": "rank2Dp4aKtWords if dim(shapes.y, 0) == 1 else tunables.KT_WORDS" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aBPerColumn", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a", "b", "b_zero_point", "y", "params"], |
| "dispatch": { "x": "ceilDiv(dim(shapes.y, 1), tileN)", "y": "ceilDiv(dim(shapes.y, 0), tileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank2_a_zero_point_b_per_column", |
| "priority": 20, |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "requiredScalarAZeroOk and rank2PerColumnBZeroOk", "dim(shapes.a, 1) >= 64", "numel(shapes.y) >= 8192 or (dim(shapes.y, 0) == 1 and dim(shapes.y, 1) >= 512)", "dim(shapes.a, 1) % 4 == 0", "rank2Dp4aDispatchOk"], |
| "requires": { |
| "limits": { |
| "maxComputeInvocationsPerWorkgroup": 128, |
| "maxComputeWorkgroupSizeX": 128, |
| "maxComputeWorkgroupStorageSize": 8448 |
| } |
| }, |
| "tunables": { "TILE_M": { "default": 32 }, "TILE_N": { "default": 128 }, "KT_WORDS": { "default": 32 } }, |
| "derive": { |
| "aStorage": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": true, |
| "gemvKUnroll": "tunables.DP4A_GEMV_K_UNROLL", |
| "tileM": "rank2Dp4aTileM if dim(shapes.y, 0) == 1 else (16 if perColumnDp4aUseNarrowTile else tunables.TILE_M)", |
| "tileN": "rank2Dp4aTileN if dim(shapes.y, 0) == 1 else (2 * tunables.TILE_N if variableSubgroup16To32 and rank2Dp4aLargePrefill and 2 * tunables.TILE_N <= deviceWorkgroupCap else tunables.TILE_N)", |
| "ktWords": "rank2Dp4aKtWords if dim(shapes.y, 0) == 1 else tunables.KT_WORDS" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aAScalarBPerColumn", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a", "b", "a_zero_point", "b_zero_point", "y", "params"], |
| "dispatch": { "x": "ceilDiv(dim(shapes.y, 1), tileN)", "y": "ceilDiv(dim(shapes.y, 0), tileM)" } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_scalar_zero_points", |
| "when": ["present.a_zero_point", "present.b_zero_point", "((rank2MatMulShapeOk and requiredScalarAZeroOk and requiredScalarBZeroOk) or (rank1DotShapeOk and scalarAZeroOk and scalarBZeroOk))", "true"], |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params_main"], |
| "dispatch": { |
| "x": "ceil(matrixN / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceilDiv((matrixM), (tunables.WORKGROUP_SIZE_Y)), DISPATCH_FOLD_WIDTH)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_a_zero_point_only", |
| "when": ["present.a_zero_point", "not present.b_zero_point", "((rank2MatMulShapeOk and requiredScalarAZeroOk) or (rank1DotShapeOk and scalarAZeroOk))", "true"], |
| "derive": { "hasAZero": true, "hasBZero": false, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "y", "params_main"], |
| "dispatch": { |
| "x": "ceil(matrixN / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceilDiv((matrixM), (tunables.WORKGROUP_SIZE_Y)), DISPATCH_FOLD_WIDTH)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_b_zero_point_only", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "requiredScalarBZeroOk"], |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params_main"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceilDiv((dim(shapes.y, 0)), (tunables.WORKGROUP_SIZE_Y)), DISPATCH_FOLD_WIDTH)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank1_dot_b_zero_point_only", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank1DotShapeOk", "scalarBZeroOk"], |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params_m_n_k"], |
| "dispatch": { "x": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "rank1_dot_no_zero_points", |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "rank1DotShapeOk", "true"], |
| "derive": { "hasAZero": false, "hasBZero": false, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "y", "params_m_n_k"], |
| "dispatch": { "x": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_no_zero_points", |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "rank2MatMulShapeOk", "true"], |
| "derive": { "hasAZero": false, "hasBZero": false, "bZeroPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "y", "params_main"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceilDiv((dim(shapes.y, 0)), (tunables.WORKGROUP_SIZE_Y)), DISPATCH_FOLD_WIDTH)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_b_zero_point_per_column", |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "scalarAZeroOk and rank2PerColumnBZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroPerColumn": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.PerColumn", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params_main"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceilDiv((dim(shapes.y, 0)), (tunables.WORKGROUP_SIZE_Y)), DISPATCH_FOLD_WIDTH)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_b_zero_point_per_column_no_azp", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank2MatMulShapeOk", "rank2PerColumnBZeroOk"], |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroPerColumn": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.PerColumnNoAZero", |
| "shader": "quant-matmul-accumulate-rank2.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params_main"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceilDiv((dim(shapes.y, 0)), (tunables.WORKGROUP_SIZE_Y)), DISPATCH_FOLD_WIDTH)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_by_rank3_scalar_zero_points", |
| "when": ["present.a_zero_point", "present.b_zero_point", "(rank3ByRank3ShapeOk or rank3ByRank2ShapeOk)", "scalarAZeroOk and scalarBZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched3x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_by_rank3_no_zero_points", |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "(rank3ByRank3ShapeOk or rank3ByRank2ShapeOk)", "true"], |
| "derive": { "hasAZero": false, "hasBZero": false, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched3x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_by_rank3_a_zero_point_only", |
| "when": ["present.a_zero_point", "not present.b_zero_point", "(rank3ByRank3ShapeOk or rank3ByRank2ShapeOk)", "scalarAZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": false, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched3x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_by_rank3_b_zero_point_only", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "(rank3ByRank3ShapeOk or rank3ByRank2ShapeOk)", "scalarBZeroOk"], |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched3x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_by_rank3_b_zero_point_per_column_batched", |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank3ByRank3ShapeOk", "scalarAZeroOk and ranks.b_zero_point == 3 and dim(shapes.b_zero_point, 0) == dim(shapes.y, 0) and dim(shapes.b_zero_point, 1) == 1 and dim(shapes.b_zero_point, 2) == dim(shapes.y, 2)"], |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroBatchedPerColumn": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched3x3PerColumn", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_by_rank3_b_zero_point_per_column_batched_no_azp", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank3ByRank3ShapeOk", "ranks.b_zero_point == 3 and dim(shapes.b_zero_point, 0) == dim(shapes.y, 0) and dim(shapes.b_zero_point, 1) == 1 and dim(shapes.b_zero_point, 2) == dim(shapes.y, 2)"], |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroBatchedPerColumn": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched3x3PerColumn", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_by_rank3_scalar_zero_points", |
| "when": ["present.a_zero_point", "present.b_zero_point", "ranks.a == 2 and ranks.b == 3 and ranks.y == 3 and dim(shapes.a, 1) == dim(shapes.b, 1) and dim(shapes.y, 0) == dim(shapes.b, 0) and dim(shapes.y, 1) == dim(shapes.a, 0) and dim(shapes.y, 2) == dim(shapes.b, 2)", "scalarAZeroOk and scalarBZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": true, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched2x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_by_rank3_a_zero_point_only", |
| "when": ["present.a_zero_point", "not present.b_zero_point", "ranks.a == 2 and ranks.b == 3 and ranks.y == 3 and dim(shapes.a, 1) == dim(shapes.b, 1) and dim(shapes.y, 0) == dim(shapes.b, 0) and dim(shapes.y, 1) == dim(shapes.a, 0) and dim(shapes.y, 2) == dim(shapes.b, 2)", "scalarAZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": false, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched2x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "a_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_by_rank3_b_zero_point_only", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "ranks.a == 2 and ranks.b == 3 and ranks.y == 3 and dim(shapes.a, 1) == dim(shapes.b, 1) and dim(shapes.y, 0) == dim(shapes.b, 0) and dim(shapes.y, 1) == dim(shapes.a, 0) and dim(shapes.y, 2) == dim(shapes.b, 2)", "scalarBZeroOk"], |
| "derive": { "hasAZero": false, "hasBZero": true, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched2x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank2_by_rank3_no_zero_points", |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "ranks.a == 2 and ranks.b == 3 and ranks.y == 3 and dim(shapes.a, 1) == dim(shapes.b, 1) and dim(shapes.y, 0) == dim(shapes.b, 0) and dim(shapes.y, 1) == dim(shapes.a, 0) and dim(shapes.y, 2) == dim(shapes.b, 2)", "true"], |
| "derive": { "hasAZero": false, "hasBZero": false, "bZeroBatchedPerColumn": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched2x3", |
| "shader": "matmul-integer-batched.wgsl.jinja", |
| "bindings": ["a_main", "b", "y", "params__uniform"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 1) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank4_scalar_zero_points", |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank4BroadcastShapeOk", "scalarAZeroOk and scalarBZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched4", |
| "shader": "quant-matmul-accumulate-rank4.wgsl.jinja", |
| "derive": { "outputToScratch": false }, |
| "bindings": ["a_main", "b", "a_zero_point", "b_zero_point", "y", "params_rank4_no_zero_points"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 3) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0) * dim(shapes.y, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank4_a_zero_point_only", |
| "when": ["present.a_zero_point", "not present.b_zero_point", "rank4BroadcastShapeOk", "scalarAZeroOk"], |
| "derive": { "hasAZero": true, "hasBZero": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched4", |
| "shader": "quant-matmul-accumulate-rank4.wgsl.jinja", |
| "derive": { "outputToScratch": false }, |
| "bindings": ["a_main", "b", "a_zero_point", "y", "params_rank4_no_zero_points"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 3) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0) * dim(shapes.y, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank4_b_zero_point_only", |
| "when": ["not present.a_zero_point", "present.b_zero_point", "rank4BroadcastShapeOk", "scalarBZeroOk"], |
| "derive": { "hasAZero": false, "hasBZero": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched4", |
| "shader": "quant-matmul-accumulate-rank4.wgsl.jinja", |
| "derive": { "outputToScratch": false }, |
| "bindings": ["a_main", "b", "b_zero_point", "y", "params_rank4_no_zero_points"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 3) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0) * dim(shapes.y, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank4_no_zero_points", |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "rank4BroadcastShapeOk", "true"], |
| "derive": { "hasAZero": false, "hasBZero": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Batched4", |
| "shader": "quant-matmul-accumulate-rank4.wgsl.jinja", |
| "derive": { "outputToScratch": false }, |
| "bindings": ["a_main", "b", "y", "params_rank4_no_zero_points"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 3) / tunables.WORKGROUP_SIZE_X)", |
| "y": "min(ceil(dim(shapes.y, 2) / tunables.WORKGROUP_SIZE_Y), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", |
| "z": "min(dim(shapes.y, 0) * dim(shapes.y, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank3_by_rank3_scalar_zero_points", |
| "priority": 20, |
| "when": ["present.a_zero_point", "present.b_zero_point", "(rank3ByRank3ShapeOk or rank3ByRank2ShapeOk)", "scalarAZeroOk and scalarBZeroOk", "dim(shapes.a, 2) >= 64", "numel(shapes.y) >= 8192", "dim(shapes.a, 2) % 4 == 0", "rank3Dp4aDispatchOk"], |
| "derive": { |
| "aVec4": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "tileM": 8, |
| "tileN": 64, |
| "ktWords": 32 |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aBatched3x3", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a_a_vec4", "b", "a_zero_point", "b_zero_point", "y", "params_quant_dp4a_matmul"], |
| "dispatch": { "x": "ceil(dim(shapes.y, 2) / 64)", "y": "ceil(dim(shapes.y, 1) / 8)", "z": "dim(shapes.y, 0)" } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank4_scalar_zero_points", |
| "priority": 20, |
| "when": ["present.a_zero_point", "present.b_zero_point", "rank4Dp4aShapeOk", "scalarAZeroOk and scalarBZeroOk", "dim(shapes.a, 3) >= 128", "numel(shapes.y) >= 8192", "dim(shapes.a, 3) % 4 == 0", "rank4Dp4aDispatchOk"], |
| "derive": { |
| "aVec4": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": true, |
| "hasBZero": true, |
| "bZeroPerColumn": false, |
| "tileM": 8, |
| "tileN": 64, |
| "ktWords": 32 |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aBatched4Scalar", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a_a_vec4", "b", "a_zero_point", "b_zero_point", "y", "params_dp4a_rank4_no_zero_points"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 3) / 64)", |
| "y": "ceil(dim(shapes.y, 2) / 8)", |
| "z": "dim(shapes.y, 0) * dim(shapes.y, 1)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "dp4a_rank4_no_zero_points", |
| "priority": 20, |
| "when": ["not present.a_zero_point", "not present.b_zero_point", "rank4Dp4aShapeOk", "true", "dim(shapes.a, 3) >= 128", "numel(shapes.y) >= 8192", "dim(shapes.a, 3) % 4 == 0", "rank4Dp4aDispatchOk"], |
| "derive": { |
| "aVec4": "\"vec4<u32>\" if tensorDtypes.a == \"uint8\" else \"vec4<i32>\"", |
| "hasAZero": false, |
| "hasBZero": false, |
| "bZeroPerColumn": false, |
| "tileM": 8, |
| "tileN": 64, |
| "ktWords": 32 |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "MatMulInteger.Dp4aBatched4NoZero", |
| "shader": "quant-dp4a-matmul.wgsl.jinja", |
| "bindings": ["a_a_vec4", "b", "y", "params_dp4a_rank4_no_zero_points"], |
| "dispatch": { |
| "x": "ceil(dim(shapes.y, 3) / 64)", |
| "y": "ceil(dim(shapes.y, 2) / 8)", |
| "z": "dim(shapes.y, 0) * dim(shapes.y, 1)" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
