import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# From https://github.com/karpathy/minGPT/blob/master/mingpt/model.py

class Model(nn.Module):
    """
    A vanilla multi-head masked self-attention layer with a projection at the end.
    It is possible to use torch.nn.MultiheadAttention here but I am including an
    explicit implementation here to show that there is nothing too scary here.
    """

    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        assert n_embd % n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        # output projection
        self.c_proj = nn.Linear(n_embd, n_embd)
        # regularization
        self.attn_dropout = nn.Dropout(attn_pdrop)
        self.resid_dropout = nn.Dropout(resid_pdrop)
        # causal mask to ensure that attention is only applied to the left in the input sequence
        self.register_buffer("bias", torch.tril(torch.ones(max_seqlen, max_seqlen))
                                     .view(1, 1, max_seqlen, max_seqlen))
        self.n_head = n_head
        self.n_embd = n_embd

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

batch_size = 128
max_seqlen = 1024
seq_len = 512
n_embd = 768
n_head = 8
attn_pdrop = 0.0
resid_pdrop = 0.0

def get_inputs():
    return [torch.rand(batch_size, seq_len, n_embd)]

def get_init_inputs():
    return [n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen]

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor causal_attention_cuda(torch::Tensor qkv, int64_t n_head);

torch::Tensor causal_attention(torch::Tensor qkv, int64_t n_head) {
    return causal_attention_cuda(qkv, n_head);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math_constants.h>

__inline__ __device__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

// One warp computes one (batch, head, query-position) row of causal attention.
// Specialized for the small GPT head dimension used here (96), but written for D <= 128.
__global__ void causal_attn_warp_kernel(
    const float* __restrict__ qkv,
    float* __restrict__ out,
    int B, int T, int C, int H, int D,
    float scale,
    long total_rows)
{
    int lane = threadIdx.x & 31;
    int warp_in_block = threadIdx.x >> 5;
    int warps_per_block = blockDim.x >> 5;
    long row = ((long)blockIdx.x) * warps_per_block + warp_in_block;
    if (row >= total_rows) return;

    int t = (int)(row % T);
    long tmp = row / T;
    int h = (int)(tmp % H);
    int b = (int)(tmp / H);

    long q_base = ((long)b * T + t) * (long)(3 * C) + h * D;

    float qv0 = 0.0f, qv1 = 0.0f, qv2 = 0.0f, qv3 = 0.0f;
    int d0 = lane;
    int d1 = lane + 32;
    int d2 = lane + 64;
    int d3 = lane + 96;
    if (d0 < D) qv0 = qkv[q_base + d0];
    if (d1 < D) qv1 = qkv[q_base + d1];
    if (d2 < D) qv2 = qkv[q_base + d2];
    if (d3 < D) qv3 = qkv[q_base + d3];

    float m = -CUDART_INF_F;

    // Pass 1: maximum logit for a stable softmax.
    for (int j = 0; j <= t; ++j) {
        long k_base = ((long)b * T + j) * (long)(3 * C) + C + h * D;
        float dot = 0.0f;
        if (d0 < D) dot += qv0 * qkv[k_base + d0];
        if (d1 < D) dot += qv1 * qkv[k_base + d1];
        if (d2 < D) dot += qv2 * qkv[k_base + d2];
        if (d3 < D) dot += qv3 * qkv[k_base + d3];
        dot = warp_reduce_sum(dot);
        float score = __shfl_sync(0xffffffff, dot, 0) * scale;
        m = fmaxf(m, score);
    }

    float denom = 0.0f;
    float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f;

    // Pass 2: denominator and weighted sum.  Each lane owns up to 4 output channels.
    for (int j = 0; j <= t; ++j) {
        long base = ((long)b * T + j) * (long)(3 * C);
        long k_base = base + C + h * D;
        float dot = 0.0f;
        if (d0 < D) dot += qv0 * qkv[k_base + d0];
        if (d1 < D) dot += qv1 * qkv[k_base + d1];
        if (d2 < D) dot += qv2 * qkv[k_base + d2];
        if (d3 < D) dot += qv3 * qkv[k_base + d3];
        dot = warp_reduce_sum(dot);
        float score = __shfl_sync(0xffffffff, dot, 0) * scale;
        float w = expf(score - m);
        denom += w;

        long v_base = base + 2 * C + h * D;
        if (d0 < D) acc0 += w * qkv[v_base + d0];
        if (d1 < D) acc1 += w * qkv[v_base + d1];
        if (d2 < D) acc2 += w * qkv[v_base + d2];
        if (d3 < D) acc3 += w * qkv[v_base + d3];
    }

    float inv = 1.0f / denom;
    long o_base = ((long)b * T + t) * (long)C + h * D;
    if (d0 < D) out[o_base + d0] = acc0 * inv;
    if (d1 < D) out[o_base + d1] = acc1 * inv;
    if (d2 < D) out[o_base + d2] = acc2 * inv;
    if (d3 < D) out[o_base + d3] = acc3 * inv;
}

torch::Tensor causal_attention_cuda(torch::Tensor qkv, int64_t n_head) {
    TORCH_CHECK(qkv.is_cuda(), "qkv must be CUDA");
    TORCH_CHECK(qkv.scalar_type() == at::kFloat, "qkv must be float32");
    TORCH_CHECK(qkv.dim() == 3, "qkv must have shape [B,T,3C]");
    TORCH_CHECK(qkv.is_contiguous(), "qkv must be contiguous");

    int B = (int)qkv.size(0);
    int T = (int)qkv.size(1);
    int C3 = (int)qkv.size(2);
    int C = C3 / 3;
    int H = (int)n_head;
    int D = C / H;
    TORCH_CHECK(C3 == 3 * C && C == H * D, "invalid qkv/head dimensions");
    TORCH_CHECK(D <= 128, "head dimension too large for this kernel");

    auto out = torch::empty({B, T, C}, qkv.options());
    long total_rows = (long)B * H * T;
    const int threads = 128;  // 4 warps/block
    const int warps = threads / 32;
    dim3 grid((unsigned int)((total_rows + warps - 1) / warps));
    float scale = rsqrtf((float)D);

    causal_attn_warp_kernel<<<grid, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        qkv.data_ptr<float>(), out.data_ptr<float>(), B, T, C, H, D, scale, total_rows);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="causal_attn_warp_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["causal_attention"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        assert n_embd % n_head == 0
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        self.c_proj = nn.Linear(n_embd, n_embd)
        self.attn_dropout = nn.Dropout(attn_pdrop)
        self.resid_dropout = nn.Dropout(resid_pdrop)
        self.register_buffer(
            "bias",
            torch.tril(torch.ones(max_seqlen, max_seqlen)).view(1, 1, max_seqlen, max_seqlen),
        )
        self.n_head = n_head
        self.n_embd = n_embd

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x).contiguous()

        # Fast fused causal attention for the benchmark path (float32 CUDA, p=0).
        if x.is_cuda and x.dtype is torch.float32 and (C // self.n_head) <= 128 and self.attn_dropout.p == 0.0:
            y = _ext.causal_attention(qkv, self.n_head)
        else:
            q, k, v = qkv.split(self.n_embd, dim=2)
            k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
            q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
            v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)

        y = self.c_proj(y)
        y = self.resid_dropout(y)
        return y