#!POPCORN leaderboard amd-fp8-mm
from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline
import time
import os
import sys

if "PYTORCH_ROCM_ARCH" not in os.environ:
    os.environ["PYTORCH_ROCM_ARCH"] = "gfx942:xnack-"

def ref_kernel(data: input_t) -> output_t:
    """
    Highly inefficient torch reference implementation of FP8 GEMM.
    You can use this as a reference / starting template for your implementation.
    """
    # c: [m, n] is pre-allocated memory to help remove allocation overhead.
    a, b, a_scale, b_scale, c = data

    # a is M x K in column-major order, we convert here for simplicity.
    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    # constants
    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    block_shape_n = 128
    block_shape_k = 128
    scale_n = b_scale.shape[0]
    scale_k = b_scale.shape[1]

    # Apply blockwise scaling to input 'a'
    a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)  # Shape: [m, scale_k, block_shape_k]
    a_scale = a_scale.reshape(m, scale_k * block_shape_k) 
    a_scale = a_scale[:, :k]

    # Dequantize 'a', in your implementation you should do this at the end.
    a = a.to(a_scale.dtype) * a_scale 

    # Apply blockwise scaling to input 'b'
    b_scale = (
        b_scale.view(-1, 1)
        .repeat(1, block_shape_n * block_shape_k)
        .view(scale_n, scale_k, block_shape_n, block_shape_k)
        .permute(0, 2, 1, 3)  # Reorder dimensions: [scale_n, blk_n, scale_k, blk_k]
        .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
    )
    b_scale = b_scale[:n, :k]

    # Dequantize 'b', in your implementation you should do this at the end.
    b = b.to(b_scale.dtype) * b_scale 

    # Compute FP8 GEMM and write to 'c'. 
    c[...] = (a @ b.T).to(torch.bfloat16)
    return c


def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // 128],
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128],
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    ref_compiled = torch.compile(ref_kernel, mode="reduce-overhead")
    return ref_compiled(data)

# if __name__ == "__main__":
#     a = torch.randn(1024, 1024, dtype=torch.float8_e4m3fnuz)
#     b = torch.randn(1024, 1024, dtype=torch.float8_e4m3fnuz)
#     a_scale = torch.randn(1024, 1024 // 128, dtype=torch.float32)
#     b_scale = torch.randn(1024 // 128, 1024 // 128, dtype=torch.float32)
#     c = torch.randn(1024, 1024, dtype=torch.bfloat16)
#     custom_kernel((a, b, a_scale, b_scale, c))
python eval.py leaderboard /tmp/tmp23k9660r
skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 56, in ref_kernel
    c[...] = (a @ b.T).to(torch.bfloat16)

