"""Fused W4A16 GEMM for SM120.

The packed weight is consumed directly by the Triton kernel.  In particular,
there is deliberately no persistent/cached dequantized weight: every byte is
unpacked next to the tensor-core dot which consumes it.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import triton
import triton.language as tl


GROUP_SIZE = 128


@triton.jit
def _w4a16_gemv_kernel(
    x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
    N: tl.constexpr, K: tl.constexpr, BN: tl.constexpr,
):
    pid_n = tl.program_id(0)
    rn = pid_n * BN + tl.arange(0, BN)
    rp = tl.arange(0, 64)
    acc = tl.zeros((BN,), tl.float32)

    for group in tl.static_range(0, K // 128):
        packed = tl.load(
            q_ptr + (group * 64 + rp[:, None]) * N + rn[None, :]
        )
        qlo = (packed & 15).to(tl.bfloat16)
        qhi = (packed >> 4).to(tl.bfloat16)
        s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
        z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
        wlo = ((qlo - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
        whi = ((qhi - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
        xlo = tl.load(x_ptr + group * 128 + 2 * rp)
        xhi = tl.load(x_ptr + group * 128 + 2 * rp + 1)
        products = (
            xlo[:, None].to(tl.float32) * wlo.to(tl.float32)
            + xhi[:, None].to(tl.float32) * whi.to(tl.float32)
        )
        acc += tl.sum(products, axis=0)

    tl.store(out_ptr + rn, acc.to(tl.bfloat16))


@triton.jit
def _w4a16_direct_kernel(
    x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    rk = tl.arange(0, BK)
    acc = tl.zeros((BM, BN), tl.float32)

    for k0 in tl.static_range(0, K, BK):
        k = k0 + rk
        x = tl.load(
            x_ptr + rm[:, None] * K + k[None, :],
            mask=rm[:, None] < M,
            other=0.0,
        )
        packed = tl.load(q_ptr + (k[:, None] // 2) * N + rn[None, :])
        shift = (k[:, None] & 1) * 4
        q = ((packed >> shift) & 15).to(tl.bfloat16)
        group = k0 // 128
        s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
        z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
        w = ((q - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
        acc += tl.dot(x, w)

    tl.store(
        out_ptr + rm[:, None] * N + rn[None, :],
        acc.to(tl.bfloat16),
        mask=rm[:, None] < M,
    )


@triton.jit
def _w4a16_splitk_kernel(
    x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
    N: tl.constexpr, K: tl.constexpr, BN: tl.constexpr,
    SPLIT_K: tl.constexpr,
):
    split = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = tl.arange(0, 16)
    rn = pid_n * BN + tl.arange(0, BN)
    rk = tl.arange(0, 128)
    acc = tl.zeros((16, BN), tl.float32)
    groups_per_split: tl.constexpr = (K // 128) // SPLIT_K

    for gi in tl.static_range(0, groups_per_split):
        group = split * groups_per_split + gi
        k = group * 128 + rk
        x = tl.load(x_ptr + rm[:, None] * K + k[None, :], mask=rm[:, None] < 1,
                    other=0.0)
        packed = tl.load(q_ptr + (k[:, None] // 2) * N + rn[None, :])
        q = ((packed >> ((k[:, None] & 1) * 4)) & 15).to(tl.bfloat16)
        s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
        z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
        w = ((q - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
        acc += tl.dot(x, w)

    tl.atomic_add(out_ptr + rm[:, None] * N + rn[None, :], acc,
                  mask=rm[:, None] < 1)


@triton.jit
def _fp32_to_bf16_kernel(src_ptr, dst_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    tl.store(dst_ptr + offsets, tl.load(src_ptr + offsets, mask=offsets < N),
             mask=offsets < N)


@triton.jit
def _w4a16_grouped_kernel(
    x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
    BM: tl.constexpr, BN: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    rk = tl.arange(0, 128)
    acc = tl.zeros((BM, BN), tl.float32)

    for group in tl.static_range(0, K // 128):
        k = group * 128 + rk
        x = tl.load(
            x_ptr + rm[:, None] * K + k[None, :],
            mask=rm[:, None] < M,
            other=0.0,
        )
        packed = tl.load(q_ptr + (k[:, None] // 2) * N + rn[None, :])
        q = ((packed >> ((k[:, None] & 1) * 4)) & 15).to(tl.bfloat16)
        s = tl.load(scales_ptr + group * N + rn).to(tl.float32)
        z = tl.load(zeros_ptr + group * N + rn).to(tl.float32)

        qdot = tl.dot(x, q)
        xsum = tl.sum(x.to(tl.float32), axis=1)
        acc += (qdot - xsum[:, None] * z[None, :]) * s[None, :]

    tl.store(
        out_ptr + rm[:, None] * N + rn[None, :],
        acc.to(tl.bfloat16),
        mask=rm[:, None] < M,
    )


@triton.jit
def _w4a16_interleave_kernel(
    x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
    BM: tl.constexpr, BN: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    rp = tl.arange(0, 64)
    acc = tl.zeros((BM, BN), tl.float32)

    for group in tl.static_range(0, K // 128):
        packed = tl.load(
            q_ptr + (group * 64 + rp[:, None]) * N + rn[None, :]
        )
        qlo = (packed & 15).to(tl.bfloat16)
        qhi = (packed >> 4).to(tl.bfloat16)
        # interleave operates on the last dimension.  Two logical transposes
        # therefore turn [64, BN] packed rows into [128, BN] K-major nibbles
        # while retaining a single global load for each source byte.
        q = tl.trans(tl.interleave(tl.trans(qlo), tl.trans(qhi)))
        s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
        z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
        w = ((q - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
        x = tl.load(
            x_ptr + rm[:, None] * K + group * 128 + tl.arange(0, 128)[None, :],
            mask=rm[:, None] < M,
            other=0.0,
        )
        acc += tl.dot(x, w)

    tl.store(
        out_ptr + rm[:, None] * N + rn[None, :],
        acc.to(tl.bfloat16),
        mask=rm[:, None] < M,
    )


def _launch(x: torch.Tensor, q: torch.Tensor, scales: torch.Tensor,
            zeros: torch.Tensor, M: int, N: int, K: int) -> torch.Tensor:
    if M <= 1:
        partial = torch.zeros((M, N), dtype=torch.float32, device=x.device)
        if N == 4096:
            bn, warps, stages = 64, 8, 2
        else:
            bn, warps, stages = 128, 16, 4
        split_k = 16
        _w4a16_splitk_kernel[(split_k, triton.cdiv(N, bn))](
            x, q, scales, zeros, partial,
            N=N, K=K, BN=bn, SPLIT_K=split_k,
            num_warps=warps, num_stages=stages,
        )
        out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
        _fp32_to_bf16_kernel[(triton.cdiv(N, 256),)](
            partial, out, N=N, BLOCK=256, num_warps=4,
        )
        return out

    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    if M <= 16:
        bm, bn, bk, warps, stages = 16, 64, 128, 8, 2
    elif M <= 32:
        bm, bn, bk, warps, stages = 32, 64, 128, 8, 2
    else:
        bm, bn, bk, warps, stages = 32, 64, 128, 8, 2
    grid = (triton.cdiv(M, bm), triton.cdiv(N, bn))
    _w4a16_direct_kernel[grid](
        x, q, scales, zeros, out,
        M=M, N=N, K=K, BM=bm, BN=bn, BK=bk,
        num_warps=warps, num_stages=stages,
    )
    return out


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert group_size == GROUP_SIZE
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        self.register_buffer("w_q", torch.empty((K // 2, N), dtype=torch.uint8))
        self.register_buffer("scales", torch.empty((K // group_size, N), dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.empty((K // group_size, N), dtype=torch.bfloat16))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _launch(x, self.w_q, self.scales, self.zeros,
                       self.M, self.N, self.K)


M = 1
N = 12288
K = 4096


def get_inputs():
    return [torch.randn(M, K, dtype=torch.bfloat16)]


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