# 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, re
from enum import auto
import numpy as np
from abc import ABC, abstractmethod
from collections import namedtuple

###################################################################################################
# This file contains library of enumerations and classes used to build operation descritpions.
# The operation descriptions are used to generate MLIR source files, performance tuning configuration,
# reference implementations, and numpy input/output files.

# The file is organized as follows:
# 1. Enumerated `Type`s grouped together for categories, For e.g. [Arch]Type, [Data]Type etc.
# 2. Dictonaries `Names` mapping the enumeration values to their string names.
#    For e.g. [Arch]TypeNames, [Data]TypeNames etc.
# 3. `Tags` for each enumeration value to be used in the generated MLIR source files.
#    For e.g. [TranslationInfo]Tags
###################################################################################################


# Architecure types
###################################################################################################
class ArchType(enum.Enum):
    Cpu = auto()
    Gpu = auto()


ArchTypeNames = {
    ArchType.Cpu: "cpu",
    ArchType.Gpu: "gpu",
}


class GpuArchType(enum.Enum):
    nvptx = auto()
    rocm = auto()
    spirv = auto()


GpuArchTypeNames = {
    GpuArchType.nvptx: "nvptx",
    GpuArchType.rocm: "rocm",
    GpuArchType.spirv: "spirv",
}


# Operation kinds
###################################################################################################
class OperationKind(enum.Enum):
    Matmul = auto()
    BatchMatmul = auto()
    SplitkMatmul = auto()
    Conv2d = auto()


OperationKindNames = {
    OperationKind.Matmul: "matmul",
    OperationKind.SplitkMatmul: "matmul_splitk",
    OperationKind.BatchMatmul: "batch_matmul",
    OperationKind.Conv2d: "conv2d",
}


# MLIR dialects
###################################################################################################
class MlirDialect(enum.Enum):
    Linalg = auto()
    Mhlo = auto()


MlirDialectNames = {
    MlirDialect.Linalg: "linalg",
    MlirDialect.Mhlo: "mhlo",
}


# Compilation modes (verification or benchmarking/profiling)
###################################################################################################
class CompilationMode(enum.Enum):
    Verify = auto()
    Profile = auto()


CompilationModeNames = {
    CompilationMode.Verify: "verify",
    CompilationMode.Profile: "profile",
}


class CompilationConfigType(enum.Enum):
    Default = auto()
    Custom = auto()


CompilationConfigTypeName = {
    CompilationConfigType.Default: "default",
    CompilationConfigType.Custom: "custom",
}


# Enumerations for data types and layouts
###################################################################################################
class DataType(enum.Enum):
    b1 = auto()
    u4 = auto()
    u8 = auto()
    u16 = auto()
    u32 = auto()
    u64 = auto()
    s4 = auto()
    s8 = auto()
    s16 = auto()
    s32 = auto()
    s64 = auto()
    e4m3 = auto()
    e5m2 = auto()
    f16 = auto()
    bf16 = auto()
    f32 = auto()
    tf32 = auto()
    f64 = auto()
    invalid = auto()


DataTypeName = {
    DataType.b1: "b1",
    DataType.u4: "u4",
    DataType.u8: "u8",
    DataType.u16: "u16",
    DataType.u32: "u32",
    DataType.u64: "u64",
    DataType.s4: "s4",
    DataType.s8: "s8",
    DataType.s16: "s16",
    DataType.s32: "s32",
    DataType.s64: "s64",
    DataType.e4m3: "e4m3",
    DataType.e5m2: "e5m2",
    DataType.f16: "f16",
    DataType.bf16: "bf16",
    DataType.f32: "f32",
    DataType.tf32: "tf32",
    DataType.f64: "f64",
}

DataTypeNumPyTag = {
    DataType.f16: np.float16,
    DataType.f32: np.float32,
}

DataTypeSizeInBits = {
    DataType.b1: 1,
    DataType.u4: 4,
    DataType.u8: 8,
    DataType.u16: 16,
    DataType.u32: 32,
    DataType.u64: 64,
    DataType.s4: 4,
    DataType.s8: 8,
    DataType.s16: 16,
    DataType.s32: 32,
    DataType.s64: 64,
    DataType.e4m3: 8,
    DataType.e5m2: 8,
    DataType.f16: 16,
    DataType.bf16: 16,
    DataType.f32: 32,
    DataType.tf32: 32,
    DataType.f64: 64,
}


class LayoutType(enum.Enum):
    ColumnMajor = auto()
    RowMajor = auto()
    NHWC = auto()
    NCWH = auto()


# cuBLAS/cuDNN layout type names convention is followed for the layout names.
# https://docs.nvidia.com/cuda/cublas/index.html#cublasoperation-t
ShortLayoutTypeName = {
    LayoutType.ColumnMajor: "n",
    LayoutType.RowMajor: "t",
    LayoutType.NHWC: "nhwc",
    LayoutType.NCWH: "ncwh",
}


# Compilation pipelines/translation info.
###################################################################################################
class TranslationInfo(enum.Enum):
    LLVMGPUMatmulSIMT = auto()
    LLVMGPUMatmulTensorCore = auto()
    LLVMGPUMatmulTensorCoreMmaSync = auto()


TranslationInfoTag = {
    TranslationInfo.LLVMGPUMatmulSIMT: "LLVMGPUMatmulSIMT",
    TranslationInfo.LLVMGPUMatmulTensorCore: "LLVMGPUMatmulTensorCore",
    TranslationInfo.LLVMGPUMatmulTensorCoreMmaSync: "LLVMGPUMatmulTensorCoreMmaSync",
}

TranslationInfoName = {
    TranslationInfo.LLVMGPUMatmulSIMT: "simt_ffma",
    TranslationInfo.LLVMGPUMatmulTensorCore: "tensorcore_wmma",
    TranslationInfo.LLVMGPUMatmulTensorCoreMmaSync: "tensorcore_mmasync",
}


# Distribution of values in a tensor.
###################################################################################################
class Distribution(enum.Enum):
    Empty = auto()
    Zeros = auto()
    Ones = auto()
    Sequential = auto()
    Identity = auto()
    Random = auto()


DistributionName = {
    Distribution.Empty: "empty",
    Distribution.Zeros: "zeros",
    Distribution.Ones: "ones",
    Distribution.Sequential: "seq",
    Distribution.Identity: "identity",
    Distribution.Random: "random",
}

###################################################################################################
# The next part of this file contains the data structures for describing a tensor, tiles etc that
# are built using the above enumerations. These data structures are used to create compose bigger
# data structures that describe an operation or a sequence of operations, along with compilation
# pipeling to form a collection of dispatches to profiled.
###################################################################################################


class TensorDescription:
    """A class for tensor description."""

    def __init__(self, datatype, layout):
        self.datatype = datatype
        self.layout = layout

    def name(self):
        return "%s%s" % (DataTypeName[self.datatype], ShortLayoutTypeName[self.layout])


class TileDescription:
    """A class for tile description."""

    def __init__(self, threadblock_shape, stages, block_dim):
        self.threadblock_shape = threadblock_shape  # in number of elements in M, N, K
        self.stages = stages  # number of shared memory stages in tile K
        self.block_dim = block_dim  # block dimension in number of threads in x, y, z

    def name(self):
        return "%dx%d_%dx%d" % (
            self.threadblock_shape[0],
            self.threadblock_shape[1],
            self.threadblock_shape[2],
            self.stages,
        )


###################################################################################################
# The following part contains utility functions for which are used by the profiler tool.
# These function may be moved as the need for create a proper structure for the
# functionality they provide becomes apparent and necessary as we move forward.
###################################################################################################
def get_np_array(tensor_description, shape, dist):
    """Returns a numpy array based on the distribution and shape."""
    # Fix the seed for reproducibility.
    np.random.seed(42)

    # Generate the numpy array based on the distribution.
    if dist == Distribution.Empty:
        return np.empty(shape)
    elif dist == Distribution.Zeros:
        return np.zeros(shape)
    elif dist == Distribution.Ones:
        return np.ones(shape)
    elif dist == Distribution.Sequential:
        return np.arange(np.prod(shape)).reshape(shape)
    elif dist == Distribution.Identity:
        return np.eye(shape[0], shape[1])
    elif dist == Distribution.Random:
        if tensor_description.datatype == DataType.s8:
            return np.random.randint(-2, 3, shape)
        elif tensor_description.datatype == DataType.u8:
            return np.random.randint(0, 4, shape)
        elif (
            tensor_description.datatype == DataType.f16
            or tensor_description.datatype == DataType.bf16
        ):
            return np.random.randint(-3, 4, shape)
        elif tensor_description.datatype == DataType.f32:
            return np.random.randint(-7, 8, shape)


###################################################################################################
def SubstituteTemplate(template, values):
    """Substitutes values into a template string."""
    text = template
    for key, value in values.items():
        regex = "\\$\\{%s\\}" % key
        newtext = re.sub(regex, value, text)
        text = newtext
    return text


###################################################################################################
class ReferenceOpInterface(ABC):
    """Interface for reference implementations."""

    @abstractmethod
    def get_input_filepaths(self):
        """Returns the list of inputs."""
        pass

    @abstractmethod
    def get_output_filepaths(self):
        """Returns the list of outputs/."""
        pass

    @abstractmethod
    def __call__(self):
        """Runs the reference implementation."""
        pass

    def is_cached(self):
        """Returns whether the reference run is cached."""

        # Returns False if any of the reference input are missing.
        for input_filepath in self.get_input_filepaths():
            if not input_filepath.exists():
                return False

        # Returns False if any of the reference output are missing.
        for output_filepath in self.get_output_filepaths():
            if not output_filepath.exists():
                return False

        # Returns True if all the reference inputs and outputs are cached.
        return True

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