"""FP8 e4m3 x e4m3 GEMM for NVIDIA B200 (SM100, Blackwell).

Real fp8 x fp8 tensor-core MMA via Triton's TMA-descriptor persistent kernel:
cp.async.bulk.tensor (TMA) loads + tcgen05.mma (5th-gen tensor cores) with fp32
accumulate, and the per-output-channel dequant scale folded into a subtile
epilogue. Matches reference.py's interface (same weight/weight_scale buffers).

K must be a multiple of 32 (the fp8 MMA native K-dim) AND a multiple of 16 for
the TMA descriptor stride requirement; when K isn't (e.g. K=4127) we zero-pad K
to the next multiple of 128 (zero contributes nothing to the dot product).
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl

E4M3_MAX = 448.0
_PAD_MULT = 128
NUM_SMS = torch.cuda.get_device_properties(0).multi_processor_count


def _set_tma_allocator():
    def _alloc(size, align, stream):
        return torch.empty(size, device="cuda", dtype=torch.int8)
    triton.set_allocator(_alloc)


_set_tma_allocator()


# --------------------------------------------------------------------------
# Pad kernel: zero-pad a 2D (rows, K) fp8 tensor along dim=1 to (rows, K_pad).
# --------------------------------------------------------------------------
@triton.jit
def _pad_kernel(src, dst, ROWS, K, K_PAD,
                stride_sm, stride_sk, stride_dm, stride_dk,
                BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr):
    pid = tl.program_id(0)
    num_k = tl.cdiv(K_PAD, BLOCK_K)
    pid_m = pid // num_k
    pid_k = pid % num_k
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rk = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
    mmask = rm < ROWS
    valid = (rk < K)[None, :]              # src region (zero elsewhere)
    dstk = (rk < K_PAD)[None, :]           # stay inside dst row (no wraparound)
    gptr = src + rm[:, None] * stride_sm + rk[None, :] * stride_sk
    vals = tl.load(gptr, mask=mmask[:, None] & valid, other=0.0)
    dptr = dst + rm[:, None] * stride_dm + rk[None, :] * stride_dk
    tl.store(dptr, vals, mask=mmask[:, None] & dstk)


def _pad_k(t: torch.Tensor, K: int, K_pad: int) -> torch.Tensor:
    if K_pad == K:
        return t
    rows = t.shape[0]
    out = torch.empty((rows, K_pad), dtype=t.dtype, device=t.device)
    # K_pad is always a multiple of 128, so BLOCK_K=128 tiles it exactly.
    BLOCK_M, BLOCK_K = 32, 128
    grid = (triton.cdiv(rows, BLOCK_M) * triton.cdiv(K_pad, BLOCK_K),)
    _pad_kernel[grid](t, out, rows, K, K_pad,
                      t.stride(0), t.stride(1), out.stride(0), out.stride(1),
                      BLOCK_M=BLOCK_M, BLOCK_K=BLOCK_K, num_warps=4)
    return out


# --------------------------------------------------------------------------
# Persistent TMA FP8 GEMM kernel.
# --------------------------------------------------------------------------
@triton.jit
def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_M, NUM_SMS):
    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
    return pid_m, pid_n


def _gemm_configs():
    cfgs = []
    # Compute-bound: 128x256 subtile epilogue is the sweet spot on B200.
    for ns in (3, 4):
        for gm in (8,):
            cfgs.append(triton.Config(
                {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "GROUP_M": gm,
                 "EPILOGUE_SUBTILE": True, "NUM_SMS": NUM_SMS}, num_warps=8, num_stages=ns))
            cfgs.append(triton.Config(
                {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": gm,
                 "EPILOGUE_SUBTILE": False, "NUM_SMS": NUM_SMS}, num_warps=8, num_stages=ns))
    # Skinny / decode (small M).
    for bm in (32, 64):
        for bn in (128, 256):
            for nw in (4, 8):
                epi = bn >= 256
                cfgs.append(triton.Config(
                    {"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": 128, "GROUP_M": 8,
                     "EPILOGUE_SUBTILE": epi, "NUM_SMS": NUM_SMS}, num_warps=nw, num_stages=4))
    return cfgs


@triton.autotune(configs=_gemm_configs(), key=["M", "N", "K"])
@triton.jit
def _fp8_gemm_kernel(
    a_ptr, b_ptr, c_ptr, scale_ptr,
    M, N, K,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr, EPILOGUE_SUBTILE: tl.constexpr, NUM_SMS: tl.constexpr,
):
    """C[M,N] = (A[M,K] @ weight.T) * scale[N], A=x fp8, B=weight (N,K) fp8."""
    start_pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    k_tiles = tl.cdiv(K, BLOCK_K)
    num_tiles = num_pid_m * num_pid_n

    a_desc = tl.make_tensor_descriptor(a_ptr, shape=[M, K], strides=[K, 1], block_shape=[BLOCK_M, BLOCK_K])
    b_desc = tl.make_tensor_descriptor(b_ptr, shape=[N, K], strides=[K, 1], block_shape=[BLOCK_N, BLOCK_K])
    c_desc = tl.make_tensor_descriptor(c_ptr, shape=[M, N], strides=[N, 1],
                                       block_shape=[BLOCK_M, BLOCK_N if not EPILOGUE_SUBTILE else BLOCK_N // 2])

    tile_id_c = start_pid - NUM_SMS
    num_pid_in_group = GROUP_M * num_pid_n

    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True, warp_specialize=True):
        pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_M, NUM_SMS)
        offs_am = pid_m * BLOCK_M
        offs_bn = pid_n * BLOCK_N
        accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        for ki in range(k_tiles):
            offs_k = ki * BLOCK_K
            a = a_desc.load([offs_am, offs_k])
            b = b_desc.load([offs_bn, offs_k])
            accumulator = tl.dot(a, b.T, accumulator)

        tile_id_c += NUM_SMS
        pid_m, pid_n = _compute_pid(tile_id_c, num_pid_in_group, num_pid_m, GROUP_M, NUM_SMS)
        offs_cm = pid_m * BLOCK_M
        offs_cn = pid_n * BLOCK_N
        if EPILOGUE_SUBTILE:
            acc = tl.reshape(accumulator, (BLOCK_M, 2, BLOCK_N // 2))
            acc = tl.permute(acc, (0, 2, 1))
            acc0, acc1 = tl.split(acc)
            cols0 = offs_cn + tl.arange(0, BLOCK_N // 2)
            cols1 = offs_cn + BLOCK_N // 2 + tl.arange(0, BLOCK_N // 2)
            s0 = tl.load(scale_ptr + cols0)
            s1 = tl.load(scale_ptr + cols1)
            c0 = (acc0 * s0[None, :]).to(tl.bfloat16)
            c_desc.store([offs_cm, offs_cn], c0)
            c1 = (acc1 * s1[None, :]).to(tl.bfloat16)
            c_desc.store([offs_cm, offs_cn + BLOCK_N // 2], c1)
        else:
            cols = offs_cn + tl.arange(0, BLOCK_N)
            s = tl.load(scale_ptr + cols)
            c = (accumulator * s[None, :]).to(tl.bfloat16)
            c_desc.store([offs_cm, offs_cn], c)


def _fp8_gemm(x: torch.Tensor, w: torch.Tensor, scale: torch.Tensor,
              out: torch.Tensor) -> torch.Tensor:
    M, K = x.shape
    N = w.shape[0]
    grid = (NUM_SMS,)
    _fp8_gemm_kernel[grid](x, w, out, scale, M, N, K)
    return out


# --------------------------------------------------------------------------
# Split-K path for skinny M (decode-like, memory/latency bound). Each (N-tile,
# K-split) CTA writes an fp32 partial; a reduction kernel sums splits, applies
# the per-channel scale, and casts to bf16. Splitting K multiplies the CTA count
# so all SMs stay busy streaming the weight from HBM.
# --------------------------------------------------------------------------
@triton.jit
def _splitk_kernel(
    a_ptr, b_ptr, part_ptr, M, N, K,
    K_PER_SPLIT: tl.constexpr,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid = tl.program_id(0)
    num_n = tl.cdiv(N, BLOCK_N)
    split_id = pid // num_n
    pid_n = pid % num_n
    k_start = split_id * K_PER_SPLIT
    k_end = k_start + K_PER_SPLIT
    if k_end > K:
        k_end = K

    offs_am = tl.arange(0, BLOCK_M)
    offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + offs_am[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_bn[:, None] * K + offs_k[None, :]

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for ki in range(k_start, k_end, BLOCK_K):
        kk = ki + offs_k
        a = tl.load(a_ptrs + ki[None, :], mask=kk[None, :] < k_end, other=0.0)
        b = tl.load(b_ptrs + ki[None, :], mask=kk[:, None] < k_end, other=0.0)
        acc = tl.dot(a, b.T, acc)

    part_ptrs = part_ptr + split_id * M * N + offs_am[:, None] * N + offs_bn[None, :]
    tl.store(part_ptrs, acc)


@triton.jit
def _splitk_reduce(part_ptr, c_ptr, scale_ptr, M, N,
                   SK: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for s in tl.static_range(0, SK):
        p = part_ptr + s * M * N + rm[:, None] * N + rn[None, :]
        acc += tl.load(p)
    sc = tl.load(scale_ptr + rn)
    out = (acc * sc[None, :]).to(tl.bfloat16)
    tl.store(c_ptr + rm[:, None] * N + rn[None, :], out)


def _choose_sk(M, N, K, BN, BK):
    """Pick a K split count so num_n * SK ~= 2-3x NUM_SMS; SK must divide K//BK."""
    num_n = (N + BN - 1) // BN
    target = max(1, (2 * NUM_SMS) // max(1, num_n))
    k_per_bk = K // BK
    best, best_sk = 1, 1
    for sk in (1, 2, 4, 8, 16, 32):
        if k_per_bk % sk != 0:
            continue
        if abs(sk - target) < abs(best_sk - target):
            best_sk, best = sk, sk
    return max(1, best_sk)


def _fp8_gemm_splitk(x: torch.Tensor, w: torch.Tensor, scale: torch.Tensor,
                     part_buf: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
    M, K = x.shape
    N = w.shape[0]
    BM, BN, BK = 32, 128, 128
    SK = _choose_sk(M, N, K, BN, BK)
    KPS = K // SK
    num_n = triton.cdiv(N, BN)
    _splitk_kernel[(num_n * SK,)](
        x, w, part_buf, M, N, K, K_PER_SPLIT=KPS,
        BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK, num_warps=4, num_stages=4,
    )
    _splitk_reduce[(triton.cdiv(M, 32), triton.cdiv(N, 256))](
        part_buf, out, scale, M, N, SK=SK, BLOCK_M=32, BLOCK_N=256, num_warps=4,
    )
    return out


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.zeros(N, K, dtype=torch.float8_e4m3fn)
        self.register_buffer("weight", w)
        self.register_buffer("weight_scale", torch.ones(N, dtype=torch.float32))
        self._w_pad = None
        self._w_version = -1
        self._part_buf = None
        self._out_buf = None
        self._out_shape = None

    def _get_out(self, M, N, device):
        if self._out_buf is None or self._out_shape != (M, N):
            self._out_buf = torch.empty((M, N), dtype=torch.bfloat16, device=device)
            self._out_shape = (M, N)
        return self._out_buf

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M, K = x.shape
        N = self.N
        out = self._get_out(M, N, x.device)
        K_pad = ((K + _PAD_MULT - 1) // _PAD_MULT) * _PAD_MULT
        if K_pad != K:
            x = _pad_k(x, K, K_pad)
            if self.weight._version != self._w_version or self._w_pad is None:
                self._w_pad = _pad_k(self.weight, K, K_pad)
                self._w_version = self.weight._version
            w = self._w_pad
        else:
            w = self.weight
        # Skinny M (decode-like, M<=32): split-K to saturate HBM bandwidth.
        # Split-K kernel uses BLOCK_M=32 so it covers all M rows here.
        if M <= 32:
            BN, BK = 128, 128
            SK = _choose_sk(M, N, K, BN, BK)
            if self._part_buf is None or self._part_buf.shape != (SK, M, N):
                self._part_buf = torch.empty((SK, M, N), dtype=torch.float32, device=x.device)
            return _fp8_gemm_splitk(x, w, self.weight_scale, self._part_buf, out)
        return _fp8_gemm(x, w, self.weight_scale, out)


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]