| #!/usr/bin/env python3 | |
| # Copyright 2025 The IREE Authors | |
| # | |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| """compilation_info support in e2e matmul tests""" | |
| from typing import Optional | |
| import enum | |
| import dataclasses | |
| import typing | |
| from tests.e2e.matmul.common import * | |
| # Describes a workgroup and tiling schedule to target a specific MMA intrinsic. | |
| @dataclasses.dataclass | |
| class MMASchedule: | |
| intrinsic: str | |
| m_count: int # Number of subgroups per workgroup along M | |
| n_count: int # Number of subgroups per workgroup along N | |
| m_tile_count: int | |
| n_tile_count: int | |
| k_tile_count: int | |
| batch_tile_count: int = ( | |
| 1 # Number of block-intrinsic tiles per subgroup along batch | |
| ) | |
| def get_subgroup_basis(self) -> str: | |
| return f"[[{self.m_count}, {self.n_count}, 1], [0, 1, 2]]" | |
| def get_batch_subgroup_basis(self) -> str: | |
| return f"[[{self.batch_tile_count}, {self.m_count}, {self.n_count}, 1], [0, 1, 2, 3]]" | |
| def get_subgroup_tile(self) -> str: | |
| """Returns subgroup tile sizes for TileAndFuse pipeline.""" | |
| return f"[{self.m_count}, {self.n_count}, 0]" | |
| def get_batch_subgroup_tile(self) -> str: | |
| """Returns subgroup tile sizes for TileAndFuse pipeline for batch matmul.""" | |
| return f"[{self.batch_tile_count}, {self.m_count}, {self.n_count}, 0]" | |
| # Enumerates of the collections of compilation info that we can generate tests | |
| # for. The values are the accepted values for the --compilation_info= flag. | |
| @enum.unique | |
| class CompilationInfoId(enum.Enum): | |
| NONE = "" | |
| LLVMGPUVectorDistributeMFMA = "LLVMGPUVectorDistributeMFMA" | |
| LLVMGPUVectorDistributeWMMAR3 = "LLVMGPUVectorDistributeWMMAR3" | |
| LLVMGPUVectorDistributeWMMAR4 = "LLVMGPUVectorDistributeWMMAR4" | |
| LLVMGPUVectorDistributeWMMA1250 = "LLVMGPUVectorDistributeWMMA1250" | |
| LLVMGPUTileAndFuseMFMA = "LLVMGPUTileAndFuseMFMA" | |
| LLVMGPUTileAndFuseWMMAR3 = "LLVMGPUTileAndFuseWMMAR3" | |
| LLVMGPUTileAndFuseWMMAR4 = "LLVMGPUTileAndFuseWMMAR4" | |
| LLVMGPUTileAndFuseWMMA1250 = "LLVMGPUTileAndFuseWMMA1250" | |
| LLVMGPUVectorDistributeMFMABlockBatch = "LLVMGPUVectorDistributeMFMABlockBatch" | |
| LLVMGPUTileAndFuseMFMABlockBatch = "LLVMGPUTileAndFuseMFMABlockBatch" | |
| LLVMGPUTileAndFuseVDMFMA_CDNA3 = "LLVMGPUTileAndFuseVDMFMA_CDNA3" | |
| LLVMGPUVectorDistributeCUDA = "LLVMGPUVectorDistributeCUDA" | |
| LLVMGPUTileAndFuseCUDA = "LLVMGPUTileAndFuseCUDA" | |
| SPIRVCooperativeMatrixVectorize = "SPIRVCooperativeMatrixVectorize" | |
| SPIRVVectorizeMali = "SPIRVVectorizeMali" | |
| SPIRVVectorizeNVIDIA = "SPIRVVectorizeNVIDIA" | |
| # Describes how to construct compilation info for the testcase. | |
| @dataclasses.dataclass | |
| class CompilationInfo: | |
| # Compilation info | |
| workgroup_size: typing.List[int] | |
| subgroup_size: Optional[int] | |
| # Translation info | |
| dispatch_lowering_pass_pipeline: str | |
| # Prints the workgroup size | |
| def workgroup_size_str(self): | |
| return "workgroup_size = [" + ", ".join(map(str, self.workgroup_size)) + "]" | |
| def get_compilation_info_attr(self) -> str: ... | |
| @dataclasses.dataclass | |
| class IREEGPUCompilationInfo(CompilationInfo): | |
| # Lowering Config | |
| workgroup_tile: list[int] | |
| reduction_tile: list[int] | |
| # Translation Info | |
| mma_schedule: Optional[MMASchedule] | |
| is_batch: bool = False # True for batch matmul with block intrinsics | |
| def get_compilation_info_attr(self) -> str: | |
| requested_pipeline = self.dispatch_lowering_pass_pipeline | |
| compiler_pipeline = requested_pipeline | |
| subgroup_size_str = "" | |
| if self.subgroup_size is not None: | |
| subgroup_size_str = f"subgroup_size = {self.subgroup_size}" | |
| # Virtual intrinsics (VMFMA, VDMFMA) use virtual_mma_layout; | |
| # physical intrinsics (MFMA, WMMA, NV_MMA_SYNC) use mma_layout. | |
| is_virtual = self.mma_schedule.intrinsic.startswith("V") | |
| mma_attr = "virtual_mma_layout" if is_virtual else "mma_layout" | |
| if compiler_pipeline == "#iree_gpu.pipeline<TileAndFuse>": | |
| # Add convert_acc_gemm for NVIDIA mma.sync intrinsics | |
| convert_acc_gemm = "" | |
| if self.mma_schedule.intrinsic.startswith("NV_MMA_SYNC"): | |
| convert_acc_gemm = "convert_acc_gemm, " | |
| subgroup_tile = ( | |
| self.mma_schedule.get_batch_subgroup_tile() | |
| if self.is_batch | |
| else self.mma_schedule.get_subgroup_tile() | |
| ) | |
| lowering_config = ( | |
| f" lowering_config = #iree_gpu.lowering_config<{{" | |
| f" mma_kind = #iree_gpu.{mma_attr}<{self.mma_schedule.intrinsic}>, " | |
| f" subgroup = {subgroup_tile}, " | |
| f" {convert_acc_gemm}" | |
| f" promote_operands = [0, 1], " | |
| f" workgroup = {self.workgroup_tile}, " | |
| f" reduction = {self.reduction_tile} }}>,\n" | |
| ) | |
| else: | |
| subgroup_basis = ( | |
| self.mma_schedule.get_batch_subgroup_basis() | |
| if self.is_batch | |
| else self.mma_schedule.get_subgroup_basis() | |
| ) | |
| lowering_config = ( | |
| f" lowering_config = #iree_gpu.lowering_config<{{" | |
| f" mma_kind = #iree_gpu.{mma_attr}<{self.mma_schedule.intrinsic}>, " | |
| f" subgroup_basis = {subgroup_basis}, " | |
| f" workgroup = {self.workgroup_tile}, " | |
| f" reduction = {self.reduction_tile} }}>,\n" | |
| ) | |
| return ( | |
| "#iree_codegen.compilation_info<\n" | |
| f"{lowering_config}" | |
| f" translation_info = #iree_codegen.translation_info<pipeline = {compiler_pipeline} {self.workgroup_size_str()}\n" | |
| f" {subgroup_size_str}>>\n" | |
| ) | |
| @dataclasses.dataclass | |
| class LegacyCompilationInfo(CompilationInfo): | |
| # Lowering Config | |
| tile_sizes: typing.List[typing.List[int]] | |
| # Translation Info | |
| software_pipeline_depth: int | |
| def get_compilation_info_attr(self) -> str: | |
| requested_pipeline = self.dispatch_lowering_pass_pipeline | |
| compiler_pipeline = requested_pipeline | |
| if requested_pipeline == "SPIRVVectorizeMali": | |
| compiler_pipeline = "#iree_gpu.spirv_pipeline<BaseVectorize>" | |
| elif requested_pipeline == "SPIRVCooperativeMatrixVectorize": | |
| compiler_pipeline = "#iree_gpu.spirv_pipeline<CooperativeMatrixVectorize>" | |
| elif requested_pipeline == "SPIRVVectorizeNVIDIA": | |
| # TODO: change to test SPIRVMatmulPromoteVectorize too | |
| compiler_pipeline = "#iree_gpu.spirv_pipeline<BaseVectorize>" | |
| subgroup_size_str = "" | |
| if self.subgroup_size is not None: | |
| subgroup_size_str = f"subgroup_size = {self.subgroup_size}" | |
| return ( | |
| "#iree_codegen.compilation_info<\n" | |
| f" lowering_config = #iree_codegen.lowering_config<tile_sizes = {self.tile_sizes}>,\n" | |
| f" translation_info = #iree_codegen.translation_info<pipeline = {compiler_pipeline} {self.workgroup_size_str()}\n" | |
| f" {subgroup_size_str},\n" | |
| f" {{ pipeline_depth = {self.software_pipeline_depth}, store_stage = 1}}>>" | |
| ) | |
| def generate_compilation_info_string_and_attr(compilation_info: CompilationInfo): | |
| if not compilation_info: | |
| return ("", "") | |
| index = generate_compilation_info_string_and_attr.compilation_index | |
| generate_compilation_info_string_and_attr.compilation_index += 1 | |
| return ( | |
| f"#compilation{index} = {compilation_info.get_compilation_info_attr()}", | |
| f"{{compilation_info = #compilation{index}}} ", | |
| ) | |
| generate_compilation_info_string_and_attr.compilation_index = 0 | |
| @dataclasses.dataclass | |
| class TileWorkgroupSizePair: | |
| tile_size: typing.List[typing.List[int]] | |
| workgroup_size: typing.List[int] | |
| # Constructs a TileWorkgroupSizePair for SPIR-V targets enforcing the | |
| # constraints between the workgroup_size and tile size | |
| def get_spirv_tile_workgroup_size_pair( | |
| workgroup_size, t_tile_k, t_tile_m=4, t_tile_n=4 | |
| ): | |
| x, y, z = workgroup_size | |
| wg_tile_m = y * t_tile_m | |
| wg_tile_n = x * t_tile_n | |
| return TileWorkgroupSizePair( | |
| [[wg_tile_m, wg_tile_n], [t_tile_m, t_tile_n], [0, 0, t_tile_k]], workgroup_size | |
| ) | |
| # Returns all the TileWorkgroupSizePairs for a given SPIRV Target | |
| def get_all_spirv_tile_workgroup_size_pairs(t_tile_k): | |
| tile_workgroup_size_pairs = [ | |
| get_spirv_tile_workgroup_size_pair([32, 8, 1], t_tile_k), | |
| get_spirv_tile_workgroup_size_pair([16, 8, 1], t_tile_k), | |
| get_spirv_tile_workgroup_size_pair([64, 2, 1], t_tile_k), | |
| get_spirv_tile_workgroup_size_pair([8, 8, 1], t_tile_k), | |
| get_spirv_tile_workgroup_size_pair([32, 1, 1], t_tile_k), | |
| get_spirv_tile_workgroup_size_pair([16, 2, 1], t_tile_k), | |
| get_spirv_tile_workgroup_size_pair([32, 1, 1], t_tile_k), | |
| ] | |
| return tile_workgroup_size_pairs | |
| def get_rocm_test_compilation_infos( | |
| compilation_info_id: CompilationInfoId, lhs_rhs_type: MatrixElemTypeId | |
| ): | |
| vecdist = "#iree_gpu.pipeline<VectorDistribute>" | |
| tileandfuse = "#iree_gpu.pipeline<TileAndFuse>" | |
| id_to_intrinsic_and_pipeline = { | |
| CompilationInfoId.LLVMGPUVectorDistributeMFMA: ("MFMA", vecdist), | |
| CompilationInfoId.LLVMGPUVectorDistributeWMMAR3: ("WMMAR3", vecdist), | |
| CompilationInfoId.LLVMGPUVectorDistributeWMMAR4: ("WMMAR4", vecdist), | |
| CompilationInfoId.LLVMGPUVectorDistributeWMMA1250: ("WMMA1250", vecdist), | |
| CompilationInfoId.LLVMGPUTileAndFuseMFMA: ("MFMA", tileandfuse), | |
| CompilationInfoId.LLVMGPUTileAndFuseVDMFMA_CDNA3: ("VDMFMA_CDNA3", tileandfuse), | |
| CompilationInfoId.LLVMGPUTileAndFuseWMMAR3: ("WMMAR3", tileandfuse), | |
| CompilationInfoId.LLVMGPUTileAndFuseWMMAR4: ("WMMAR4", tileandfuse), | |
| CompilationInfoId.LLVMGPUTileAndFuseWMMA1250: ("WMMA1250", tileandfuse), | |
| } | |
| if compilation_info_id not in id_to_intrinsic_and_pipeline: | |
| raise ValueError("Unknown pipeline for rocm") | |
| intrinsic, pipeline = id_to_intrinsic_and_pipeline[compilation_info_id] | |
| use_tile_and_fuse = pipeline == "#iree_gpu.pipeline<TileAndFuse>" | |
| schedules = [] | |
| if intrinsic == "MFMA": | |
| schedules = [ | |
| MMASchedule("MFMA_F32_16x16x4_F32", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x4_F32", 1, 1, 1, 1, 2), | |
| MMASchedule("MFMA_F32_16x16x4_F32", 1, 1, 1, 2, 1), | |
| MMASchedule("MFMA_F32_16x16x4_F32", 1, 1, 2, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x4_F32", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 1, 1, 1, 1, 2), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 1, 1, 1, 2, 1), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 1, 1, 2, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 2, 2, 1, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 2, 4, 2, 1, 2), | |
| MMASchedule("MFMA_F32_16x16x16_F16", 4, 2, 4, 2, 2), | |
| MMASchedule("MFMA_F32_32x32x8_F16", 1, 1, 1, 2, 2), | |
| MMASchedule("MFMA_F32_32x32x8_F16", 2, 2, 1, 1, 1), | |
| MMASchedule("MFMA_F32_32x32x8_F16", 1, 4, 2, 1, 2), | |
| MMASchedule("MFMA_F32_32x32x8_F16", 4, 2, 1, 2, 4), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FNUZ", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FNUZ", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FNUZ", 4, 1, 4, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FNUZ", 4, 2, 4, 2, 1), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FNUZ", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FNUZ", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FNUZ", 4, 1, 1, 2, 2), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FNUZ", 4, 2, 2, 2, 2), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FN", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FN", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FN", 4, 1, 4, 1, 1), | |
| MMASchedule("MFMA_F32_16x16x32_F8E4M3FN", 4, 2, 4, 2, 1), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FN", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FN", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FN", 4, 1, 1, 2, 2), | |
| MMASchedule("MFMA_F32_32x32x16_F8E4M3FN", 4, 2, 2, 2, 2), | |
| MMASchedule("MFMA_I32_16x16x32_I8", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_I32_16x16x32_I8", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_I32_16x16x32_I8", 4, 1, 4, 1, 1), | |
| MMASchedule("MFMA_I32_16x16x32_I8", 4, 2, 4, 2, 1), | |
| MMASchedule("MFMA_I32_32x32x16_I8", 1, 1, 1, 1, 1), | |
| MMASchedule("MFMA_I32_32x32x16_I8", 2, 2, 1, 1, 2), | |
| MMASchedule("MFMA_I32_32x32x16_I8", 4, 1, 1, 2, 2), | |
| MMASchedule("MFMA_I32_32x32x16_I8", 4, 2, 2, 2, 2), | |
| MMASchedule("VMFMA_F32_16x16x32_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("VMFMA_F32_16x16x32_F16", 4, 2, 1, 2, 4), | |
| MMASchedule("VMFMA_F32_32x32x16_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("VMFMA_F32_32x32x16_F16", 4, 2, 1, 2, 4), | |
| MMASchedule("VMFMA_F32_16x16x32_F8E4M3FNUZ", 1, 1, 1, 1, 1), | |
| MMASchedule("VMFMA_F32_16x16x32_F8E4M3FNUZ", 4, 1, 4, 1, 1), | |
| ] | |
| elif intrinsic == "VDMFMA_CDNA3": | |
| # CDNA3 (gfx94x) VDMFMA: virtual dense MFMAs for skinny matmuls. | |
| # F16/BF16: M=8, N=16, K=64 (x2 underlying SMFMACs) | |
| schedules = [ | |
| MMASchedule("VDMFMA_F32_8x16x64x2_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("VDMFMA_F32_8x16x64x2_F16", 1, 2, 1, 2, 2), | |
| MMASchedule("VDMFMA_F32_8x16x64x2_BF16", 1, 1, 1, 1, 1), | |
| MMASchedule("VDMFMA_F32_8x16x64x2_BF16", 1, 2, 1, 2, 2), | |
| ] | |
| elif intrinsic == "WMMAR3": | |
| schedules = [ | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMAR3_F32_16x16x16_F16", 4, 2, 4, 2, 2), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMAR3_I32_16x16x16_I8", 4, 2, 4, 2, 2), | |
| ] | |
| elif intrinsic == "WMMAR4": | |
| schedules = [ | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMAR4_F32_16x16x16_F16", 4, 2, 4, 2, 2), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMAR4_F32_16x16x16_F8E4M3FN", 4, 2, 4, 2, 2), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMAR4_I32_16x16x16_I8", 4, 2, 4, 2, 2), | |
| ] | |
| elif intrinsic == "WMMA1250": | |
| # gfx1250 WMMA intrinsics: 16x16 tiles with various K sizes. | |
| # F16: K=32, F8E4M3FN: K=64 or K=128, I8: K=64 | |
| schedules = [ | |
| MMASchedule("WMMA_F32_16x16x32_F16", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x32_F16", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMA_F32_16x16x32_F16", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMA_F32_16x16x32_F16", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x32_F16", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x32_F16", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMA_F32_16x16x32_F16", 4, 2, 4, 2, 2), | |
| # K=64 F8E4M3FN | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMA_F32_16x16x64_F8E4M3FN", 4, 2, 4, 2, 2), | |
| # K=128 F8E4M3FN | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMA_F32_16x16x128_F8E4M3FN", 4, 2, 4, 2, 2), | |
| # I8 | |
| MMASchedule("WMMA_I32_16x16x64_I8", 1, 1, 1, 1, 1), | |
| MMASchedule("WMMA_I32_16x16x64_I8", 1, 1, 1, 1, 2), | |
| MMASchedule("WMMA_I32_16x16x64_I8", 1, 1, 1, 2, 1), | |
| MMASchedule("WMMA_I32_16x16x64_I8", 1, 1, 2, 1, 1), | |
| MMASchedule("WMMA_I32_16x16x64_I8", 2, 2, 1, 1, 1), | |
| MMASchedule("WMMA_I32_16x16x64_I8", 2, 4, 2, 1, 2), | |
| MMASchedule("WMMA_I32_16x16x64_I8", 4, 2, 4, 2, 2), | |
| ] | |
| else: | |
| raise NotImplementedError("unhandled intrinsic case") | |
| # MFMA and VDMFMA use 64-lane subgroups (CDNA). | |
| # WMMAR3, WMMAR4, WMMA1250 all use 32. | |
| subgroup_size = 64 if intrinsic in ("MFMA", "VDMFMA_CDNA3") else 32 | |
| infos = [] | |
| for schedule in schedules: | |
| # Skip schedules with an intrinsic which element type does not | |
| # match the requested one. | |
| # Extracts the input type from strings. The naming convention is | |
| # [output_type]_MxNxK_[input_type]. | |
| input_type = schedule.intrinsic.split("_")[-1] | |
| if lhs_rhs_type.value.upper() != input_type: | |
| continue | |
| if schedule.intrinsic == "MFMA_F32_16x16x4_F32": | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 4 | |
| elif schedule.intrinsic == "MFMA_F32_16x16x16_F16": | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 16 | |
| elif schedule.intrinsic == "MFMA_F32_32x32x8_F16": | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 32 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 32 | |
| wg_tile_k = schedule.k_tile_count * 8 | |
| elif ( | |
| schedule.intrinsic == "VMFMA_F32_16x16x32_F16" | |
| or schedule.intrinsic == "MFMA_I32_16x16x32_I8" | |
| or schedule.intrinsic == "MFMA_F32_16x16x32_F8E4M3FNUZ" | |
| or schedule.intrinsic == "VMFMA_F32_16x16x32_F8E4M3FNUZ" | |
| or schedule.intrinsic == "MFMA_F32_16x16x32_F8E4M3FN" | |
| or schedule.intrinsic == "VMFMA_F32_16x16x32_F8E4M3FN" | |
| ): | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 32 | |
| elif schedule.intrinsic in ( | |
| "VDMFMA_F32_8x16x64x2_F16", | |
| "VDMFMA_F32_8x16x64x2_BF16", | |
| ): | |
| # VDMFMA skinny intrinsics: M=8, N=16, K=64 | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 8 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 64 | |
| elif ( | |
| schedule.intrinsic == "VMFMA_F32_32x32x16_F16" | |
| or schedule.intrinsic == "MFMA_F32_32x32x16_F8E4M3FNUZ" | |
| or schedule.intrinsic == "MFMA_F32_32x32x16_F8E4M3FN" | |
| or schedule.intrinsic == "MFMA_I32_32x32x16_I8" | |
| ): | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 32 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 32 | |
| wg_tile_k = schedule.k_tile_count * 16 | |
| elif schedule.intrinsic in ( | |
| "WMMAR3_F32_16x16x16_F16", | |
| "WMMAR3_I32_16x16x16_I8", | |
| "WMMAR4_F32_16x16x16_F16", | |
| "WMMAR4_F32_16x16x16_F8E4M3FN", | |
| "WMMAR4_I32_16x16x16_I8", | |
| ): | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 16 | |
| elif schedule.intrinsic == "WMMA_F32_16x16x32_F16": | |
| # gfx1250: M=16, N=16, K=32 | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 32 | |
| elif schedule.intrinsic in ( | |
| "WMMA_F32_16x16x64_F8E4M3FN", | |
| "WMMA_I32_16x16x64_I8", | |
| ): | |
| # gfx1250: M=16, N=16, K=64 | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 64 | |
| elif schedule.intrinsic == "WMMA_F32_16x16x128_F8E4M3FN": | |
| # gfx1250: M=16, N=16, K=128 | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 16 | |
| wg_tile_k = schedule.k_tile_count * 128 | |
| else: | |
| raise NotImplementedError("unhandled intrinsic case") | |
| workgroup_tile = [wg_tile_m, wg_tile_n, 0] | |
| reduction_k = schedule.k_tile_count if use_tile_and_fuse else wg_tile_k | |
| reduction_tile = [0, 0, reduction_k] | |
| workgroup_size = [schedule.n_count * subgroup_size, schedule.m_count, 1] | |
| infos.append( | |
| IREEGPUCompilationInfo( | |
| workgroup_tile=workgroup_tile, | |
| reduction_tile=reduction_tile, | |
| dispatch_lowering_pass_pipeline=pipeline, | |
| workgroup_size=workgroup_size, | |
| mma_schedule=schedule, | |
| subgroup_size=subgroup_size, | |
| ) | |
| ) | |
| return infos | |
| def get_rocm_block_batch_test_compilation_infos( | |
| compilation_info_id: CompilationInfoId, lhs_rhs_type: MatrixElemTypeId | |
| ): | |
| """Generate compilation infos for CDNA block intrinsic batch matmul tests.""" | |
| vecdist = "#iree_gpu.pipeline<VectorDistribute>" | |
| tileandfuse = "#iree_gpu.pipeline<TileAndFuse>" | |
| if compilation_info_id == CompilationInfoId.LLVMGPUVectorDistributeMFMABlockBatch: | |
| pipeline = vecdist | |
| elif compilation_info_id == CompilationInfoId.LLVMGPUTileAndFuseMFMABlockBatch: | |
| pipeline = tileandfuse | |
| else: | |
| raise ValueError("Unknown block batch pipeline") | |
| use_tile_and_fuse = pipeline == tileandfuse | |
| subgroup_size = 64 | |
| # (intrinsic, block_size, MFMA_M, MFMA_N, MFMA_K) | |
| block_intrinsic_specs = { | |
| "F16": [ | |
| ("MFMA_F32_16x16x4x4B_F16", 4, 16, 16, 4), | |
| ("MFMA_F32_32x32x4x2B_F16", 2, 32, 32, 4), | |
| ("MFMA_F32_4x4x4x16B_F16", 16, 4, 4, 4), | |
| ], | |
| "BF16": [ | |
| ("MFMA_F32_16x16x4x4B_BF16", 4, 16, 16, 4), | |
| ("MFMA_F32_32x32x4x2B_BF16", 2, 32, 32, 4), | |
| ("MFMA_F32_4x4x4x16B_BF16", 16, 4, 4, 4), | |
| ], | |
| "I8": [ | |
| ("MFMA_I32_16x16x4x4B_I8", 4, 16, 16, 4), | |
| ("MFMA_I32_32x32x4x2B_I8", 2, 32, 32, 4), | |
| ("MFMA_I32_4x4x4x16B_I8", 16, 4, 4, 4), | |
| ], | |
| } | |
| input_type = lhs_rhs_type.value.upper() | |
| if input_type not in block_intrinsic_specs: | |
| return [] | |
| schedules_and_specs = [] | |
| for intrinsic, block_size, mm, mn, mk in block_intrinsic_specs[input_type]: | |
| for schedule in [ | |
| MMASchedule(intrinsic, 1, 1, 1, 1, 1, batch_tile_count=1), | |
| MMASchedule(intrinsic, 1, 1, 2, 2, 1, batch_tile_count=1), | |
| MMASchedule(intrinsic, 2, 2, 1, 1, 1, batch_tile_count=1), | |
| ]: | |
| schedules_and_specs.append((schedule, block_size, mm, mn, mk)) | |
| infos = [] | |
| for schedule, block_size, mm, mn, mk in schedules_and_specs: | |
| batch_wg = schedule.batch_tile_count * block_size | |
| m_wg = schedule.m_count * schedule.m_tile_count * mm | |
| n_wg = schedule.n_count * schedule.n_tile_count * mn | |
| k_wg = schedule.k_tile_count * mk | |
| workgroup_tile = [batch_wg, m_wg, n_wg, 0] | |
| reduction_k = schedule.k_tile_count if use_tile_and_fuse else k_wg | |
| reduction_tile = [0, 0, 0, reduction_k] | |
| workgroup_size = [schedule.n_count * schedule.m_count * subgroup_size, 1, 1] | |
| infos.append( | |
| IREEGPUCompilationInfo( | |
| workgroup_tile=workgroup_tile, | |
| reduction_tile=reduction_tile, | |
| dispatch_lowering_pass_pipeline=pipeline, | |
| workgroup_size=workgroup_size, | |
| mma_schedule=schedule, | |
| subgroup_size=subgroup_size, | |
| is_batch=True, | |
| ) | |
| ) | |
| return infos | |
| def get_cuda_test_compilation_infos( | |
| compilation_info_id: CompilationInfoId, | |
| lhs_rhs_type: MatrixElemTypeId, | |
| acc_type: Optional[MatrixElemTypeId] = None, | |
| ): | |
| """Generate compilation infos for CUDA/NVIDIA GPU tests.""" | |
| # Only F16 input is supported for NV_MMA_SYNC intrinsics | |
| if lhs_rhs_type != MatrixElemTypeId.F16: | |
| return [] | |
| # Determine the pipeline based on compilation_info_id | |
| if compilation_info_id == CompilationInfoId.LLVMGPUVectorDistributeCUDA: | |
| pipeline = "#iree_gpu.pipeline<VectorDistribute>" | |
| elif compilation_info_id == CompilationInfoId.LLVMGPUTileAndFuseCUDA: | |
| pipeline = "#iree_gpu.pipeline<TileAndFuse>" | |
| else: | |
| raise ValueError("Unknown pipeline for CUDA") | |
| if acc_type == MatrixElemTypeId.F16: | |
| intrinsic = "NV_MMA_SYNC_F16_16x8x16_F16" | |
| else: | |
| # Default to F32 accumulator | |
| intrinsic = "NV_MMA_SYNC_F32_16x8x16_F16" | |
| schedules = [ | |
| # Basic single subgroup configurations | |
| MMASchedule(intrinsic, 1, 1, 1, 1, 1), | |
| MMASchedule(intrinsic, 1, 1, 1, 1, 2), | |
| MMASchedule(intrinsic, 1, 1, 1, 2, 1), | |
| MMASchedule(intrinsic, 1, 1, 2, 1, 1), | |
| # Multiple subgroups | |
| MMASchedule(intrinsic, 2, 2, 1, 1, 1), | |
| MMASchedule(intrinsic, 2, 2, 2, 2, 2), | |
| MMASchedule(intrinsic, 2, 4, 2, 1, 2), | |
| MMASchedule(intrinsic, 4, 2, 4, 2, 2), | |
| ] | |
| subgroup_size = 32 | |
| infos = [] | |
| for schedule in schedules: | |
| # NV_MMA_SYNC intrinsics: M=16, N=8, K=16 | |
| wg_tile_m = schedule.m_count * schedule.m_tile_count * 16 | |
| wg_tile_n = schedule.n_count * schedule.n_tile_count * 8 | |
| wg_tile_k = schedule.k_tile_count * 16 | |
| workgroup_tile = [wg_tile_m, wg_tile_n, 0] | |
| reduction_tile = [0, 0, wg_tile_k] | |
| workgroup_size = [schedule.n_count * subgroup_size, schedule.m_count, 1] | |
| infos.append( | |
| IREEGPUCompilationInfo( | |
| workgroup_tile=workgroup_tile, | |
| reduction_tile=reduction_tile, | |
| dispatch_lowering_pass_pipeline=pipeline, | |
| workgroup_size=workgroup_size, | |
| mma_schedule=schedule, | |
| subgroup_size=subgroup_size, | |
| ) | |
| ) | |
| return infos | |
| # Returns the list of CompilationInfo's to use for the CompilationInfoId. | |
| def get_test_compilation_infos( | |
| compilation_info_id: CompilationInfoId, | |
| lhs_rhs_type: MatrixElemTypeId, | |
| acc_type: Optional[MatrixElemTypeId] = None, | |
| ) -> typing.List[typing.Optional[CompilationInfo]]: | |
| if compilation_info_id == CompilationInfoId.NONE: | |
| return [None] | |
| if compilation_info_id in [ | |
| CompilationInfoId.LLVMGPUVectorDistributeMFMA, | |
| CompilationInfoId.LLVMGPUVectorDistributeWMMAR3, | |
| CompilationInfoId.LLVMGPUVectorDistributeWMMAR4, | |
| CompilationInfoId.LLVMGPUVectorDistributeWMMA1250, | |
| CompilationInfoId.LLVMGPUTileAndFuseMFMA, | |
| CompilationInfoId.LLVMGPUTileAndFuseVDMFMA_CDNA3, | |
| CompilationInfoId.LLVMGPUTileAndFuseWMMAR3, | |
| CompilationInfoId.LLVMGPUTileAndFuseWMMAR4, | |
| CompilationInfoId.LLVMGPUTileAndFuseWMMA1250, | |
| ]: | |
| return get_rocm_test_compilation_infos(compilation_info_id, lhs_rhs_type) | |
| if compilation_info_id in [ | |
| CompilationInfoId.LLVMGPUVectorDistributeMFMABlockBatch, | |
| CompilationInfoId.LLVMGPUTileAndFuseMFMABlockBatch, | |
| ]: | |
| return get_rocm_block_batch_test_compilation_infos( | |
| compilation_info_id, lhs_rhs_type | |
| ) | |
| if compilation_info_id in [ | |
| CompilationInfoId.LLVMGPUVectorDistributeCUDA, | |
| CompilationInfoId.LLVMGPUTileAndFuseCUDA, | |
| ]: | |
| return get_cuda_test_compilation_infos( | |
| compilation_info_id, lhs_rhs_type, acc_type | |
| ) | |
| software_pipeline_depth = 0 | |
| tile_workgroup_size_pairs = [] | |
| if compilation_info_id == CompilationInfoId.SPIRVCooperativeMatrixVectorize: | |
| tile_workgroup_size_pairs = [ | |
| TileWorkgroupSizePair( | |
| [[64, 128], [32, 64], [0, 0, 32], [16, 16, 16]], [64, 2, 1] | |
| ) | |
| ] | |
| elif compilation_info_id == CompilationInfoId.SPIRVVectorizeNVIDIA: | |
| tile_workgroup_size_pairs = get_all_spirv_tile_workgroup_size_pairs(32) | |
| elif compilation_info_id == CompilationInfoId.SPIRVVectorizeMali: | |
| tile_workgroup_size_pairs = get_all_spirv_tile_workgroup_size_pairs(4) | |
| compilation_infos = [] | |
| for tile_workgroup_size_pair in tile_workgroup_size_pairs: | |
| compilation_infos.append( | |
| LegacyCompilationInfo( | |
| tile_sizes=tile_workgroup_size_pair.tile_size, | |
| dispatch_lowering_pass_pipeline=compilation_info_id.value, | |
| workgroup_size=tile_workgroup_size_pair.workgroup_size, | |
| subgroup_size=None, | |
| software_pipeline_depth=software_pipeline_depth, | |
| ) | |
| ) | |
| return compilation_infos |