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

from task import input_t, output_t

# best version 3

# fp8_gemm_configs = [
#     Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128, "GROUP_SIZE_M": 4}, num_stages=num_stages, num_warps=8)
#     for block_m in [16, 32, 64] for block_n in [32, 64, 128] for num_stages in [2, 3, 4, 5]
# ]

# @triton.autotune(configs=fp8_gemm_configs, key=['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) * 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 += 1
    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
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.t().contiguous()
    b_scale = b_scale.contiguous()
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)

    if M < 2048:
        BLOCK_SIZE_M = 32
    else:
        BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 128
    GROUP_SIZE_M = 16
    num_stages = 2
    num_warps = 8
    waves_per_eu = 2

    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)
    fp8_gemm_kernel_begin_time = time.perf_counter_ns()
    torch.cuda.synchronize()
    fp8_gemm_kernel[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,)
    torch.cuda.synchronize()
    fp8_gemm_kernel_end_time = time.perf_counter_ns()

    test_m = M // 64
    test_n = N // 64
    test_k = K // 64
    if test_m == 96 and test_n == 72 and test_k == 112:
        print(fp8_gemm_kernel_end_time - fp8_gemm_kernel_begin_time)
    # bm = fp8_gemm_kernel.best_config.kwargs["BLOCK_SIZE_M"]
    # bn = fp8_gemm_kernel.best_config.kwargs["BLOCK_SIZE_N"]
    # ns = fp8_gemm_kernel.best_config.num_stages
    # m = M / 64
    # n = N / 64
    # k = K / 64
    # if (m != 16):
    #     print(f"{N},{K},{bm},{bn},{ns}")
    return c

python eval.py benchmark /tmp/tmpihtkm0p8
226646917
2030228
1756783
1737317
1703295
1713330
1693333
1691448
1704495
1715747
1677548
1706433
1697236
1703805
1693406
1685113
1716172
1677803
1683589
1718968
1715963
1699991
1720941
1741365
1744761
1690639
1684896
1700377
1715691
1678650
1718553
1731001
1730269
1680997
1679790
1703677
1714099
1705849
1720672
1717354
1722932
1688927
1700770
1708896
1744812
1707237
1731958
1749710
1713511
1698833
1697377
1705649
1726810
1677667
1706148
1759926
1723521
1692344
1699202
1714963
1726693
1690379
1689688
1746483
1731352
1688634
1682428
1716916
1741802
1697667
1692484
1700071
1733605
1695615
1700583
1737013
1721483
1699038
1681730
1747466
1729732
1739336
1724923
1740114
1745326
1707652
1729008
1733473
1748454
1726331
1727264
1721013
1735395
1739359
1707671
1712699
1712958
1740379
1712351
1706432
1721479
