"""FP8 e4m3 GEMM solution: y = ((x @ w.T) * weight_scale).to(bf16).

Custom Triton kernels running real fp8 x fp8 WGMMA tensor-core MMAs
(fp8_e4m3 inputs, fp32 accumulate) on Hopper SM90:

  * Large / aligned shapes: persistent TMA kernel with grouped tile ordering;
    per-channel scale applied in the epilogue.
  * Odd K (row stride not 16-byte aligned, which TMA requires): fuse a
    zero-padding copy into aligned buffers, then run the same TMA kernel.
  * Skinny M (decode-like, memory-bound): transposed-operand split-K kernel
    (w as the M-tiled operand so WGMMA gets full m64 tiles) with fp32 atomic
    accumulation into a workspace plus a tiny scale+convert epilogue kernel.
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl
from triton.tools.tensor_descriptor import TensorDescriptor

E4M3_MAX = 448.0


# ---------------------------------------------------------------------------
# Kernel 1: persistent TMA fp8 GEMM (requires 16-byte-aligned row strides).
# ---------------------------------------------------------------------------
@triton.jit
def _fp8_gemm_tma(
    a_desc, b_desc, c_desc, S,
    M, N, K,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
    GROUP_M: tl.constexpr, NUM_SMS: tl.constexpr,
):
    start_pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BM)
    num_pid_n = tl.cdiv(N, BN)
    k_tiles = tl.cdiv(K, BK)
    num_tiles = num_pid_m * num_pid_n

    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
        num_pid_in_group = GROUP_M * num_pid_n
        group_id = tile_id // 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 + (tile_id % group_size_m)
        pid_n = (tile_id % num_pid_in_group) // group_size_m

        off_m = pid_m * BM
        off_n = pid_n * BN
        acc = tl.zeros((BM, BN), dtype=tl.float32)
        for ki in range(k_tiles):
            a = a_desc.load([off_m, ki * BK])
            b = b_desc.load([off_n, ki * BK])
            acc = tl.dot(a, b.T, acc)

        rn = off_n + tl.arange(0, BN)
        scale = tl.load(S + rn, mask=rn < N, other=0.0)
        acc = acc * scale[None, :]
        c_desc.store([off_m, off_n], acc.to(tl.bfloat16))


# ---------------------------------------------------------------------------
# Kernel 2: zero-padding copy (for K whose row stride is not 16-byte aligned).
# ---------------------------------------------------------------------------
@triton.jit
def _pad_copy(SRC, DST, K, KPAD, BLK: tl.constexpr):
    pid = tl.program_id(0)
    nblk = tl.cdiv(KPAD, BLK)
    row = pid // nblk
    col = (pid % nblk) * BLK + tl.arange(0, BLK)
    v = tl.load(SRC + row * K + col, mask=col < K, other=0.0)
    tl.store(DST + row * KPAD + col, v, mask=col < KPAD)


# ---------------------------------------------------------------------------
# Kernel 3: transposed split-K fp8 GEMM for skinny M (TMA loads).
# Computes y.T tiles: (n tile) x (all m) so the WGMMA M dimension is large.
# ---------------------------------------------------------------------------
@triton.jit
def _fp8_gemm_skinny(
    w_desc, x_desc, WS,
    M, N, K,
    SPLIT_K: tl.constexpr,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_n = tl.cdiv(N, BM)
    pid_n = pid % num_pid_n
    pid_k = pid // num_pid_n

    k_per_split = tl.cdiv(K, SPLIT_K * BK) * BK
    k_start = pid_k * k_per_split
    k_end = min(k_start + k_per_split, K)

    off_n = pid_n * BM
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for k in range(k_start, k_end, BK):
        a = w_desc.load([off_n, k])   # (BM, BK) tile of w
        b = x_desc.load([0, k])       # (BN, BK) tile of x
        acc = tl.dot(a, b.T, acc)

    rn = off_n + tl.arange(0, BM)
    rm = tl.arange(0, BN)
    mask = rn[:, None] < N
    WS_ptr = WS + rn[:, None] * M + rm[None, :]
    tl.atomic_add(WS_ptr, acc, mask=mask, sem="relaxed")


@triton.jit
def _epilogue_skinny(WS, Y, S, M, N, BM: tl.constexpr, BN: tl.constexpr):
    pid = tl.program_id(0)
    num_pid_n = tl.cdiv(N, BN)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    mask = (rm[:, None] < M) & (rn[None, :] < N)
    v = tl.load(WS + rn[None, :] * M + rm[:, None], mask=mask, other=0.0)
    sc = tl.load(S + rn, mask=rn < N, other=0.0)
    tl.store(Y + rm[:, None] * N + rn[None, :], (v * sc[None, :]).to(tl.bfloat16), mask=mask)


def _tma_aligned(t: torch.Tensor) -> bool:
    # TMA requires every stride except the innermost to be a multiple of 16 bytes.
    return t.stride(0) % 16 == 0 and t.data_ptr() % 16 == 0


# Hand-tuned configs on H100 PCIe (BM, BN, BK, GROUP_M, num_warps, num_stages).
_LARGE_CONFIGS = {
    (4096, 4096, 4096): (128, 128, 128, 8, 4, 4),
    (4096, 14336, 4096): (128, 256, 128, 16, 8, 3),
}
# (BM, BN(tile over M), BK, SPLIT_K, num_warps, num_stages)
_SKINNY_CONFIGS = {
    (32, 8192, 8192): (128, 32, 256, 4, 4, 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
        w = torch.empty(N, K, dtype=torch.bfloat16)
        nn.init.normal_(w, std=0.02)
        s = (w.float().abs().amax(dim=1, keepdim=True) / E4M3_MAX).clamp(min=1e-12)
        w_fp8 = (w.float() / s).to(torch.float8_e4m3fn)
        self.register_buffer("weight", w_fp8)
        self.register_buffer("weight_scale", s.squeeze(1).to(torch.float32))
        self._workspace = None
        self._pad_bufs = None
        self._NUM_SMS = None

    def _num_sms(self, device):
        if self._NUM_SMS is None:
            self._NUM_SMS = torch.cuda.get_device_properties(device).multi_processor_count
        return self._NUM_SMS

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

        if M <= 64:
            return self._forward_skinny(x, w, ws, y, M, N, K)

        if _tma_aligned(x) and _tma_aligned(w):
            return self._forward_tma(x, w, ws, y, M, N, K)
        return self._forward_padded(x, w, ws, y, M, N, K)

    # -- large aligned -------------------------------------------------------
    def _large_config(self, M, N, K):
        cfg = _LARGE_CONFIGS.get((M, N, K))
        if cfg is not None:
            return cfg
        if N >= 8192:
            return (128, 256, 128, 16, 8, 3)
        return (128, 128, 128, 8, 4, 4)

    def _forward_tma(self, x, w, ws, y, M, N, K):
        BM, BN, BK, GROUP_M, num_warps, num_stages = self._large_config(M, N, K)
        a_desc = TensorDescriptor.from_tensor(x, [BM, BK])
        b_desc = TensorDescriptor.from_tensor(w, [BN, BK])
        c_desc = TensorDescriptor.from_tensor(y, [BM, BN])
        NUM_SMS = self._num_sms(x.device)
        grid = (min(NUM_SMS, triton.cdiv(M, BM) * triton.cdiv(N, BN)),)
        _fp8_gemm_tma[grid](
            a_desc, b_desc, c_desc, ws, M, N, K,
            BM, BN, BK, GROUP_M, NUM_SMS,
            num_warps=num_warps, num_stages=num_stages,
        )
        return y

    # -- odd K: pad-copy into aligned buffers, then TMA -----------------------
    def _forward_padded(self, x, w, ws, y, M, N, K):
        KPAD = ((K + 127) // 128) * 128
        bufs = self._pad_bufs
        if bufs is None or bufs[0].shape[0] != M or bufs[0].shape[1] != KPAD or bufs[1].shape[0] != N:
            xp = torch.empty(M, KPAD, dtype=torch.float8_e4m3fn, device=x.device)
            wp = torch.empty(N, KPAD, dtype=torch.float8_e4m3fn, device=x.device)
            bufs = (xp, wp)
            self._pad_bufs = bufs
        xp, wp = bufs
        BLK = 1024
        _pad_copy[(M * triton.cdiv(KPAD, BLK),)](x, xp, K, KPAD, BLK, num_warps=4)
        _pad_copy[(N * triton.cdiv(KPAD, BLK),)](w, wp, K, KPAD, BLK, num_warps=4)
        return self._forward_tma(xp, wp, ws, y, M, N, K)

    # -- skinny ----------------------------------------------------------------
    def _forward_skinny(self, x, w, ws, y, M, N, K):
        cfg = _SKINNY_CONFIGS.get((self.M, self.N, self.K)) or (128, max(32, M), 256, 4, 4, 3)
        BM, BN, BK, SPLIT_K, num_warps, num_stages = cfg
        BN = max(BN, triton.next_power_of_2(M))

        need = N * M
        ws_buf = self._workspace
        if ws_buf is None or ws_buf.numel() < need:
            ws_buf = torch.empty(need, dtype=torch.float32, device=x.device)
            self._workspace = ws_buf
        ws_buf = ws_buf[:need].view(N, M)
        ws_buf.zero_()

        w_desc = TensorDescriptor.from_tensor(w, [BM, BK])
        x_desc = TensorDescriptor.from_tensor(x, [BN, BK])
        grid = (triton.cdiv(N, BM) * SPLIT_K,)
        _fp8_gemm_skinny[grid](
            w_desc, x_desc, ws_buf, M, N, K,
            SPLIT_K, BM, BN, BK,
            num_warps=num_warps, num_stages=num_stages,
        )
        EBM, EBN = 32, 128
        egrid = (triton.cdiv(M, EBM) * triton.cdiv(N, EBN),)
        _epilogue_skinny[egrid](ws_buf, y, ws, M, N, EBM, EBN, num_warps=4)
        return y


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


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


M = 4096
N = 4096
K = 4096