import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
import torch


def fp8_gemm_highprec(a_fp8: torch.Tensor,
                           b_fp8: torch.Tensor,
                           a_scale: torch.Tensor,
                           b_scale: torch.Tensor) -> torch.Tensor:
    """
    FP8→FP32 decode, FP32 GEMM with per-128 scaling, final to BF16.
    All ops stay on device; no huge repeat_interleave.
    """
    # assert a_fp8.is_cuda and b_fp8.is_cuda

    M, K = a_fp8.shape
    N    = b_fp8.shape[0]
    KS   = K // 128
    SN   = (N + 127) // 128

    # 1) decode to FP32
    a_f32 = a_fp8.to(torch.float32)
    b_f32 = b_fp8.to(torch.float32)

    # 2a) apply row‐scale to A:
    #    view A as (M, KS, 128), then broadcast‐mul with a_scale (M,KS,1)
    a_f32 = a_f32.view(M, KS, 128)
    a_f32.mul_( a_scale.unsqueeze(2) )   # (M,KS,1) broadcasts over the 128 axis
    a_f32 = a_f32.reshape(M, K)

    # 2b) apply col‐scale to B:
    #    view B as (N, KS, 128), pick each row's block index to lookup b_scale
    b_f32 = b_f32.view(N, KS, 128)
    # block index per row: 0 for rows 0–127, 1 for 128–255, etc.
    row_blocks = torch.arange(N, device=b_f32.device) // 128  # shape (N,)
    # gather the per‐slice scales for each row: (N,KS)
    bs_rows = b_scale[row_blocks]                            # (N,KS)
    b_f32.mul_( bs_rows.unsqueeze(2) )                       # broadcast over 128
    b_f32 = b_f32.reshape(N, K)

    # 3) full‐precision GEMM, then final cast to BF16
    c_f32 = torch.matmul(a_f32, b_f32.transpose(0,1))
    return c_f32.to(torch.bfloat16)


# In your custom_kernel:
def custom_kernel(data: input_t) -> output_t:
    a_fp8, b_fp8, a_scale, b_scale, c = data
    c.copy_(fp8_gemm_highprec(a_fp8, b_fp8, a_scale, b_scale))
    return c
python eval.py benchmark /tmp/tmpsw1bmjsg