"""FP8 e4m3 GEMM for RTX PRO 6000 (SM120 Blackwell).

Real fp8 x fp8 tensor-core MMA (e4m3 inputs, fp32 accumulate) with a fused
per-output-channel (N) dequant scale:  y = (x @ w.T) * weight_scale  -> bf16.

Key implementation facts for SM120 / Triton 3.6:
  * K-tile tail predication is catastrophically bad in the Triton codegen here:
    any K not a multiple of BLOCK_K runs ~14x slower. We avoid it entirely by
    padding the fp8 operands along K to a multiple of BLOCK_K with zeros (zero
    padding does not change the matmul), so the inner loop has no K predicate.
  * Output is written in bfloat16 to cut the write traffic of the memory-bound
    decode shape.
  * Per-column scale is fused into the epilogue of each tile.
  * Config is chosen per regime (compute-bound vs decode) from a small curated
    set that was measured to be near-roofline on SM120; a tiny autotune over
    that set picks the winner per (M, N, Kpad) without the noise of a large
    search space.
"""

import torch
import torch.nn as nn

import triton
import triton.language as tl

OP_TYPE = "gemm"
SUPPORTED_PRECISIONS = ["fp8_e4m3"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
E4M3_MAX = 448.0


def _configs():
    # Curated, measured-near-roofline configs on SM120. Two regimes:
    #   compute-bound (large M, N): large tiles, 8 warps, BK=64.
    #   decode / skinny-M: small BM so the single M-tile still fills the GPU;
    #     smaller BN yields more concurrent thread blocks to saturate DRAM BW,
    #     and more warps keep the memory pipeline full.
    return [
        # --- compute bound ---
        triton.Config({"BM": 128, "BN": 256, "BK": 64}, num_stages=3, num_warps=8),
        triton.Config({"BM": 256, "BN": 128, "BK": 64}, num_stages=3, num_warps=8),
        triton.Config({"BM": 128, "BN": 256, "BK": 64}, num_stages=4, num_warps=8),
        triton.Config({"BM": 256, "BN": 128, "BK": 64}, num_stages=4, num_warps=8),
        # --- decode / bandwidth bound ---
        triton.Config({"BM": 32, "BN": 32, "BK": 128}, num_stages=5, num_warps=8),
        triton.Config({"BM": 32, "BN": 64, "BK": 128}, num_stages=5, num_warps=8),
        triton.Config({"BM": 64, "BN": 32, "BK": 128}, num_stages=5, num_warps=8),
        triton.Config({"BM": 64, "BN": 64, "BK": 128}, num_stages=5, num_warps=4),
    ]


@triton.autotune(configs=_configs(), key=["M", "N", "Kpad"])
@triton.jit
def _gemm_kernel(
    a_ptr, b_ptr, s_ptr, c_ptr,
    M, N, Kpad,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_cn,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pidm = tl.program_id(0)
    pidn = tl.program_id(1)

    offm = pidm * BM + tl.arange(0, BM)
    offn = pidn * BN + tl.arange(0, BN)
    offk = tl.arange(0, BK)

    a_ptrs = a_ptr + offm[:, None] * stride_am + offk[None, :] * stride_ak
    b_ptrs = b_ptr + offn[:, None] * stride_bn + offk[None, :] * stride_bk

    m_mask = offm < M
    n_mask = offn < N

    acc = tl.zeros((BM, BN), dtype=tl.float32)

    # K is guaranteed a multiple of BK by the caller, so this loop carries no
    # K predicate at all.
    for _ in range(0, Kpad, BK):
        a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
        b = tl.load(b_ptrs, mask=n_mask[:, None], other=0.0)
        acc += tl.dot(a, b.T)
        a_ptrs += BK * stride_ak
        b_ptrs += BK * stride_bk

    # Fused per-output-channel dequant scale.
    scale = tl.load(s_ptr + offn, mask=n_mask, other=0.0)
    acc = acc * scale[None, :]

    c_ptrs = c_ptr + offm[:, None] * stride_cn + offn[None, :]
    tl.store(c_ptrs, acc.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])


class Model(nn.Module):
    """y = ((x @ w.T) * weight_scale).to(bf16).

    x: fp8_e4m3 (M, K).  w: fp8_e4m3 (N, K) normalized to the e4m3 range.
    weight_scale: (N,) per-output-channel dequant scale.
    """

    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)                           # (N, K) fp8
        self.register_buffer("weight_scale", s.squeeze(1).to(torch.float32))  # (N,)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight
        M, N, K = x.shape[0], w.shape[0], w.shape[1]
        assert x.shape[1] == K

        # Pad K to a multiple of 128 (the LCM of every BK we autotune), zero
        # filled. Two reasons:
        #   * It makes the inner loop carry no K predicate at all — the fix for
        #     the SM120 codegen cliff where a predicated K tail runs ~14x slow.
        #   * Because every BK in the config set divides 128, Kpad is always an
        #     exact multiple of the chosen BK, so the loop never overshoots the
        #     padded buffer (no out-of-bounds reads on the tail tile).
        BK_MIN = 128
        Kpad = triton.cdiv(K, BK_MIN) * BK_MIN
        if Kpad != K:
            a = torch.nn.functional.pad(x, (0, Kpad - K))
            b = torch.nn.functional.pad(w, (0, Kpad - K))
        else:
            a, b = x, w

        c = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)

        def _grid(meta):
            return (triton.cdiv(M, meta["BM"]), triton.cdiv(N, meta["BN"]))

        _gemm_kernel[_grid](
            a, b, self.weight_scale, c,
            M, N, Kpad,
            a.stride(0), a.stride(1),
            b.stride(0), b.stride(1),
            c.stride(0),
        )
        return c


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]