# 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