import torch
import triton
import triton.language as tl
from triton import Config

from task import input_t, output_t


# fast transpose


@triton.jit
def transpose_kernel(
    M: tl.constexpr,
    N: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    X,
    stride_xm: tl.constexpr,
    stride_xn: tl.constexpr,
    Y,
    stride_ym: tl.constexpr,
    stride_yn: tl.constexpr,
):
    pid_m = tl.program_id(0).to(tl.int64)
    pid_n = tl.program_id(1).to(tl.int64)

    m_exact: tl.constexpr = (M % BLOCK_M) == 0
    n_exact: tl.constexpr = (N % BLOCK_N) == 0

    off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    off_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    mask_m = off_m < M
    mask_n = off_n < N

    if m_exact:
        if n_exact:
            mask = None
            other = None
        else:
            mask = mask_n[None, :]
            other = 0
    else:
        if n_exact:
            mask = mask_m[:, None]
            other = 0
        else:
            mask = mask_m[:, None] & mask_n[None, :]
            other = 0

    X_ptrs = X + off_m[:, None] * stride_xm + off_n[None, :] * stride_xn
    Y_ptrs = Y + off_m[:, None] * stride_ym + off_n[None, :] * stride_yn

    tile = tl.load(X_ptrs, mask=mask, other=other)
    tl.store(Y_ptrs, tile, mask=mask)


def transpose(X, Y, BLOCK_M=128, BLOCK_N=128, num_warps=16) -> None:
    if X.dtype.itemsize == 1:
        X = X.view(torch.int8)
    if Y.dtype.itemsize == 1:
        Y = Y.view(torch.int8)

    # check compatibility:
    assert X.shape == Y.shape
    assert X.dtype == Y.dtype

    M = X.shape[0]
    N = X.shape[1]

    block_num = ((M + BLOCK_M - 1) // BLOCK_M) * ((N + BLOCK_N - 1) // BLOCK_N)
    iter_num = block_num / 304
    if iter_num < 1:
        BLOCK_M //= 2
        BLOCK_N //= 2

    stride_xm = X.stride(0)
    stride_xn = X.stride(1)

    stride_ym = Y.stride(0)
    stride_yn = Y.stride(1)

    grid = ((M + BLOCK_M - 1) // BLOCK_M, (N + BLOCK_N - 1) // BLOCK_N)
    transpose_kernel[grid](
        M,
        N,
        BLOCK_M,
        BLOCK_N,
        X,
        stride_xm,
        stride_xn,
        Y,
        stride_ym,
        stride_yn,
        num_warps=num_warps,
    )

fp8_gemm_configs = [
    Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, "GROUP_SIZE_M": 8, "waves_per_eu": 2}, num_stages=2, num_warps=8),
    Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, "GROUP_SIZE_M": 8, "waves_per_eu": 2}, num_stages=2, num_warps=8)
]

# @triton.autotune(configs=fp8_gemm_configs, key=['M', 'N', 'K'])
@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr,
                    GROUP_SIZE_M: tl.constexpr):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    # pid_m = pid // num_pid_n
    # pid_n = pid % num_pid_n

    NUM_XCDS: tl.constexpr = 8
    GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    ## pid remapping on xcds
    # Number of pids per XCD in the new arrangement
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
    # We calculate the number of xcds that have pids_per_xcd pids as
    # tall_xcds
    tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    # Compute current XCD and local pid within the XCD
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    # Calculate new pid based on the new grouping
    # Note that we need to consider the following two cases:
    # 1. the current pid is on a tall xcd
    # 2. the current pid is on a short xcd
    if xcd < tall_xcds:
        pid = xcd * pids_per_xcd + local_pid
    else:
        pid = (
            tall_xcds * pids_per_xcd
            + (xcd - tall_xcds) * (pids_per_xcd - 1)
            + local_pid
        )

    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    k = tl.cdiv(K, BLOCK_SIZE_K)
    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K)
    b_s_step = tl.cdiv(N, BLOCK_SIZE_K)

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += M
        b_s_ptrs += b_s_step
    c = accumulator.to(c_ptr.dtype.element_ty)
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    K = a.size(-1)
    M = a.size(0)
    N = b.size(0)

    ac = torch.empty(a.shape, dtype=a.dtype, device=a.device)
    bc = torch.empty(b.shape, dtype=b.dtype, device=b.device)
    transpose(a, ac)
    transpose(b, bc)
    if M > 1024:
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 64
        BLOCK_SIZE_K = 128
        GROUP_SIZE_M = 8
        num_stages = 2
        num_warps = 8
        waves_per_eu = 2
    else:
        BLOCK_SIZE_M = 128
        BLOCK_SIZE_N = 128
        BLOCK_SIZE_K = 128
        GROUP_SIZE_M = 8
        num_stages = 2
        num_warps = 8
        waves_per_eu = 2

    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
    fp8_gemm_kernel[grid](ac, bc, c, a_scale, b_scale, M, N, K,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=GROUP_SIZE_M,
        num_stages=num_stages,
        num_warps=num_warps,
        waves_per_eu=waves_per_eu,
        kpack=1,
        matrix_instr_nonkdim=16,)

    return c

python eval.py benchmark /tmp/tmpkzo1egdm