# Copyright 2022 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
from library import *
from dispatch import *
from matmul import MatmulOperation, MatmulCompilationInfo, CudaMatmulGenerator


class BatchMatmulOperation(MatmulOperation):
  """Data structure to describe a batch matrix multiplication operation."""

  def __init__(self, bmm_shape, lhs, rhs, result):
    assert len(bmm_shape) == 4, "Batch matmul shape must be 4D"
    super().__init__(bmm_shape[1:], lhs, rhs, result, bmm_shape[0], 1,
                     OperationKind.BatchMatmul)

  def name(self):
    return f'{OperationKindNames[self.operation_kind]}_'\
           f'{self.batch_count}x{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 lhs_npy_shape(self):
    return f'{self.batch_count}x{super().lhs_npy_shape()}'

  def rhs_npy_shape(self):
    return f'{self.batch_count}x{super().rhs_npy_shape()}'

  def result_npy_shape(self):
    return f'{self.batch_count}x{super().result_npy_shape()}'


class EmitLinalgBatchMatmulDispatch:
  """Emitters for the `linalg.batch_matmul` dispatch."""

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

    self.linalg_row_row_matmul_template = """
// Dispatch linalg.batch_matmul row-row layout 
func.func @${operation_name}_${compilation_info_name}(
  %lhs: tensor<${batch_count}x${problem_m}x${problem_k}x${datatype_lhs}>,
  %rhs: tensor<${batch_count}x${problem_k}x${problem_n}x${datatype_rhs}>) -> tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>
{
  %c0 = arith.constant 0.0 : ${datatype_result}
  %init = tensor.empty() : tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>
  %inital_result = linalg.fill ins(%c0 : ${datatype_result}) outs(%init : tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>) -> tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>
  %result = linalg.batch_matmul ${compilation_info_attribute} 
                     ins(%lhs, %rhs: tensor<${batch_count}x${problem_m}x${problem_k}x${datatype_lhs}>, tensor<${batch_count}x${problem_k}x${problem_n}x${datatype_rhs}>)
                     outs(%inital_result: tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>) -> tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>
  return %result : tensor<${batch_count}x${problem_m}x${problem_n}x${datatype_result}>
}
"""

  def emit(self, 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': dispatch.configuration.name()})
    compilation_info_attribute = compilation_info_attribute_str \
      if dispatch.configuration.config_type != CompilationConfigType.Default else ""

    values = {
        'operation_name': dispatch.operation.name(),
        'compilation_info_attribute': compilation_info_attribute,
        'batch_count': str(dispatch.operation.batch_count),
        'problem_m': str(dispatch.operation.M),
        'problem_n': str(dispatch.operation.N),
        'problem_k': str(dispatch.operation.K),
        'datatype_lhs': DataTypeName[dispatch.operation.lhs.datatype],
        'datatype_rhs': DataTypeName[dispatch.operation.rhs.datatype],
        'datatype_result': DataTypeName[dispatch.operation.result.datatype],
        'compilation_info_name': dispatch.configuration.name()
    }

    return SubstituteTemplate(self.linalg_row_row_matmul_template, values)


class ReferenceBatchMatmulOp(ReferenceOpInterface):
  """Reference implementation for the batch matmul operation in numpy."""

  def __init__(self, bmm_operation, op_reference_cache_path, dist_lhs,
               dist_rhs):
    self.bmm_operation = bmm_operation
    self.op_reference_cache_path = op_reference_cache_path

    if not self.op_reference_cache_path.exists():
      self.op_reference_cache_path.mkdir()

    # Problem shape.
    self.batch_count = bmm_operation.batch_count
    self.M = bmm_operation.M
    self.N = bmm_operation.N
    self.K = bmm_operation.K

    # Data type for the input and result matrices.
    self.dtype_lhs = DataTypeNumPyTag[bmm_operation.lhs.datatype]
    self.dtype_rhs = DataTypeNumPyTag[bmm_operation.rhs.datatype]
    self.dtype_result = DataTypeNumPyTag[bmm_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 = "batch_count{batch_count}xm{problem_m}xk{problem_k}_"\
      "{tensor_description}_{dist}_lhs.npy".format(
      batch_count=self.batch_count,
      problem_m=self.M,
      problem_k=self.K,
      tensor_description=self.bmm_operation.lhs.name(),
      dist=DistributionName[self.dist_lhs])

    # Filename for the right hand side input tensor.
    self.filename_rhs = "batch_count{batch_count}xk{problem_k}xn{problem_n}_"\
      "{tensor_description}_{dist}_rhs.npy".format(
      batch_count=self.batch_count,
      problem_k=self.K,
      problem_n=self.N,
      tensor_description=self.bmm_operation.rhs.name(),
      dist=DistributionName[self.dist_rhs])

    # Filename for the reference result tensor.
    self.filename_reference_result = "batch_count{batch_count}xm{problem_m}xn{problem_n}_"\
      "{tensor_description}_reference_result.npy".format(
      batch_count=self.batch_count,
      problem_m=self.M,
      problem_n=self.N,
      tensor_description=self.bmm_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.bmm_operation.lhs,
                                (self.batch_count, self.M, self.K),
                                self.dist_lhs)
    rhs_np_array = get_np_array(self.bmm_operation.rhs,
                                (self.batch_count, 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 CudaBatchMatmulGenerator(CudaMatmulGenerator):
  """Batch matmul dispatch generator class. """

  def __init__(self, args):
    """Initializes the batch matmul dispatch generator."""
    super().__init__(args)

    # Predefined batch matmul problem shapes.
    self.batch_matmul_shapes = [[16, 512, 64, 512]]

    # Batch matmul dispatches collection.
    self.dispatches_collection_list = []

  def _append_matmul_dispatch_collection(self, bmm_shapes, data_type,
                                         configuration_list):
    """Update the batch matmul dispatch collection with the given configuration list."""

    # Create dispatches collection for each problem shape with the configuration list.
    for bmm_shape in bmm_shapes:
      operation = BatchMatmulOperation(
        bmm_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.BatchMatmul,
                                  CompilationConfigType.Default))

      # Append the dispatches collection.
      self.dispatches_collection_list.append(DispatchCollection(\
        operation, supported_configuration_list))

  def _cuda_matmul_tensor_cores_f16(self):
    """Appends a list of matmul dispatches for GPU TensorCore F16 data type."""
    configuration_list = self._get_matmul_custom_compilation_info_list(
        self.tile_descriptions_tensor_cores_f16, self.translation_infos,
        OperationKind.BatchMatmul)
    data_type = [DataType.f16, DataType.f16, DataType.f16]
    self._append_matmul_dispatch_collection(self.batch_matmul_shapes, data_type,
                                            configuration_list)

  def _cuda_matmul_tensor_cores_f32(self):
    """Appends a list of matmul dispatches for GPU TensorCore F32 data type."""
    configuration_list = self._get_matmul_custom_compilation_info_list(
        self.tile_descriptions_tensor_cores_f32, self.translation_infos,
        OperationKind.BatchMatmul)
    data_type = [DataType.f32, DataType.f32, DataType.f32]
    self._append_matmul_dispatch_collection(self.batch_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()
    return self.dispatches_collection_list