from enum import Enum
import functools, sys, ctypes, os, errno
import numpy as np
from functools import partialmethod
from mlir import ir
from mlir.dialects import arith, func, gpu, memref, nvgpu, scf, nvvm
from mlir.extras import types as T
from mlir import runtime as rt
from tools import nvgpucompiler

MLIR_DYNAMIC = -9223372036854775808


def const(value: int, ty=None):
    ty = T.index() if ty is None else ty
    if isinstance(value, ir.Value) and (
        value.type.isinstance(value.type) or T.bool().isinstance(value.type)
    ):
        return value
    return arith.constant(ty, value)


def get_type_size(ty):
    if ir.MemRefType.isinstance(ty):
        size = get_type_size(ty.element_type)
        for sz in ty.shape:
            size *= sz
        return size
    if ir.FloatType.isinstance(ty):
        return ir.FloatType(ty).width // 8
    if ir.IntegerType.isinstance(ty):
        return ir.IntegerType(ty).width // 8
    raise NotImplementedError(ty)


def get_mlir_func_obj_ty(inputArgs):
    args = []
    c_int_p = ctypes.c_int * 1
    c_float_p = ctypes.c_float * 1
    c_bool_p = ctypes.c_bool * 1
    for arg in inputArgs:
        if isinstance(arg, bool):
            args.append(c_bool_p(arg))
        elif isinstance(arg, int):
            args.append(c_int_p(arg))
        elif isinstance(arg, float):
            args.append(c_float_p(arg))
        elif isinstance(arg, np.ndarray):
            args.append(
                ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(arg)))
            )
        else:
            raise NotImplementedError(arg)
    return args


class Mbarriers:
    def __init__(self, number_of_barriers=1):
        self.mbar_ty = ir.Type.parse(
            "!nvgpu.mbarrier.group<memorySpace=#gpu.address_space<workgroup>, num_barriers = "
            + str(number_of_barriers)
            + ">"
        )
        self.mbar_group_op = nvgpu.mbarrier_create(self.mbar_ty)
        self.number_of_barriers = number_of_barriers

    def __getitem__(self, key):
        self.id_op = const(key)
        return self

    def init(self, count: int, predicate=None):
        count_op = const(count)
        if predicate is None:
            nvgpu.mbarrier_init(self.mbar_group_op, count_op, self.id_op)
        else:
            nvgpu.mbarrier_init(
                self.mbar_group_op, count_op, self.id_op, predicate=predicate
            )

    def arrive(self, txcount: int = 0, predicate=None):
        if txcount != 0:
            txcount_op = const(txcount)
            nvgpu.mbarrier_arrive_expect_tx(
                self.mbar_group_op, txcount_op, self.id_op, predicate=predicate
            )
        else:
            nvgpu.mbarrier_arrive(
                ir.Type.parse("!nvgpu.mbarrier.token"), self.mbar_group_op, self.id_op
            )

    def try_wait(self, phase: bool = False, ticks: int = 10000000):
        ticks_op = const(ticks)
        phase_op = const(phase, T.bool())
        nvgpu.MBarrierTryWaitParityOp(
            self.mbar_group_op,
            phase_op,
            ticks_op,
            mbarId=self.id_op,
        )


class TMA:
    """A class that builds a TMA descriptor."""

    def __init__(
        self,
        tma_box_shape,
        memref_ty,
        swizzle=nvgpu.TensorMapSwizzleKind.SWIZZLE_NONE,
        l2promo=nvgpu.TensorMapL2PromoKind.L2PROMO_NONE,
        oob=nvgpu.TensorMapOOBKind.OOB_ZERO,
        interleave=nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE,
    ):
        self.swizzle = swizzle  # mlir.nvgpu.TensorMapSwizzleKind
        self.l2promo = l2promo  # mlir.nvgpu.TensorMapL2PromoKind
        self.oob = oob  # mlir.nvgpu.TensorMapOOBKind
        self.interleave = interleave  # mlir.nvgpu.TensorMapInterleaveKind
        self.tma_box_shape = tma_box_shape
        self.memref_ty = memref_ty  # MemRefType
        self.tma_memref = ir.MemRefType.get(tma_box_shape, memref_ty.element_type)

    @property
    def tensormap_descriptor_ty(self):
        """Returns a tensormap descriptor type."""
        tensorMemrefType = ir.MemRefType.get(
            self.tma_box_shape,
            self.memref_ty.element_type,
            memory_space=ir.Attribute.parse("3"),
        )
        return nvgpu.TensorMapDescriptorType.get(
            tensorMemrefType,
            self.swizzle,
            self.l2promo,
            self.oob,
            self.interleave,
        )

    def create_descriptor(self, device_ptr):
        tma_descriptor_ty = self.tensormap_descriptor_ty
        device_unranked_memref = memref.CastOp(
            ir.UnrankedMemRefType.get(
                self.memref_ty.element_type, self.memref_ty.memory_space
            ),
            device_ptr,
        )
        self.tma_descriptor = nvgpu.TmaCreateDescriptorOp(
            tma_descriptor_ty, device_unranked_memref, map(const, self.tma_box_shape)
        )
        return self.tma_descriptor.result

    def prefetch(self, predicate=None):
        nvgpu.tma_prefetch_descriptor(self.tma_descriptor, predicate=predicate)

    def load(self, dest, mbarrier: Mbarriers, coords=[0], predicate=None):
        nvgpu.TmaAsyncLoadOp(
            dest,
            mbarrier.mbar_group_op,
            self.tma_descriptor,
            coordinates=map(const, coords),
            mbarId=mbarrier.id_op,
            predicate=predicate,
        )


WARP_GROUP_SIZE = 128  # Number of threads in a warpgroup


class Warpgroup:
    def __init__(self, primary_thread, register_size):
        assert (primary_thread % WARP_GROUP_SIZE) == 0
        tidx = gpu.thread_id(gpu.Dimension.x)
        self.primary_thread = primary_thread
        self.register_size = register_size
        self.is_wg_primary = (tidx % WARP_GROUP_SIZE) == 0
        self.wg_id = tidx / WARP_GROUP_SIZE
        self.is_me = self.wg_id == (primary_thread // WARP_GROUP_SIZE)

    def __enter__(self):
        if_op = scf.IfOp(self.is_me)
        self.ipoint_op = ir.InsertionPoint(if_op.then_block)
        self.ipoint_op.__enter__()
        if self.register_size < 64:
            nvvm.setmaxregister(self.register_size, nvvm.SetMaxRegisterAction.decrease)
        else:
            nvvm.setmaxregister(self.register_size, nvvm.SetMaxRegisterAction.increase)

    def __exit__(self, exc_type, exc_value, traceback):
        scf.yield_([])
        self.ipoint_op.__exit__(exc_type, exc_value, traceback)
        return True


class WGMMAType(Enum):
    Accumulator = 1
    Descriptor = 2


class WGMMAMatrix:
    def __init__(
        self,
        matrix_type: WGMMAType,
        shape: list = None,
        desc: TMA = None,
        smem=None,
        ty=None,
        acc_op=None,
    ):
        if acc_op is None:
            self.M = shape[0]
            self.N = shape[1]
            self.ty = ty
            self.matrix_type = matrix_type
            self.desc = desc
            self.smem = smem
            if matrix_type is WGMMAType.Accumulator:
                self.acc_op = nvgpu.warpgroup_mma_init_accumulator(self.acc_ty)
        elif acc_op:
            self.acc_op = acc_op
            self.matrix_type = WGMMAType.Accumulator

    @property
    def acc_ty(self):
        parse_str = f"!nvgpu.warpgroup.accumulator<fragmented=vector<{self.M}x{self.N}x{self.ty}>>"
        return ir.Type.parse(parse_str)

    @property
    def wgmma_ty(self):
        parse_str = f"!nvgpu.warpgroup.descriptor<tensor=memref<{self.M}x{self.N}x{self.desc.memref_ty.element_type}, #gpu.address_space<workgroup>>>"
        return ir.Type.parse(parse_str)

    def store_accumulator(self, dest):
        assert self.matrix_type == WGMMAType.Accumulator
        nvgpu.warpgroup_mma_store(self.acc_op, dest)

    def update_smem(self, smem):
        self.smem = smem

    def update_accumulator(self, acc_op):
        self.acc_op = acc_op

    def __matmul__(self, rhs):
        lhs = nvgpu.warpgroup_generate_descriptor(
            self.wgmma_ty, self.smem, self.desc.tma_descriptor
        )
        rhs = nvgpu.warpgroup_generate_descriptor(
            rhs.wgmma_ty, rhs.smem, rhs.desc.tma_descriptor
        )
        return [lhs, rhs]

    def __iadd__(self, matmulResult):
        lhs = matmulResult[0]
        rhs = matmulResult[1]
        acc_op = nvgpu.WarpgroupMmaOp(
            self.acc_op.type, lhs, rhs, self.acc_op, transposeB=True
        )
        return WGMMAMatrix(WGMMAType.Accumulator, acc_op=acc_op)


def get_dynamic_shared_memory(shape=None, ty=None, offset: int = 0):
    smem_space_str = "#gpu.address_space<workgroup>"
    smem_space = ir.Attribute.parse(smem_space_str)
    dynamic_smem = gpu.dynamic_shared_memory(
        ir.MemRefType.get((MLIR_DYNAMIC,), T.i8(), memory_space=smem_space)
    )
    if shape is None:
        return dynamic_smem
    memref_ty = ir.MemRefType.get(shape, ty, memory_space=smem_space)
    return memref.view(
        ir.MemRefType.get(
            memref_ty.shape, memref_ty.element_type, memory_space=smem_space
        ),
        dynamic_smem,
        const(offset),
        [],
    )


def get_mlir_ty(arg):
    def get_mlir_ty_from_np(dtype):
        if dtype == np.float16:
            return T.f16()
        if dtype == np.float32:
            return T.f32()
        if dtype == np.float64:
            return T.f64()
        if dtype == np.int32:
            return T.i32()
        if dtype == np.int64:
            return T.i64()
        raise NotImplementedError(dtype)

    if isinstance(arg, bool):
        return T.bool()
    elif isinstance(arg, int):
        return T.index()
    elif isinstance(arg, float):
        return T.f32()
    elif isinstance(arg, np.ndarray):
        descriptor = rt.get_ranked_memref_descriptor(arg)
        dtype = get_mlir_ty_from_np(arg.dtype)
        shape = descriptor.shape
        return memref.MemRefType.get(shape, dtype)
    raise NotImplementedError(arg)


class NVDSL:
    @staticmethod
    def mlir_gpu_launch(grid=(1, 1, 1), block=(1, 1, 1), smem=0):
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                launch_op = gpu.LaunchOp(
                    None,
                    [],
                    *map(const, grid),
                    *map(const, block),
                    dynamicSharedMemorySize=arith.constant(T.i32(), smem),
                )
                launch_op.body.blocks.append(*([T.index()] * 12))
                with ir.InsertionPoint(launch_op.body.blocks[0]):
                    result = func(*args, **kwargs)
                    gpu.terminator()
                    return result

            return wrapper

        return decorator

    @staticmethod
    def mlir_func(funcBody):
        @functools.wraps(funcBody)
        def wrapper(*args, **kwargs):
            function_name = funcBody.__name__

            def saveIR(module):
                """Save generated IR"""
                if True:  # self.saveIR:
                    # print(mlir_nvgpu_module)
                    original_stdout = sys.stdout
                    with open("nvdsl.mlir", "w") as f:
                        sys.stdout = f
                        print(module)
                        sys.stdout = original_stdout

            def _binary_op(lhs, rhs, op: str, predAtt="") -> "ArithValue":
                """Generate MLIR's Arith dialects binary operations."""
                rhs = const(rhs)
                if arith._is_float_type(lhs.type) and arith._is_float_type(rhs.type):
                    op += "F"
                    if op.startswith("Cmp"):
                        predicateAttr = getattr(arith, f"CmpFPredicate").__dict__[
                            predAtt
                        ]
                elif arith._is_integer_like_type(
                    lhs.type
                ) and arith._is_integer_like_type(lhs.type):
                    if op == "Div" or op == "Rem":
                        op += "U"
                    op += "I"
                    if op.startswith("Cmp"):
                        predicateAttr = getattr(arith, f"CmpIPredicate").__dict__[
                            predAtt
                        ]
                else:
                    raise NotImplementedError(
                        f"Unsupported '{op}' operands: {lhs}, {rhs}"
                    )

                if op.startswith("Cmp"):
                    op = getattr(arith, f"{op}Op")

                    return op(predicateAttr, lhs, rhs).result
                else:
                    op = getattr(arith, f"{op}Op")
                    return op(lhs, rhs).result

            @ir.register_value_caster(ir.IndexType.static_typeid)
            @ir.register_value_caster(ir.F32Type.static_typeid)
            @ir.register_value_caster(ir.F16Type.static_typeid)
            @ir.register_value_caster(ir.F64Type.static_typeid)
            @ir.register_value_caster(ir.IntegerType.static_typeid)
            class ArithValue(ir.Value):
                """Overloads operators for MLIR's Arith dialects binary operations."""

                def __init__(self, v):
                    super().__init__(v)

                __add__ = partialmethod(_binary_op, op="Add")
                __sub__ = partialmethod(_binary_op, op="Sub")
                __mul__ = partialmethod(_binary_op, op="Mul")
                __truediv__ = partialmethod(_binary_op, op="Div")
                __mod__ = partialmethod(_binary_op, op="Rem")
                __xor__ = partialmethod(_binary_op, op="XOr")
                __lt__ = partialmethod(_binary_op, op="Cmp", predAtt="ult")
                __le__ = partialmethod(_binary_op, op="Cmp", predAtt="ule")
                __eq__ = partialmethod(_binary_op, op="Cmp", predAtt="eq")
                __ne__ = partialmethod(_binary_op, op="Cmp", predAtt="ne")
                __gt__ = partialmethod(_binary_op, op="Cmp", predAtt="ugt")
                __ge__ = partialmethod(_binary_op, op="Cmp", predAtt="uge")
                __and__ = partialmethod(_binary_op, op="And")
                __or__ = partialmethod(_binary_op, op="Or")

                def __str__(self):
                    return (
                        super()
                        .__str__()
                        .replace(ir.Value.__name__, ArithValue.__name__)
                    )

            # Generate MLIR Context and start generating IR
            with ir.Context(), ir.Location.unknown():
                types = []
                for arg in args:
                    types.append(get_mlir_ty(arg))

                # Build IR
                module = ir.Module.create()
                with ir.InsertionPoint(module.body):
                    fop = func.FuncOp(function_name, (types, []))
                    fop.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
                    with ir.InsertionPoint(fop.add_entry_block()):
                        fargs = []
                        for i, a in enumerate(types):
                            fargs.append(fop.arguments[i])

                        # Call user function body
                        result = funcBody(*fargs, **kwargs)
                        func.ReturnOp([])

                # Save IR in a file
                # saveIR(module)

                # Verify the module
                module.operation.verify()

                # Compile and JIT MLIR module
                options = f"cubin-chip=sm_90a cubin-features=+ptx80 opt-level=3"
                support_lib = os.getenv("SUPPORT_LIB")
                if not os.path.exists(support_lib):
                    raise FileNotFoundError(
                        errno.ENOENT, os.strerror(errno.ENOENT), support_lib
                    )
                compiler = nvgpucompiler.NvgpuCompiler(
                    options, opt_level=3, shared_libs=[support_lib]
                )
                engine = compiler.compile_and_jit(module)

            # Convert input arguments to MLIR arguments
            newArgs = get_mlir_func_obj_ty(args)

            # Run the compiled program
            engine.invoke(function_name, *newArgs)

            return result

        return wrapper