import enum
import os.path
import shutil
import functools
import operator
import collections
import subprocess
from library import *
from dispatch import *


################################################################################
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 passes and tiling configuration. 
     It strictly contains the parameter that changes the functionality of matmul 
     operation.
  """

  def __init__(self, problem_shape, lhs, rhs, result):
    """Initializes a matrix multiplication operation.
    Matrix-multiple operation: `result[M, N] = lhs[M, K] * rhs[K, N]`
    problem_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.
    """
    self.problem_shape = problem_shape  # Matmul problem shape [M, N, K]
    self.M = problem_shape[0]
    self.N = problem_shape[1]
    self.K = problem_shape[2]

    self.lhs = lhs  # TensorDescription
    self.rhs = rhs  # TensorDescription
    self.result = result  # TensorDescription
    self.operation_kind = OperationKind.Matmul

  def __eq__(self, other):
    return self.problem_shape == other.problem_shape and \
           self.lhs == other.lhs and \
           self.rhs == other.rhs and \
           self.result == other.result

  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 "matmul_{m}x{n}x{k}_"\
      "{typeLhs}{layoutLhs}_"\
      "{typeRhs}{layoutRhs}_"\
      "{typeResult}{layoutResult}".format(
      m=self.problem_shape[0],
      n=self.problem_shape[1],
      k=self.problem_shape[2],
      typeLhs=DataTypeName[self.lhs.datatype],
      layoutLhs=ShortLayoutTypeName[self.lhs.layout],
      typeRhs=DataTypeName[self.rhs.datatype],
      layoutRhs=ShortLayoutTypeName[self.rhs.layout],
      typeResult=DataTypeName[self.result.datatype],
      layoutResult=ShortLayoutTypeName[self.result.layout]
    )

  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

  def flops(self):
    """Returns the number of floating point operations performed by the matmul operation."""
    return 2 * self.M * self.N * self.K


##############################################################################
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,
               config_type=CompilationConfigType.Custom):
    self.tile_description = tile_description  # TileDescription
    self.translation_info = translation_info  # TranslationInfo
    self.config_type = config_type  # CompilationConfigType

  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):
    self.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]
>
"""

  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]),
    }

    return SubstituteTemplate(self.compilation_info_template, values)


###############################################################################
class EmitLinalgMatmulDispatch:
  """Emitters for the `linalg.matmul` dispatch."""

  def __init__(self):
    self.mlir_dialect = MlirDialect.Linalg

    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.problem_shape[0]),
        'problem_n':
            str(matmul_dispatch.operation.problem_shape[1]),
        'problem_k':
            str(matmul_dispatch.operation.problem_shape[2]),
        '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)


# Emit `mhlo.matmul` operation.
# TODO: Add support for testing lowering matmul op from other dialect.
class EmitMhloMatmulOperation:

  def __init__(self):
    self.mlir_dialect = MlirDialect.Mhlo

    self.linalg_row_row_matmul_template = """
// mhlo.matmul operation row-row layout
"""


###############################################################################
class ReferenceMatmulOperation:
  """Reference implementation for the matmul operation in numpy.
      ReferenceMatmulOperation class has the following responsibilities:
       1) Generates matmul operation inputs as np.array for a desired 
          distribution.
       2) Runs the matmul reference operation in np.matmul.
       3) Generates the matmul operation expected output as np.array.
       4) Additional, generate input and output filename strings.
  """

  def __init__(self, matmul_operation, dist_lhs, dist_rhs):
    self.matmul_operation = matmul_operation
    # Problem shape.
    self.M = matmul_operation.problem_shape[0]
    self.N = matmul_operation.problem_shape[1]
    self.K = matmul_operation.problem_shape[2]

    # 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.reference_filename_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())

  # Generates input data, runs reference numpy.matmul, and save npy files to the output directory.
  def run_and_save(self, output_dir="."):

    # 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(os.path.join(output_dir, self.filename_lhs),\
             np.array(lhs_np_array, dtype = self.dtype_lhs))
    np.save(os.path.join(output_dir, self.filename_rhs),\
             np.array(rhs_np_array, dtype = self.dtype_rhs))

    # Save the expected result as an np.array.
    np.save(os.path.join(output_dir, self.reference_filename_result),\
             np.array(result, dtype = self.dtype_result))


###############################################################################
class MatmulOperationLauncher:
  """Launches the compilation and execution of the matmul operation.
  MatmulOperationLauncher class has the following responsibilities:
    1) Launches compilation of the matmul for a verification or profiling runs.
    2) Launches the verification by running the IREE compiled matmul operation
       and the python nmpy.matmul.
    3) Launches the profiling by running the IREE compiled matmul operation.
  """

  def __init__(self, args, operation):
    self.operation = operation

    # Variables from top-level argparse.
    self.generated_path = os.path.join(args.build_dir, 'generated',
                                       args.mlir_dialect)
    self.args = args
    self.benchmark_dispatch_repeat_count = args.batch_size
    self.batch_size = args.batch_size

    # Additional paths.
    self.matmul_path = os.path.join(self.generated_path, 'matmul')
    self.operation_path = os.path.join(self.matmul_path, operation.name())
    self.source_mlir_file = os.path.join(self.operation_path,
                                         operation.name() + '.mlir')

    # path to iree-compile tool. (for compiling the input mlir file to vmfb)
    self.iree_compile_path = os.path.join(args.build_dir, 'tools',
                                          'iree-compile')

    # path to iree-benchmark-module tool. (for performance benchmarking and profiling)
    self.iree_benchmark_module_path = os.path.join(args.build_dir, 'tools',
                                                   'iree-benchmark-module')

    # path to iree-run-module tool. (for verification)
    self.iree_run_module_path = os.path.join(args.build_dir, 'tools',
                                             'iree-run-module')

    # output vmfb files for the operation.
    self.vmfb_verify_file = os.path.join(self.operation_path,
                                         self.operation.name() + '_verify.vmfb')
    self.vmfb_benchmark_file = os.path.join(
        self.operation_path,
        self.operation.name() + '_benchmark.vmfb')

  def compile(self, compilation_mode):
    """Compiles the matmul operation to a vmfb file for profiling."""

    benchmark_dispatch_repeat_count = self.benchmark_dispatch_repeat_count if compilation_mode == CompilationMode.Profile else 1
    vmfb_file = self.vmfb_benchmark_file if compilation_mode == CompilationMode.Profile else self.vmfb_verify_file

    # Base iree-compile commandline
    cmd = [self.iree_compile_path, self.source_mlir_file, "-o", f"{vmfb_file}"]

    # General compilation options
    cmd += [f"--iree-hal-target-backends={self.args.device}"]
    cmd += [f"--iree-hal-cuda-llvm-target-arch={self.args.cuda_arch}"]
    if self.args.split_k_slices != "":
      cmd += [f"--iree-flow-split-matmul-reduction={self.args.split_k_slices}"]
    if self.args.use_mma_sync:
      cmd += [f"--iree-codegen-llvmgpu-use-mma-sync"]
    if self.args.use_wmma:
      cmd += [f"--iree-codegen-llvmgpu-use-wmma"]

    # Compilation options for profiling
    cmd += [
        f"--iree-hal-benchmark-dispatch-repeat-count={benchmark_dispatch_repeat_count}"
    ]

    if not os.path.exists(vmfb_file) or self.args.force_compile:
      print(
          f">> Compilation command for {CompilationModeNames[compilation_mode]} : {' '.join(cmd)}"
      )
      subprocess.check_output(cmd)

    elif self.args.verbose:
      print("Skipping compilation of matmul operation: " + vmfb_file +
            " since it already exists.")

  def verify(self, configuration):
    """Verifies the matmul operation with a given configuration."""
    # First compile the operation to a vmfb file.
    self.compile(CompilationMode.Verify)

    # Verify using random data distribution.
    # TODO 1) make input distribution configurable through command line.
    # TODO 2) make the reference run to check if reference npy files are present,
    #         then do not re-run the reference.
    reference_op = ReferenceMatmulOperation(self.operation,\
                                         Distribution.Random,\
                                         Distribution.Random)

    lhs_npy_file = os.path.join(self.operation_path, reference_op.filename_lhs)
    rhs_npy_file = os.path.join(self.operation_path, reference_op.filename_rhs)
    expected_result_npy_file = os.path.join(
        self.operation_path, reference_op.reference_filename_result)

    # If the reference numpy do not exists, run the reference implementation
    # and generate npy files.
    if not os.path.exists(lhs_npy_file) or \
       not os.path.exists(rhs_npy_file) or \
       not os.path.exists(expected_result_npy_file):
      reference_op.run_and_save(self.operation_path)

    # Commandline `iree-run-module` for verification.
    cmd = [
        self.iree_run_module_path, f'--module={self.vmfb_verify_file}',
        f'--device={self.args.device}'
    ]

    # Operation-specific verification command-line.
    # TODO: abstract the operation-specific verification command-line out of verification.
    cmd.append(f'--function={self.operation.name()}_{configuration.name()}')
    cmd.append(f'--input=@{lhs_npy_file}')
    cmd.append(f'--input=@{rhs_npy_file}')
    cmd.append(f'--expected_output=@{expected_result_npy_file}')

    # Print the command if verbose.
    if self.args.verbose:
      print(">> Verification command: " + ' '.join(cmd))

    # Launch verification.
    cmd_output = subprocess.check_output(cmd, text=True)

    # Parse the verification output.
    m = re.search(r"\[(?P<verification_result>[a-zA-Z]+)\]", cmd_output)
    if m is None:
      raise Exception(
          "Failed to parse verification output by iree-run-module: " +
          cmd_output)
    verification_result = m.group('verification_result')

    if self.args.verbose or verification_result != "SUCCESS":
      print(cmd_output)

    return verification_result

  def profile(self, configuration):
    """Profiles the matmul operation with a given configuration."""
    # First compile the operation to a vmfb file.
    self.compile(CompilationMode.Profile)

    # Commandline `iree-benchmark-module` for profiling.
    cmd = [
        self.iree_benchmark_module_path, f'--module={self.vmfb_benchmark_file}',
        f'--device={self.args.device}'
    ]

    # Profiling specific flags.
    cmd += [f'--benchmark_repetitions={self.args.benchmark_repetitions}']
    cmd += [f'--batch_size={self.batch_size}']

    # Operation-specific profiling command-line.
    cmd += [f'--function={self.operation.name()}_{configuration.name()}']
    cmd += [f'--input={self.operation.lhs_npy_shape()}']
    cmd += [f'--input={self.operation.rhs_npy_shape()}']

    # Print the command if verbose.
    if self.args.verbose:
      print(">> Profiling command: " + ' '.join(cmd))

    # Launch profiling.
    cmd_output = subprocess.check_output(cmd,
                                         text=True,
                                         stderr=subprocess.STDOUT)

    # Parse the profiling output.
    m = re.search(r"real_time_median\s+(?P<runtime>\d+.\d+)\s+ms", cmd_output)
    if m is None:
      raise Exception("Failed to parse runtime from benchmark result: " +
                      cmd_output)
    runtime_in_ms = float(m.group('runtime'))
    return runtime_in_ms


##############################################################################
class MatmulGenerator:
  """Matmul dispatch generator class.
  Generates a list of pre-definied 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):
    self.args = args

    self.default_config = False if args.default_config in [
        'False', 'false', '0'
    ] else True

    self.translation_infos = [
        #TranslationInfo.LLVMGPUMatmulSimt,  # CUDA Core (SMIT)
        #TranslationInfo.LLVMGPUMatmulTensorCore, # Tensor Core (WMMA)
        TranslationInfo.
        LLVMGPUMatmulTensorCoreMmaSync,  # Tensor Core (MMA.SYNC)
    ]

    self.problem_shapes = [[128, 256, 8192]]
    """
    self.problem_shapes = [[128, 128, 256], [256, 512, 128], [1024, 512, 2048],
                           [2560, 2560, 2560], [3456, 1024, 2048]]
    """

    # List of pre-definied matmul dispatch collections.
    self.dispatches_collection_list = []

    # CUDA specific constants.
    self.cuda_warp_size = 32
    self.cuda_smem_capacity_in_bytes_sm80 = 192 << 10

  def _is_tile_aligned_shape(self, dispatch):
    """Checks if the given dispatch is valid for CUDA."""
    problem_shape = dispatch.operation.problem_shape
    threadblock_shape = dispatch.configuration.tile_description.threadblock_shape
    if len(problem_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(problem_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_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_sm80

  def _cuda_supported_configuration_list(self, operation, configuration_list):
    """Returns a list of supported configurations for CUDA."""
    supported_configuration_list = []
    for configuration in configuration_list:
      dispatch = Dispatch(operation, configuration)
      if not self._is_tile_aligned_shape(dispatch):
        print(f"Warning: {dispatch.name()} is not aligned is being skipped.")
        continue
      if not self._is_cuda_smem_avialable(dispatch):
        print(f"Warning: {dispatch.name()} requires {self._cuda_smem_required_in_bytes(dispatch)} "\
              f"bytes of shared memory, which is larger than the SM80 capacity "\
              f"{self.cuda_smem_capacity_in_bytes_sm80} bytes.")
        continue

      # If all checks pass, add the configuration to the supported list.
      supported_configuration_list.append(configuration)

    return supported_configuration_list

  def _cuda_matmul_tensor_cores_f16(self):
    """Appends a list of matmul dispatches for GPU TensorCore F16 data type."""

    tile_descriptions = [
        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]),
    ]

    # Create configuration list from the tile descriptions and translation infos.
    configuration_list = []

    for tile_description in tile_descriptions:
      for translation_info in self.translation_infos:
        configuration_list.append(
            MatmulCompilationInfo(tile_description, translation_info))

    # Create dispatches collection for each problem shape with the configuration list.
    for problem_shape in self.problem_shapes:
      operation = MatmulOperation(
        problem_shape,\
        TensorDescription(DataType.f16, LayoutType.RowMajor), \
        TensorDescription(DataType.f16, LayoutType.RowMajor), \
        TensorDescription(DataType.f16, 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 requested.
      if self.default_config:
        supported_configuration_list.append(
            MatmulCompilationInfo([], [], CompilationConfigType.Default))

      self.dispatches_collection_list.append(DispatchCollection(\
        operation, supported_configuration_list))

  ################################################################################
  def _cuda_matmul_tensor_cores_f32(self):
    """Appends a list of matmul dispatches for GPU TensorCore F32 data type."""

    tile_descriptions = [
        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([64, 64, 64], 3, [64, 2, 1]),
    ]

    # Create dispatches collection for each problem shape with the configuration list.
    configuration_list = []

    for tile_description in tile_descriptions:
      for translation_info in self.translation_infos:
        configuration_list.append(
            MatmulCompilationInfo(tile_description, translation_info))

    for problem_shape in self.problem_shapes:
      operation = MatmulOperation(
        problem_shape,\
        TensorDescription(DataType.f32, LayoutType.RowMajor), \
        TensorDescription(DataType.f32, LayoutType.RowMajor), \
        TensorDescription(DataType.f32, 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 requested.
      if self.default_config:
        supported_configuration_list.append(
            MatmulCompilationInfo([], [], CompilationConfigType.Default))

      self.dispatches_collection_list.append(DispatchCollection(\
        operation, supported_configuration_list))

  def generate(self):
    """Generates a list of matmul operations."""
    if self.args.device == 'cuda':
      self._cuda_matmul_tensor_cores_f16()
      self._cuda_matmul_tensor_cores_f32()
    return self.dispatches_collection_list


##############################################################################