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, stride_xn,
    Y, stride_ym, stride_yn,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    # 构建 block ptr
    x_block_ptr = tl.make_block_ptr(
        base=X,
        shape=(M, N),
        strides=(stride_xm, stride_xn),
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(0, 1)
    )

    y_block_ptr = tl.make_block_ptr(
        base=Y,
        shape=(N, M),
        strides=(stride_yn, stride_ym),
        offsets=(pid_n * BLOCK_N, pid_m * BLOCK_M),
        block_shape=(BLOCK_N, BLOCK_M),
        order=(0, 1)
    )
    # y_block_ptr = tl.make_block_ptr(
    #     base=Y,
    #     shape=(M, N),
    #     strides=(stride_ym, stride_yn),
    #     offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
    #     block_shape=(BLOCK_M, BLOCK_N),
    #     order=(1, 0)
    # )

    tile = tl.load(x_block_ptr)
    tile = tl.trans(tile)
    tl.store(y_block_ptr, tile)


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,
    )

@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_XCDS: tl.constexpr = 8
    # GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    # pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    # tall_xcds = GRID_MN % NUM_XCDS
    # tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    # xcd = pid % NUM_XCDS
    # local_pid = pid // NUM_XCDS
    # 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_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    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
    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)

    a_block_ptr = tl.make_block_ptr(
        base=a_ptr,
        shape=(M, K),
        strides=(K, 1),
        offsets=(pid_m * BLOCK_SIZE_M, 0),
        block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K),
        order=(1, 0)
    )
    b_block_ptr = tl.make_block_ptr(
        base=b_ptr,
        shape=(K, N),
        strides=(1, K),
        offsets=(0, pid_n * BLOCK_SIZE_N),
        block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N),
        order=(0, 1)
    )

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_block_ptr)
        b = tl.load(b_block_ptr)
        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_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K))
        b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))
        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)


@triton.jit
def fp8_gemm_kernel_with_trans(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)

    # # 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_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    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
    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)

    a_block_ptr = tl.make_block_ptr(
        base=a_ptr,
        shape=(M, K),
        strides=(1, M),
        offsets=(pid_m * BLOCK_SIZE_M, 0),
        block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K),
        order=(0, 1)
    )
    b_block_ptr = tl.make_block_ptr(
        base=b_ptr,
        shape=(K, N),
        strides=(N, 1),
        offsets=(0, pid_n * BLOCK_SIZE_N),
        block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N),
        order=(1, 0)
    )

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_block_ptr)
        b = tl.load(b_block_ptr)
        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_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K))
        b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))
        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)
 
    need_transpose = True
    BLOCK_SIZE_K = 128
    if M <= 1024 and N < 1024:
        BLOCK_SIZE_M = 32
        BLOCK_SIZE_N = 32
        GROUP_SIZE_M = 4
        num_stages = 3
        num_warps = 1
        waves_per_eu = 16
        # fused trans
        need_transpose = False
    elif M <= 1024 and N < 2048:
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 32
        GROUP_SIZE_M = 8
        num_stages = 2
        num_warps = 1
        waves_per_eu = 1
        need_transpose = False
    elif M <= 1024 and N < 4096:
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 64
        GROUP_SIZE_M = 1
        num_stages = 2 # K?
        num_warps = 1
        waves_per_eu = 1
        # fused trans
        need_transpose = False
    elif M <= 1024 and N < 6144 and K < 4096:
        BLOCK_SIZE_M = 128
        BLOCK_SIZE_N = 64
        GROUP_SIZE_M = 1
        num_stages = 1 
        num_warps = 2
        waves_per_eu = 1
        # fused trans
        need_transpose = False
    elif M <= 1024 and N < 6144 and K >= 4096:    
        BLOCK_SIZE_M = 128
        BLOCK_SIZE_N = 64
        GROUP_SIZE_M = 8
        num_stages = 2
        num_warps = 8
        waves_per_eu = 2
    elif M <= 1024:
        BLOCK_SIZE_M = 128
        BLOCK_SIZE_N = 64
        GROUP_SIZE_M = 8
        num_stages = 1
        num_warps = 2
        waves_per_eu = 2
        # fused trans
        need_transpose = False
    elif K < 1024:
        BLOCK_SIZE_M = 128
        BLOCK_SIZE_N = 64
        GROUP_SIZE_M = 1
        num_stages = 1
        num_warps = 2
        waves_per_eu = 2
        need_transpose = False
    elif K < 2048:
        BLOCK_SIZE_M = 256
        BLOCK_SIZE_N = 128
        GROUP_SIZE_M = 8
        num_stages = 2
        num_warps = 4
        waves_per_eu = 16
        need_transpose = False
    else:
        BLOCK_SIZE_M = 128
        BLOCK_SIZE_N = 64
        GROUP_SIZE_M = 8
        num_stages = 2
        num_warps = 8
        waves_per_eu = 2
    
    if need_transpose:
        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)
        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=2,
            matrix_instr_nonkdim=16,)
    else:
        grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
        fp8_gemm_kernel_with_trans[grid](a, b, 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=2,
            matrix_instr_nonkdim=16,)
    return c

python eval.py leaderboard /tmp/tmptjag2oeb