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


@torch.jit.script
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.
    """
    M, K = a_fp8.shape
    N = b_fp8.shape[0]
    KS = K // 128  # K must be divisible by 128

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

    # 2a) Apply row-wise scaling to A using out-of-place ops
    a_f32 = a_f32.view(M, KS, 128) * a_scale.unsqueeze(2)
    a_f32 = a_f32.reshape(M, K)

    # 2b) Apply column-wise scaling to B using out-of-place ops
    b_f32 = b_f32.view(N, KS, 128)
    row_blocks = torch.arange(N, device=b_f32.device) // 128
    bs_rows = b_scale[row_blocks]
    b_f32 = b_f32 * bs_rows.unsqueeze(2)
    b_f32 = b_f32.reshape(N, K)

    # 3) FP32 GEMM and convert to BF16
    c_f32 = torch.matmul(a_f32, b_f32.transpose(0, 1))
    return c_f32.to(torch.bfloat16)


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 test /tmp/tmpu1hy7qnp