import torch
import triton
import triton.language as tl
# import triton.profiler as proton
from task import input_t, output_t
# from ref import ref_kernel, generate_input
DEVICE = triton.runtime.driver.active.get_active_torch_device()

@triton.jit
def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS):
    group_id = tile_id // 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 + (tile_id % group_size_m)
    pid_n = (tile_id % num_pid_in_group) // group_size_m
    return pid_m, pid_n

@triton.jit
def swizzle_tile_v2(pid, M, N, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, GROUP_SIZE_M: tl.constexpr):
    grid_m     = tl.cdiv(M, BLOCK_SIZE_M)
    grid_n     = tl.cdiv(N, BLOCK_SIZE_N)
    width      = GROUP_SIZE_M * grid_m
    group_id   = pid // width
    group_size = tl.minimum(grid_n - group_id * GROUP_SIZE_M, GROUP_SIZE_M)
    pid_n      = group_id * GROUP_SIZE_M + (pid % group_size)
    pid_m      = (pid % width) // group_size
    return pid_m, pid_n

@triton.jit
def fp8_gemm_kernel(
    # tensor pointers and parameters
    a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_ascale_m, stride_ascale_k,
    stride_bscale_n, stride_bscale_k,
    stride_cm, stride_cn,
    # Triton config
    BLOCK_SIZE_M: tl.constexpr, 
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    # tiling arrangement
    pid = tl.program_id(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_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
    # pid_m, pid_n = _compute_pid(pid, num_pid_in_group, num_pid_m, 1, 1)
    pid_m, pid_n = swizzle_tile_v2(pid, M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M)
    
    # offt
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M))
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk)
    # # scale ptr
    a_scale_ptrs = a_scale_ptr + (offs_am[:, None] * stride_ascale_m)
    b_scale_ptrs = b_scale_ptr + pid_n[:, None] * stride_bscale_n 
    # accumulator
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # K dimension iteration
    for k_block in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k_block * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[None, :] < K - k_block * BLOCK_SIZE_K, other=0.0)
        # a_scale = tl.load(a_scale_ptrs).broadcast_to(BLOCK_SIZE_M, BLOCK_SIZE_K)
        a_scale = tl.load(a_scale_ptrs)
        b_scale = tl.load(b_scale_ptrs)[:,None]
        acc += tl.dot(a, b.T)*a_scale*b_scale
        # acc = tl.dot_scaled(a, a_scale, 'bf16', b, b_scale, 'bf16', acc)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
        a_scale_ptrs += stride_ascale_k
        b_scale_ptrs += stride_bscale_k

    c = acc.to(tl.bfloat16)
    offs_cm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M))
    offs_cn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
    c_ptrs = c_ptr + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn)
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

# @triton.jit
# def fp8_gemm_kernel_kparallel(
#     # tensor pointers and parameters
#     a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr,
#     M, N, K,
#     stride_am, stride_ak,
#     stride_bn, stride_bk,
#     stride_ascale_m, stride_ascale_k,
#     stride_bscale_n, stride_bscale_k,
#     stride_cm, stride_cn,
#     # Triton config
#     BLOCK_SIZE_M: tl.constexpr, 
#     BLOCK_SIZE_N: tl.constexpr,
#     BLOCK_SIZE_K: tl.constexpr,
#     SPLIT_K: tl.constexpr,
# ):
#     # tiling arrangment
#     pid = tl.program_id(0)
#     pid_k = tl.program_id(1)
#     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
#     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
#     num_pid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
#     pid_m = pid // num_pid_n
#     pid_n = pid % num_pid_n
    
#     # offt
#     offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M))
#     offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
#     offs_k = (tl.arange(0, BLOCK_SIZE_K) + pid_k * BLOCK_SIZE_K)
#     offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M)
#     offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N)
#     offs_am = tl.where(offs_am < M, offs_am, 0)
#     offs_bn = tl.where(offs_bn < N, offs_bn, 0)
#     offs_k = tl.where(offs_k < K, offs_k, 0)
#     a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
#     b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk)
#     # # scale ptr
#     a_scale_ptrs = a_scale_ptr + (offs_am[:, None] * stride_ascale_m) + pid_k * stride_ascale_k
#     b_scale_ptrs = b_scale_ptr + pid_n[:, None] * stride_bscale_n + pid_k * stride_bscale_k
#     # accumulator
#     acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
#     # K dimension iteration
#     for k_block in range(0, num_pid_k):
#         tl.multiple_of(a_ptrs, [BLOCK_SIZE_M, BLOCK_SIZE_K])
#         tl.multiple_of(b_ptrs, [BLOCK_SIZE_N, BLOCK_SIZE_K])
#         a = tl.load(a_ptrs)
#         b = tl.load(b_ptrs)
#         a_scale = tl.load(a_scale_ptrs)
#         b_scale = tl.load(b_scale_ptrs)[:,None]
#         acc += tl.dot(a, b.T)*a_scale*b_scale
#         a_ptrs += BLOCK_SIZE_K * stride_ak
#         b_ptrs += BLOCK_SIZE_K * stride_bk
#         a_scale_ptrs += stride_ascale_k
#         b_scale_ptrs += stride_bscale_k

#     c = acc.to(tl.bfloat16)
#     offs_cm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M))
#     offs_cn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
#     c_ptrs = c_ptr + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn)
#     c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
#     if (SPLIT_K == 1):
#         tl.store(c_ptrs, c, mask=c_mask)
#     else:
#         tl.atomic_add(c_ptrs, c, mask=c_mask)

# @triton.jit
# def fp8_gemm_kernel_persistent(
#     # tensor pointers and parameters
#     a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr,
#     M, N, K,
#     stride_am, stride_ak,
#     stride_bn, stride_bk,
#     stride_ascale_m, stride_ascale_k,
#     stride_bscale_n, stride_bscale_k,
#     stride_cm, stride_cn,
#     # Triton config
#     BLOCK_SIZE_M: tl.constexpr, 
#     BLOCK_SIZE_N: tl.constexpr,
#     BLOCK_SIZE_K: tl.constexpr,
#     GROUP_SIZE_M: tl.constexpr,
#     NUM_SMS: tl.constexpr,
# ):
#     start_pid = tl.program_id(0)
#     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
#     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
#     k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
#     num_tiles = num_pid_m * num_pid_n

#     tile_id_c = start_pid - NUM_SMS
#     offs_k = tl.arange(0, BLOCK_SIZE_K)
#     num_pid_in_group = GROUP_SIZE_M * num_pid_n

#     for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten = True):
#         pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS)
#         start_m = pid_m * BLOCK_SIZE_M
#         start_n = pid_n * BLOCK_SIZE_N
#         # offt
#         offs_am = (start_m + tl.arange(0, BLOCK_SIZE_M))
#         offs_bn = (start_n + tl.arange(0, BLOCK_SIZE_N))
#         offs_am = tl.where(offs_am < M, offs_am, 0)
#         offs_bn = tl.where(offs_bn < N, offs_bn, 0)
#         offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M)
#         offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N)
#         a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
#         b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk)
#         # scale ptr
#         a_scale_ptrs = a_scale_ptr + (offs_am[:, None] * stride_ascale_m)
#         b_scale_ptrs = b_scale_ptr + pid_n[:, None] * stride_bscale_n
#         # accumulator
#         acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
#         for ki in range(k_tiles):
#             # a_scale = tl.load(a_scale_ptrs).broadcast_to(BLOCK_SIZE_M, BLOCK_SIZE_K)
#             tl.multiple_of(a_ptrs, [BLOCK_SIZE_M, BLOCK_SIZE_K])
#             tl.multiple_of(b_ptrs, [BLOCK_SIZE_N, BLOCK_SIZE_K])
#             a = tl.load(a_ptrs)
#             b = tl.load(b_ptrs)
#             a_scale = tl.load(a_scale_ptrs)
#             b_scale = tl.load(b_scale_ptrs)[:,None]
#             acc += tl.dot(a, b.T)*a_scale*b_scale
#             a_ptrs += BLOCK_SIZE_K * stride_ak
#             b_ptrs += BLOCK_SIZE_K * stride_bk
#             a_scale_ptrs += stride_ascale_k
#             b_scale_ptrs += stride_bscale_k
#         c = acc.to(tl.bfloat16)
#         tile_id_c += NUM_SMS
#         pid_m, pid_n = _compute_pid(tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS)
#         offs_cm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M))
#         offs_cn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
#         c_ptrs = c_ptr + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn)
#         c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
#         tl.store(c_ptrs, c, mask=c_mask)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    
    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()
    
    m, k = a.shape
    n = b.shape[0]

    num_warps = 4
    
    grid = lambda meta: (triton.cdiv(m, meta['BLOCK_SIZE_M']) * triton.cdiv(n, meta['BLOCK_SIZE_N']), )
    
    fp8_gemm_kernel[grid](
        a, b, a_scale, b_scale, c,
        m, n, k,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        a_scale.stride(0), a_scale.stride(1),
        b_scale.stride(0), b_scale.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=128,
        BLOCK_SIZE_N=128,
        BLOCK_SIZE_K=128,
        GROUP_SIZE_M=8,
        num_warps=num_warps,
    )
    return c


# def kparallel_kernel(data: input_t) -> output_t:
#     a, b, a_scale, b_scale, c = data
    
#     a = a.contiguous()
#     a_scale = a_scale.contiguous()
#     b_scale = b_scale.contiguous()
    
#     m, k = a.shape
#     n = b.shape[0]
#     split_k = 4

#     num_warps = 8
    
#     grid = lambda meta: (triton.cdiv(m, meta['BLOCK_SIZE_M']) * triton.cdiv(n, meta['BLOCK_SIZE_N']), split_k)
    
#     fp8_gemm_kernel_kparallel[grid](
#         a, b, a_scale, b_scale, c,
#         m, n, k,
#         a.stride(0), a.stride(1),
#         b.stride(0), b.stride(1),
#         a_scale.stride(0), a_scale.stride(1),
#         b_scale.stride(0), b_scale.stride(1),
#         c.stride(0), c.stride(1),
#         BLOCK_SIZE_M=BLOCK_SIZE_M,
#         BLOCK_SIZE_N=BLOCK_SIZE_N,
#         BLOCK_SIZE_K=BLOCK_SIZE_K,
#         SPLIT_K=split_k,
#     )
#     return c

# def persistent_kernel(data: input_t) -> output_t:
#     a, b, a_scale, b_scale, c = data
    
#     a = a.contiguous()
#     a_scale = a_scale.contiguous()
#     b_scale = b_scale.contiguous()
    
#     m, k = a.shape
#     n = b.shape[0]

#     num_warps = 4
    

#     NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count
#     grid = lambda meta: (min(NUM_SMS, triton.cdiv(m, meta['BLOCK_SIZE_M']) * triton.cdiv(n, meta['BLOCK_SIZE_N'])), )
#     fp8_gemm_kernel_persistent[grid](
#         a, b, a_scale, b_scale, c,
#         m, n, k,
#         a.stride(0), a.stride(1),
#         b.stride(0), b.stride(1),
#         a_scale.stride(0), a_scale.stride(1),
#         b_scale.stride(0), b_scale.stride(1),
#         c.stride(0), c.stride(1),
#         BLOCK_SIZE_M=BLOCK_SIZE_M,
#         BLOCK_SIZE_N=BLOCK_SIZE_N,
#         BLOCK_SIZE_K=BLOCK_SIZE_K,
#         GROUP_SIZE_M=16,
#         NUM_SMS=NUM_SMS,
#         num_warps=num_warps,
#     )
#     return c

# def test(dut, ref):
#     # test the Triton kernel
#     m, n, k = 256, 256, 256
#     data = generate_input(m, n, k, seed=42)
#     data_ker = dut(data).clone()
#     print("Triton kernel output:", data_ker)
#     data_ref = ref(data)
#     print("Reference kernel output:", data_ref)
#     print("Max difference:", torch.max(torch.abs(data_ker - data_ref)))
#     assert torch.allclose(data_ker, data_ref, atol=1e-2), "Output does not match reference implementation"

# @triton.testing.perf_report(
#     triton.testing.Benchmark(
#         x_names=["M"],
#         x_vals=[2**i for i in range(7, 13)],
#         line_arg = "provider",
#         line_vals = ["custom", "kparallel", "persistent", "torch"],
#         line_names = ["Custom", "Kparallel", "Persistent", "Torch"],
#         # styles=[("green", "solid"), ("blue", "dashed")],
#         styles = [("green", "solid"), ("blue", "dashed"), ("red", "dotted"), ("black", "dotted")],
#         ylabel="runtime (ms)",
#         plot_name="FP8 GEMM Performance",
#         args={},
#     )
# )
# def benchmark(M, provider):
#     N,K=M,M
#     data = generate_input(M, N, K, seed=42)
#     quantiles = [0.5, 0.2, 0.8]
#     if provider == "custom":
#         # Triton kernel
#         ms, min_ms, max_ms = triton.testing.do_bench(
#             lambda: custom_kernel(data),
#             quantiles=quantiles
#         )
#     if provider == "kparallel":
#         # Triton kernel
#         ms, min_ms, max_ms = triton.testing.do_bench(
#             lambda: kparallel_kernel(data),
#             quantiles=quantiles
#         )
#     if provider == "persistent":
#         # Triton kernel
#         ms, min_ms, max_ms = triton.testing.do_bench(
#             lambda: persistent_kernel(data),
#             quantiles=quantiles
#         )
#     if provider == "torch":
#         # Reference kernel
#         ms, min_ms, max_ms = triton.testing.do_bench(
#             lambda: ref_kernel(data),
#             quantiles=quantiles
#         )
#     return ms, min_ms, max_ms

# def run_single_kernel(kernel=custom_kernel):
#     # Run a single kernel to warm up the GPU
#     m, n, k = 2048, 2048, 2048
#     data = generate_input(m, n, k, seed=4242)
#     kernel(data)

# if __name__ == "__main__":
#     # test(custom_kernel, ref_kernel)
#     benchmark.run(show_plots=True, print_data=True)
#     # test(kparallel_kernel, ref_kernel)
#     # run_single_kernel(persistent_kernel)
#     # run_single_kernel()
python eval.py test /tmp/tmp12bh0wko