| # Copyright 2023 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 | |
| import enum, shutil, functools, operator, collections, subprocess | |
| from library import * | |
| from dispatch import * | |
| from options import get_cmd_line_argument_list | |
| ################################################################################ | |
| class MatmulOperation: | |
| """Data structure to describe a matrix multiplication operation. | |
| This includes the shape, datatype, and layout of the operands. This data | |
| structure is *independent* of the compilation* and tiling configuration. | |
| It "mostly" contains the parameter that changes the functionality of matmul | |
| operation. The only exception is the split_k_slices parameter, which is | |
| changes the performance of the matmul operation and not the functionality. | |
| """ | |
| def __init__( | |
| self, | |
| matmul_shape, | |
| lhs, | |
| rhs, | |
| result, | |
| batch_count=1, | |
| split_k_slices=1, | |
| operation_kind=OperationKind.Matmul, | |
| ): | |
| """Initializes a matrix multiplication operation. | |
| Matrix-multiple operation: `result[M, N] = lhs[M, K] * rhs[K, N]` | |
| matmul_shape: A tuple representing the matrix multiplication problem shape | |
| in the format (M, N, K), where M is the number of rows in the lhs matrix, | |
| N is the number of columns in the rhs matrix, and K is the number of columns | |
| in the lhs matrix and rows in the rhs matrix. | |
| lhs: A TensorDescription object representing the left-hand-side matrix operand. | |
| rhs: A TensorDescription object representing the right-hand-side matrix operand. | |
| result: A TensorDescription object representing the result matrix operand. | |
| """ | |
| # Parameters that change the matmul operation *functionally*. | |
| self.operation_kind = operation_kind | |
| self.matmul_shape = matmul_shape | |
| self.M = matmul_shape[0] | |
| self.N = matmul_shape[1] | |
| self.K = matmul_shape[2] | |
| self.batch_count = batch_count | |
| self.lhs = lhs # TensorDescription | |
| self.rhs = rhs # TensorDescription | |
| self.result = result # TensorDescription | |
| # Parameters that change the matmul operation *performance*. | |
| self.split_k_slices = split_k_slices | |
| def __eq__(self, other): | |
| """Returns true if the matmul operation is *functionally* the same.""" | |
| return ( | |
| self.matmul_shape == other.matmul_shape | |
| and self.lhs == other.lhs | |
| and self.rhs == other.rhs | |
| and self.result == other.result | |
| and self.batch_count == other.batch_count | |
| ) | |
| def name(self): | |
| """Procedurally generated name for the matmul operation. | |
| The name uniquely identifies a matmul operation with matmul shape, | |
| lhs dataype and layout, rhs datatype and layout, and result | |
| datatype and layout. | |
| """ | |
| return ( | |
| f"{OperationKindNames[self.operation_kind]}_" | |
| f"{self.M}x{self.N}x{self.K}_" | |
| f"{DataTypeName[self.lhs.datatype]}{ShortLayoutTypeName[self.lhs.layout]}_" | |
| f"{DataTypeName[self.rhs.datatype]}{ShortLayoutTypeName[self.rhs.layout]}_" | |
| f"{DataTypeName[self.result.datatype]}{ShortLayoutTypeName[self.result.layout]}" | |
| ) | |
| def get_argument_dict(self): | |
| """Returns the dictionary of matmul arguments (shape, datatypes, split_k_slices).""" | |
| split_k_mode = ( | |
| "parallel" if self.operation_kind == OperationKind.SplitkMatmul else "N/A" | |
| ) | |
| split_k_slices = ( | |
| self.split_k_slices | |
| if self.operation_kind == OperationKind.SplitkMatmul | |
| else "N/A" | |
| ) | |
| return { | |
| "batch_count": self.batch_count, | |
| "m": self.M, | |
| "n": self.N, | |
| "k": self.K, | |
| "lhs": self.lhs.name(), | |
| "rhs": self.rhs.name(), | |
| "result": self.result.name(), | |
| "split_k_mode": split_k_mode, | |
| "split_k_slices": split_k_slices, | |
| } | |
| def get_dict_entry(self): | |
| """Returns the dictionary of matmul operation summary.""" | |
| dict_entry = { | |
| "op_kind": OperationKindNames[self.operation_kind], | |
| "Operation": self.name(), | |
| "bytes": self.bytes(), | |
| "flops": self.flops(), | |
| } | |
| dict_entry.update(self.get_argument_dict()) | |
| return dict_entry | |
| def lhs_npy_shape(self): | |
| """Returns the shape of the lhs numpy array as a string in the format "MxKxDataType".""" | |
| return f"{self.M}x{self.K}x{DataTypeName[self.lhs.datatype]}" | |
| def rhs_npy_shape(self): | |
| """Returns the shape of the rhs numpy array as a string in the format "KxNxDataType".""" | |
| return f"{self.K}x{self.N}x{DataTypeName[self.rhs.datatype]}" | |
| def result_npy_shape(self): | |
| """Returns the shape of the result numpy array as a string in the format "MxNxDataType".""" | |
| return f"{self.M}x{self.N}x{DataTypeName[self.result.datatype]}" | |
| def bytes(self): | |
| """Returns the number of bytes read/written by the matmul operation.""" | |
| bytes = ( | |
| (DataTypeSizeInBits[self.lhs.datatype] * self.M // 8) * self.K | |
| + (DataTypeSizeInBits[self.rhs.datatype] * self.K // 8) * self.N | |
| + (DataTypeSizeInBits[self.result.datatype] * self.M // 8) * self.N | |
| ) | |
| return bytes * self.batch_count | |
| def flops(self): | |
| """Returns the number of floating point operations performed by the matmul operation.""" | |
| return 2 * self.M * self.N * self.K * self.batch_count | |
| ############################################################################## | |
| class MatmulCompilationInfo: | |
| """Data structure strictly describes the compilation passes and the tiling configurations. | |
| For a matrix multiplication operation, compilation passes and tiling configuration | |
| influences the performance of the compiled matmul operation, but the functionality. | |
| This data structure should be independent of the matmul operation functionality. | |
| Any change in this data structure should not affect the functionality of the matmul operation, i.e., | |
| we should be able to use the same reference results for a matrix operation compiled with different | |
| compilation info. | |
| """ | |
| def __init__( | |
| self, | |
| tile_description, | |
| translation_info, | |
| operation_kind=OperationKind.Matmul, | |
| config_type=CompilationConfigType.Custom, | |
| ): | |
| self.tile_description = tile_description # TileDescription | |
| self.translation_info = translation_info # TranslationInfo | |
| self.operation_kind = operation_kind # OperationKind | |
| self.config_type = config_type # CompilationConfigType | |
| def get_dict_entry(self): | |
| """Returns the dictionary entry for the matmul compilation info.""" | |
| if self.config_type == CompilationConfigType.Default: | |
| return { | |
| "Tile config": "Default", | |
| "Core class": "Default", | |
| "Instruction class": "Default", | |
| } | |
| translation_info_name = TranslationInfoName[self.translation_info] | |
| return { | |
| "Tile config": self.tile_description.name(), | |
| "Core class": translation_info_name.split("_")[0], | |
| "Instruction class": translation_info_name.split("_")[1], | |
| } | |
| def name(self): | |
| """Procedurally generated name for the matmul compilation info.""" | |
| if self.config_type == CompilationConfigType.Default: | |
| return "tile_config_default" | |
| return "tile_config_{tbm}x{tbn}_{tbk}x{stages}_{translation_info}".format( | |
| tbm=self.tile_description.threadblock_shape[0], | |
| tbn=self.tile_description.threadblock_shape[1], | |
| tbk=self.tile_description.threadblock_shape[2], | |
| stages=self.tile_description.stages, | |
| translation_info=TranslationInfoName[self.translation_info], | |
| ) | |
| ################################################################################ | |
| class EmitMatmulCompilationInfo: | |
| """Emitters for the matmul compilation info.""" | |
| def __init__(self): | |
| # matmul compilation info template | |
| self.matmul_compilation_info_template = """ | |
| // matmul compilation info (tile configuration, translation info, workgroup size) | |
| #${compilation_info_name} = #iree_codegen.compilation_info< | |
| lowering_config = <tile_sizes = [[${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}]]>, | |
| translation_info = <${translation_info} pipeline_depth = ${stages}>, | |
| workgroup_size = [${block_dim_x} : index, ${block_dim_y} : index, ${block_dim_z} : index] | |
| > | |
| """ | |
| # batch matmul and split-k matmul compilation info template | |
| self.batch_matmul_compilation_info_template = """ | |
| // batch matmul compilation info (tile configuration, translation info, workgroup size) | |
| #${compilation_info_name} = #iree_codegen.compilation_info< | |
| lowering_config = <tile_sizes = [[1, ${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}]]>, | |
| translation_info = <${translation_info} pipeline_depth = ${stages}>, | |
| workgroup_size = [${block_dim_x} : index, ${block_dim_y} : index, ${block_dim_z} : index] | |
| > | |
| """ | |
| def emit(self, compilation_info): | |
| """Emits the matmul compilation info as a string.""" | |
| if compilation_info.config_type == CompilationConfigType.Default: | |
| return "" | |
| values = { | |
| "compilation_info_name": compilation_info.name(), | |
| "translation_info": TranslationInfoTag[compilation_info.translation_info], | |
| "threadblock_shape_m": str( | |
| compilation_info.tile_description.threadblock_shape[0] | |
| ), | |
| "threadblock_shape_n": str( | |
| compilation_info.tile_description.threadblock_shape[1] | |
| ), | |
| "threadblock_shape_k": str( | |
| compilation_info.tile_description.threadblock_shape[2] | |
| ), | |
| "stages": str(compilation_info.tile_description.stages), | |
| "block_dim_x": str(compilation_info.tile_description.block_dim[0]), | |
| "block_dim_y": str(compilation_info.tile_description.block_dim[1]), | |
| "block_dim_z": str(compilation_info.tile_description.block_dim[2]), | |
| } | |
| # linalg.matmul (without split-k) compilation info template. | |
| compilation_info_template = self.matmul_compilation_info_template | |
| # linalg.batch_matmul and linalg.matmul (with split-k) have different | |
| # compilation info template from the linalg.matmul (without split-k). | |
| if ( | |
| compilation_info.operation_kind == OperationKind.BatchMatmul | |
| or compilation_info.operation_kind == OperationKind.SplitkMatmul | |
| ): | |
| compilation_info_template = self.batch_matmul_compilation_info_template | |
| return SubstituteTemplate(compilation_info_template, values) | |
| ############################################################################### | |
| class EmitLinalgMatmulDispatch: | |
| """Emitters for the `linalg.matmul` dispatch.""" | |
| def __init__(self): | |
| self.mlir_dialect = MlirDialect.Linalg | |
| # linalg.matmul mlir template | |
| self.linalg_row_row_matmul_template = """ | |
| // Dispatch linalg.matmul row-row layout | |
| func.func @${operation_name}_${compilation_info_name}( | |
| %lhs: tensor<${problem_m}x${problem_k}x${datatype_lhs}>, | |
| %rhs: tensor<${problem_k}x${problem_n}x${datatype_rhs}>) -> tensor<${problem_m}x${problem_n}x${datatype_result}> | |
| { | |
| %c0 = arith.constant 0.0 : ${datatype_result} | |
| %init = tensor.empty() : tensor<${problem_m}x${problem_n}x${datatype_result}> | |
| %inital_result = linalg.fill ins(%c0 : ${datatype_result}) outs(%init : tensor<${problem_m}x${problem_n}x${datatype_result}>) -> tensor<${problem_m}x${problem_n}x${datatype_result}> | |
| %result = linalg.matmul ${compilation_info_attribute} | |
| ins(%lhs, %rhs: tensor<${problem_m}x${problem_k}x${datatype_lhs}>, tensor<${problem_k}x${problem_n}x${datatype_rhs}>) | |
| outs(%inital_result: tensor<${problem_m}x${problem_n}x${datatype_result}>) -> tensor<${problem_m}x${problem_n}x${datatype_result}> | |
| return %result : tensor<${problem_m}x${problem_n}x${datatype_result}> | |
| } | |
| """ | |
| def emit(self, matmul_dispatch): | |
| """Emit the matmul operation in the MLIR dialect for a single compilation info""" | |
| compilation_info_attribute_template = ( | |
| """{compilation_info = #${compilation_info_name}}""" | |
| ) | |
| compilation_info_attribute_str = SubstituteTemplate( | |
| compilation_info_attribute_template, | |
| {"compilation_info_name": matmul_dispatch.configuration.name()}, | |
| ) | |
| compilation_info_attribute = ( | |
| compilation_info_attribute_str | |
| if matmul_dispatch.configuration.config_type | |
| != CompilationConfigType.Default | |
| else "" | |
| ) | |
| values = { | |
| "operation_name": matmul_dispatch.operation.name(), | |
| "compilation_info_attribute": compilation_info_attribute, | |
| "problem_m": str(matmul_dispatch.operation.M), | |
| "problem_n": str(matmul_dispatch.operation.N), | |
| "problem_k": str(matmul_dispatch.operation.K), | |
| "datatype_lhs": DataTypeName[matmul_dispatch.operation.lhs.datatype], | |
| "datatype_rhs": DataTypeName[matmul_dispatch.operation.rhs.datatype], | |
| "datatype_result": DataTypeName[matmul_dispatch.operation.result.datatype], | |
| "compilation_info_name": matmul_dispatch.configuration.name(), | |
| } | |
| return SubstituteTemplate(self.linalg_row_row_matmul_template, values) | |
| ############################################################################### | |
| class ReferenceMatmulOp(ReferenceOpInterface): | |
| """Reference implementation for the matmul operation in numpy.""" | |
| def __init__(self, matmul_operation, op_reference_cache_path, dist_lhs, dist_rhs): | |
| self.matmul_operation = matmul_operation | |
| self.op_reference_cache_path = op_reference_cache_path | |
| # Problem shape. | |
| self.M = matmul_operation.M | |
| self.N = matmul_operation.N | |
| self.K = matmul_operation.K | |
| # Data type for the input and result matrices. | |
| self.dtype_lhs = DataTypeNumPyTag[matmul_operation.lhs.datatype] | |
| self.dtype_rhs = DataTypeNumPyTag[matmul_operation.rhs.datatype] | |
| self.dtype_result = DataTypeNumPyTag[matmul_operation.result.datatype] | |
| # Distribution of the input tensors. | |
| self.dist_lhs = dist_lhs | |
| self.dist_rhs = dist_rhs | |
| # Filename for the left hand side input tensor. | |
| self.filename_lhs = ( | |
| "m{problem_m}xk{problem_k}_" | |
| "{tensor_description}_{dist}_lhs.npy".format( | |
| problem_m=self.M, | |
| problem_k=self.K, | |
| tensor_description=self.matmul_operation.lhs.name(), | |
| dist=DistributionName[self.dist_lhs], | |
| ) | |
| ) | |
| # Filename for the right hand side input tensor. | |
| self.filename_rhs = ( | |
| "k{problem_k}xn{problem_n}_" | |
| "{tensor_description}_{dist}_rhs.npy".format( | |
| problem_k=self.K, | |
| problem_n=self.N, | |
| tensor_description=self.matmul_operation.rhs.name(), | |
| dist=DistributionName[self.dist_rhs], | |
| ) | |
| ) | |
| # Filename for the reference result tensor. | |
| self.filename_reference_result = ( | |
| "m{problem_m}xn{problem_n}_" | |
| "{tensor_description}_reference_result.npy".format( | |
| problem_m=self.M, | |
| problem_n=self.N, | |
| tensor_description=self.matmul_operation.result.name(), | |
| ) | |
| ) | |
| # Filepath for input and output files. | |
| self.filepath_lhs = self.op_reference_cache_path.joinpath(self.filename_lhs) | |
| self.filepath_rhs = self.op_reference_cache_path.joinpath(self.filename_rhs) | |
| self.filepath_reference_result = self.op_reference_cache_path.joinpath( | |
| self.filename_reference_result | |
| ) | |
| def get_input_filepaths(self): | |
| """Returns the list of input file paths.""" | |
| return [self.filepath_lhs, self.filepath_rhs] | |
| def get_output_filepaths(self): | |
| """Returns the list of expected output file paths.""" | |
| return [self.filepath_reference_result] | |
| def __call__(self): | |
| """Generates input data, runs reference numpy.matmul, and save npy files to the output directory.""" | |
| # Generate the input data as np.array for the matmul operation. | |
| lhs_np_array = get_np_array( | |
| self.matmul_operation.lhs, (self.M, self.K), self.dist_lhs | |
| ) | |
| rhs_np_array = get_np_array( | |
| self.matmul_operation.rhs, (self.K, self.N), self.dist_rhs | |
| ) | |
| # Run the reference np.matmul and generate result np.array. | |
| result = np.matmul(lhs_np_array, rhs_np_array) | |
| # Save the input data as np.array for the matmul operation. | |
| np.save(self.filepath_lhs, np.array(lhs_np_array, dtype=self.dtype_lhs)) | |
| np.save(self.filepath_rhs, np.array(rhs_np_array, dtype=self.dtype_rhs)) | |
| # Save the expected result as an np.array. | |
| np.save( | |
| self.filepath_reference_result, np.array(result, dtype=self.dtype_result) | |
| ) | |
| class CudaMatmulDispatchChecker: | |
| """Given a matmul dispatch, checks if the dispatch is supported by the target GPU.""" | |
| def __init__(self, args): | |
| self.args = args | |
| # CUDA shared memory capacity per SM in KB. | |
| self.sharedMemPerSm = { | |
| "sm_80": 163, # 1KB is reserved for the driver. | |
| "sm_86": 99, # 1KB is reserved for the driver | |
| } | |
| self.cuda_arch = self.args.cuda_arch | |
| self.cuda_smem_capacity_in_bytes = self.sharedMemPerSm[self.cuda_arch] << 10 | |
| def _is_tile_aligned_shape(self, dispatch): | |
| """Checks if the given dispatch is valid for CUDA.""" | |
| matmul_shape = dispatch.operation.matmul_shape | |
| threadblock_shape = dispatch.configuration.tile_description.threadblock_shape | |
| if len(matmul_shape) != len(threadblock_shape): | |
| raise ValueError( | |
| "Problem shape and threadblock shape must have the same rank." | |
| ) | |
| is_aligned = all(a % b == 0 for a, b in zip(matmul_shape, threadblock_shape)) | |
| return is_aligned | |
| def _cuda_smem_required_in_bytes(self, dispatch): | |
| """Returns size bytes of shared memory required for a given cuda dispatch.""" | |
| threadblock_shape = dispatch.configuration.tile_description.threadblock_shape | |
| num_stages = dispatch.configuration.tile_description.stages | |
| tile_shape_lhs = threadblock_shape[0] * threadblock_shape[2] | |
| tile_shape_rhs = threadblock_shape[2] * threadblock_shape[1] | |
| return ( | |
| ( | |
| tile_shape_lhs * DataTypeSizeInBits[dispatch.operation.lhs.datatype] | |
| + tile_shape_rhs * DataTypeSizeInBits[dispatch.operation.rhs.datatype] | |
| ) | |
| * num_stages | |
| ) // 8 | |
| def _is_problem_k_divisible_by_split_k(self, dispatch): | |
| """Checks if the given dispatch is valid for CUDA.""" | |
| return dispatch.operation.K % dispatch.operation.split_k_slices == 0 | |
| def _is_cuda_smem_avialable(self, dispatch): | |
| """Checks if the given dispatch is valid for CUDA.""" | |
| return ( | |
| self._cuda_smem_required_in_bytes(dispatch) | |
| <= self.cuda_smem_capacity_in_bytes | |
| ) | |
| def is_valid(self, dispatch): | |
| """Checks if the given dispatch is valid for CUDA.""" | |
| if not self._is_tile_aligned_shape(dispatch): | |
| if self.args.verbose: | |
| print(f"[Warning]: {dispatch.name()} is not aligned is being skipped.") | |
| return False | |
| if not self._is_cuda_smem_avialable(dispatch): | |
| if self.args.verbose: | |
| print( | |
| f"[Warning]: {dispatch.name()} requires {self._cuda_smem_required_in_bytes(dispatch)} " | |
| f"bytes of shared memory, which is larger than the {self.cuda_arch} capacity " | |
| f"{self.cuda_smem_capacity_in_bytes} bytes." | |
| ) | |
| return False | |
| if (dispatch.operation.split_k_slices > 1) and ( | |
| not self._is_problem_k_divisible_by_split_k(dispatch) | |
| ): | |
| if self.args.verbose: | |
| print( | |
| f"[Warning]: {dispatch.name()} problem k is not divisible by {dispatch.operation.split_k_slices} " | |
| f"split-k slices, which is not supported on LLVM GPU CUDA backend." | |
| ) | |
| return False | |
| return True | |
| class CudaMatmulGenerator: | |
| """Matmul dispatch generator class. | |
| Generates a list of pre-defined matmul operations with resonable tuning cofigurations. | |
| The generator function are seperated based on the target backend and the data type. | |
| Please see example `MatmulGenerator._cuda_matmul_tensor_cores_f16` for cuda target | |
| backend and f16 data type.""" | |
| def __init__(self, args): | |
| """Initializes the matmul generator.""" | |
| self.args = args | |
| self.translation_infos = [ | |
| # TranslationInfo.LLVMGPUMatmulSimt, # CUDA Core (SMIT) | |
| # TranslationInfo.LLVMGPUMatmulTensorCore, # Tensor Core (WMMA) | |
| TranslationInfo.LLVMGPUMatmulTensorCoreMmaSync, # Tensor Core (MMA.SYNC) | |
| ] | |
| # List of pre-defined threadblock tile shapes for Tensor Core. | |
| self.tile_descriptions_tensor_cores_f16 = [ | |
| TileDescription([256, 128, 32], 3, [64, 4, 1]), | |
| TileDescription([128, 256, 32], 3, [128, 2, 1]), | |
| TileDescription([128, 128, 64], 4, [64, 2, 1]), | |
| TileDescription([128, 128, 32], 5, [64, 2, 1]), | |
| TileDescription([128, 64, 32], 5, [64, 2, 1]), | |
| TileDescription([64, 64, 64], 5, [64, 2, 1]), | |
| TileDescription([64, 64, 32], 10, [64, 2, 1]), | |
| ] | |
| self.tile_descriptions_tensor_cores_f32 = [ | |
| TileDescription([128, 256, 16], 3, [128, 2, 1]), | |
| TileDescription([256, 128, 16], 3, [64, 4, 1]), | |
| TileDescription([128, 128, 16], 5, [64, 2, 1]), | |
| TileDescription([128, 128, 32], 3, [64, 2, 1]), | |
| TileDescription([128, 128, 32], 4, [64, 2, 1]), | |
| TileDescription([128, 64, 32], 3, [64, 2, 1]), | |
| TileDescription([128, 64, 16], 5, [64, 2, 1]), | |
| TileDescription([64, 64, 32], 3, [64, 2, 1]), | |
| TileDescription([64, 64, 16], 10, [64, 2, 1]), | |
| ] | |
| # Create a list of matmul problem and initialize with some *default* shapes. | |
| self.matmul_shapes = [[256, 512, 128], [2560, 2560, 2560], [3456, 1024, 2048]] | |
| # Append matmul problem with *user* provided shapes. | |
| self.add_cmd_line_shapes() | |
| # Matmul dispatches collection. | |
| self.dispatches_collection_list = [] | |
| def add_cmd_line_shapes(self): | |
| """Adds matmul shapes from command line arguments.""" | |
| m_list = get_cmd_line_argument_list(self.args.problem_m) | |
| n_list = get_cmd_line_argument_list(self.args.problem_n) | |
| k_list = get_cmd_line_argument_list(self.args.problem_k) | |
| # If no command line matmul problem shapes are provided, only | |
| # use the default shapes. | |
| if len(m_list) == 0 and len(n_list) == 0 and len(k_list) == 0: | |
| return | |
| # If any of the command line matmul problem shapes are provided, | |
| # set the default shapes to empty problem dimension. | |
| if len(m_list) == 0: | |
| m_list = [256] | |
| if len(n_list) == 0: | |
| n_list = [256] | |
| if len(k_list) == 0: | |
| k_list = [256] | |
| # Append the command line matmul problem shapes with user-proivded | |
| # matmul problem shapes. | |
| for m in m_list: | |
| for n in n_list: | |
| for k in k_list: | |
| self.matmul_shapes.append([m, n, k]) | |
| def _cuda_supported_configuration_list(self, operation, configuration_list): | |
| """Returns a list of supported configurations for CUDA.""" | |
| supported_configuration_list = [] | |
| dispatch_checker = CudaMatmulDispatchChecker(self.args) | |
| for configuration in configuration_list: | |
| if not dispatch_checker.is_valid(Dispatch(operation, configuration)): | |
| continue | |
| supported_configuration_list.append(configuration) | |
| # Return the supported configuration list. | |
| return supported_configuration_list | |
| def _get_matmul_custom_compilation_info_list( | |
| self, tile_descriptions, translation_infos, operation_kind | |
| ): | |
| """Creates a *custom* list of matmul compilation info.""" | |
| configuration_list = [] | |
| for tile_description in tile_descriptions: | |
| for translation_info in translation_infos: | |
| configuration_list.append( | |
| MatmulCompilationInfo( | |
| tile_description, | |
| translation_info, | |
| operation_kind, | |
| CompilationConfigType.Custom, | |
| ) | |
| ) | |
| return configuration_list | |
| def _append_matmul_dispatch_collection( | |
| self, matmul_shapes, data_type, configuration_list | |
| ): | |
| """Appends the matmul dispatches collection with the given configuration list.""" | |
| # Create dispatches collection for each matmul_shape x configuration list.. | |
| for matmul_shape in matmul_shapes: | |
| operation = MatmulOperation( | |
| matmul_shape, | |
| TensorDescription(data_type[0], LayoutType.RowMajor), | |
| TensorDescription(data_type[1], LayoutType.RowMajor), | |
| TensorDescription(data_type[2], LayoutType.RowMajor), | |
| ) | |
| # Filter out configurations that are not supported by LLVM GPU CUDA backend. | |
| supported_configuration_list = self._cuda_supported_configuration_list( | |
| operation, configuration_list | |
| ) | |
| # Add default configuration if enabled. | |
| if self.args.default_config: | |
| supported_configuration_list.append( | |
| MatmulCompilationInfo( | |
| [], [], OperationKind.Matmul, CompilationConfigType.Default | |
| ) | |
| ) | |
| # Append the dispatch collection. | |
| self.dispatches_collection_list.append( | |
| DispatchCollection(operation, supported_configuration_list) | |
| ) | |
| def _cuda_matmul_tensor_cores_f16(self): | |
| """Appends dispatches for TensorCore with F16 input, F16 accum, F16 output.""" | |
| configuration_list = self._get_matmul_custom_compilation_info_list( | |
| self.tile_descriptions_tensor_cores_f16, | |
| self.translation_infos, | |
| OperationKind.Matmul, | |
| ) | |
| data_type = [DataType.f16, DataType.f16, DataType.f16] | |
| self._append_matmul_dispatch_collection( | |
| self.matmul_shapes, data_type, configuration_list | |
| ) | |
| def _cuda_matmul_tensor_cores_f32(self): | |
| """Appends dispatches for TensorCore with F32 input, F32 accum, F32 output.""" | |
| configuration_list = self._get_matmul_custom_compilation_info_list( | |
| self.tile_descriptions_tensor_cores_f32, | |
| self.translation_infos, | |
| OperationKind.Matmul, | |
| ) | |
| data_type = [DataType.f32, DataType.f32, DataType.f32] | |
| self._append_matmul_dispatch_collection( | |
| self.matmul_shapes, data_type, configuration_list | |
| ) | |
| def _cuda_matmul_tensor_cores_mixed_precision(self): | |
| """Appends dispatches for TensorCore with F16 input, F32 accum, F32 output.""" | |
| configuration_list = self._get_matmul_custom_compilation_info_list( | |
| self.tile_descriptions_tensor_cores_f16, | |
| self.translation_infos, | |
| OperationKind.Matmul, | |
| ) | |
| data_type = [DataType.f16, DataType.f16, DataType.f32] | |
| self._append_matmul_dispatch_collection( | |
| self.matmul_shapes, data_type, configuration_list | |
| ) | |
| def generate(self): | |
| """Generates a list of matmul operations.""" | |
| self._cuda_matmul_tensor_cores_f16() | |
| self._cuda_matmul_tensor_cores_f32() | |
| self._cuda_matmul_tensor_cores_mixed_precision() | |
| return self.dispatches_collection_list |