"""FP8 e4m3 GEMM for B200 (SM100).

y = (x @ weight.T) * weight_scale, returned as bf16.
  x:            fp8_e4m3fn (M, K)
  weight:       fp8_e4m3fn (N, K)   (per-output-channel normalized)
  weight_scale: fp32       (N,)     (per-output-channel dequant scale)

Real Blackwell fp8 tensor-core path: fp8 x fp8 tl.dot, fp32 accumulate,
per-channel scale fused into the epilogue, bf16 output. K is padded to a
multiple of BLOCK_K so the main loop is fully unmasked (a masked K-tail
pessimizes the whole pipeline on Blackwell).
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl

BK_PAD = 128


def _configs():
    cfgs = []
    for bm, bn, bk, w, s in [
        (128, 256, 128, 8, 4),
        (128, 256, 128, 8, 3),
        (256, 128, 128, 8, 4),
        (128, 128, 128, 8, 4),
        (128, 128, 128, 8, 3),
        (64, 256, 128, 8, 4),
        (64, 128, 128, 4, 4),
        (64, 64, 128, 4, 4),
        (32, 128, 128, 4, 4),
        (32, 256, 128, 4, 4),
    ]:
        cfgs.append(triton.Config(
            {"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8},
            num_warps=w, num_stages=s,
        ))
    return cfgs


@triton.autotune(configs=_configs(), key=["M", "N", "K"])
@triton.jit
def _fp8_gemm_kernel(
    x_ptr, w_ptr, s_ptr, y_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wn, stride_wk,
    stride_ym, stride_yn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
    offs_k = tl.arange(0, BLOCK_K)

    x_ptrs = x_ptr + (offs_am[:, None] * stride_xm + offs_k[None, :] * stride_xk)
    w_ptrs = w_ptr + (offs_bn[None, :] * stride_wn + offs_k[:, None] * stride_wk)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    # K is padded to a multiple of BLOCK_K by the host -> no masked tail.
    for _ in range(0, K, BLOCK_K):
        a = tl.load(x_ptrs)
        b = tl.load(w_ptrs)
        acc = tl.dot(a, b, acc, out_dtype=tl.float32)
        x_ptrs += BLOCK_K * stride_xk
        w_ptrs += BLOCK_K * stride_wk

    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    scale = tl.load(s_ptr + offs_cn, mask=offs_cn < N, other=0.0).to(tl.float32)
    acc = acc * scale[None, :]

    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    y_ptrs = y_ptr + offs_cm[:, None] * stride_ym + offs_cn[None, :] * stride_yn
    mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(y_ptrs, acc.to(tl.bfloat16), mask=mask)


def _pad_k(t: torch.Tensor, Kp: int) -> torch.Tensor:
    M, K = t.shape
    out = t.new_zeros((M, Kp))
    out[:, :K] = t
    return out


def _fp8_gemm(x: torch.Tensor, w: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
    M, K = x.shape
    N, Kw = w.shape
    assert K == Kw
    Kp = ((K + BK_PAD - 1) // BK_PAD) * BK_PAD
    if Kp != K:
        x = _pad_k(x, Kp)
        w = _pad_k(w, Kp)
    y = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
    _fp8_gemm_kernel[grid](
        x, w, s, y,
        M, N, Kp,
        x.stride(0), x.stride(1),
        w.stride(0), w.stride(1),
        y.stride(0), y.stride(1),
    )
    return y


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        self.register_buffer("weight", torch.zeros(N, K, dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.zeros(N, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _fp8_gemm(x, self.weight, self.weight_scale)


M = 4096
N = 4096
K = 4096


def get_inputs():
    x = (torch.rand(M, K) * 8 - 4).to(torch.float8_e4m3fn)
    return [x]


def get_init_inputs():
    return [M, N, K]