"""FP8 e4m3 x e4m3 GEMM for H100 PCIe (SM90 Hopper).

y = (x @ w.T) * weight_scale   (output bf16)
  x:          fp8_e4m3 (M, K)
  w:          fp8_e4m3 (N, K)   [already quantized into the e4m3 range]
  weight_scale: fp32 (N,)       [per-output-channel dequant scale]

The MMA runs as fp8 x fp8 -> fp32 on Hopper tensor cores; the per-channel
scale is folded into the epilogue. This is mathematically equivalent to the
bf16-matmul reference (fp8->bf16 is exact) up to accumulation-order rounding,
well within the 0.2 atol/rtol tolerance.
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl


@triton.jit
def _fp8_gemm_kernel(
    A, B, C, S,
    M, N, K,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(axis=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)
    offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    # A (M,K) row-major: A[m,k]; B (N,K) row-major, used transposed -> B[k,n]=w[n,k]
    A_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    B_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        k_rem = K - k * BLOCK_K
        a = tl.load(A_ptrs, mask=offs_k[None, :] < k_rem, other=0.0)
        b = tl.load(B_ptrs, mask=offs_k[:, None] < k_rem, other=0.0)
        accumulator = tl.dot(a, b, accumulator)
        A_ptrs += BLOCK_K * stride_ak
        B_ptrs += BLOCK_K * stride_bk

    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    C_ptrs = C + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn)
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    scale = tl.load(S + offs_cn, mask=offs_cn < N, other=0.0)
    acc = accumulator * scale[None, :]
    tl.store(C_ptrs, acc.to(tl.bfloat16), mask=c_mask)


def _select_config(M: int, N: int, K: int) -> dict:
    """Hand-picked Hopper fp8 tile configs (refined via benchmarking)."""
    if M <= 64:
        # Skinny / decode: memory-bound; keep M tile small, N tile moderate.
        return dict(BLOCK_M=32, BLOCK_N=128, BLOCK_K=128, GROUP_M=8,
                    num_warps=4, num_stages=3)
    # Big-M compute-bound shapes.
    return dict(BLOCK_M=128, BLOCK_N=256, BLOCK_K=128, GROUP_M=8,
                num_warps=8, num_stages=3)


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        # Must match reference exactly (strict load_state_dict): weight (N,K) fp8,
        # weight_scale (N,) fp32.
        self.register_buffer("weight", torch.empty(N, K, dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.empty(N, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M = x.shape[0]
        N, K = self.N, self.K
        assert x.dtype == torch.float8_e4m3fn and x.is_contiguous()
        w = self.weight
        assert w.dtype == torch.float8_e4m3fn and w.is_contiguous()
        y = torch.empty(M, N, dtype=torch.bfloat16, device=x.device)

        cfg = _select_config(M, N, K)
        BM, BN, BK = cfg["BLOCK_M"], cfg["BLOCK_N"], cfg["BLOCK_K"]
        grid = (triton.cdiv(M, BM) * triton.cdiv(N, BN),)
        _fp8_gemm_kernel[grid](
            x, w, y, self.weight_scale,
            M, N, K,
            x.stride(0), x.stride(1),
            w.stride(0), w.stride(1),
            y.stride(0), y.stride(1),
            BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK, GROUP_M=cfg["GROUP_M"],
            num_warps=cfg["num_warps"], num_stages=cfg["num_stages"],
        )
        return y