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 NewGELU(nn.Module):
    """
    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
    Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
    """
    def __init__(self):
        super(NewGELU, self).__init__()
    
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))

class Model(nn.Module):
    """
    A multi-head masked self-attention layer with a projection at the end that uses ReLU instead of Softmax.
    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, 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)
        # 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.relu(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

        return y

batch_size = 16
max_seqlen = 1024
n_embd = 768  # Hidden dimension, typical for BERT-base size
n_head = 12   # Number of attention heads, typical for BERT-base size

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

def get_init_inputs():
    return [n_embd, n_head, 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 relu_causal_attn_forward(torch::Tensor qkv, int64_t n_head);
"""

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

static __forceinline__ __device__ float warp_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

// Specialized for head size 64. One block computes one (batch, head, query-position).
// Warp 0 writes output dims 0..31 and warp 1 writes dims 32..63.  The QK dot is
// recomputed by both warps, avoiding shared-memory reductions/barriers and allowing
// a uniform branch to skip V loads for negative ReLUed attention scores.
__global__ __launch_bounds__(64) void relu_causal_attn_hs64_kernel(
    const float* __restrict__ qkv,
    float* __restrict__ out,
    int B, int T, int C, int nH)
{
    const int hs = 64;
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;             // 0 or 1
    const int od = warp * 32 + lane;       // output dimension within head

    int bid = blockIdx.x;
    const int t = bid % T;
    bid /= T;
    const int h = bid % nH;
    const int b = bid / nH;

    const int head_base = h * hs;
    const long long token_stride = (long long)3 * C;
    const long long q_base = ((long long)b * T + t) * token_stride + head_base;
    const float q0 = qkv[q_base + lane];
    const float q1 = qkv[q_base + lane + 32];

    const float scale = 0.125f; // 1/sqrt(64)
    float acc = 0.0f;

    const long long b_base = (long long)b * T * token_stride;
    for (int j = 0; j <= t; ++j) {
        const long long kj = b_base + (long long)j * token_stride + C + head_base;
        float prod = fmaf(q1, qkv[kj + lane + 32], q0 * qkv[kj + lane]);
        float s = warp_sum(prod) * scale;
        // ReLU after causal masking. Branch is uniform within each warp.
        if (s > 0.0f) {
            const long long vj = b_base + (long long)j * token_stride + 2LL * C + head_base;
            acc = fmaf(s, qkv[vj + od], acc);
        }
    }

    const long long out_idx = ((long long)b * T + t) * C + head_base + od;
    out[out_idx] = acc;
}

// Generic slow fallback for unusual head sizes. Kept as CUDA code (not a PyTorch
// reference route); the benchmarked problem uses hs=64.
__global__ void relu_causal_attn_generic_kernel(
    const float* __restrict__ qkv,
    float* __restrict__ out,
    int B, int T, int C, int nH, int hs)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * T * C;
    if (idx >= total) return;
    int d = idx % C;
    int t = (idx / C) % T;
    int b = idx / (C * T);
    int h = d / hs;
    int hd = d - h * hs;
    const float scale = rsqrtf((float)hs);
    const long long token_stride = (long long)3 * C;
    const long long q_base = ((long long)b * T + t) * token_stride + h * hs;
    const long long b_base = (long long)b * T * token_stride;
    float acc = 0.0f;
    for (int j = 0; j <= t; ++j) {
        const long long k_base = b_base + (long long)j * token_stride + C + h * hs;
        float s = 0.0f;
        for (int r = 0; r < hs; ++r) s = fmaf(qkv[q_base + r], qkv[k_base + r], s);
        s *= scale;
        if (s > 0.0f) {
            const long long v_base = b_base + (long long)j * token_stride + 2LL * C + h * hs;
            acc = fmaf(s, qkv[v_base + hd], acc);
        }
    }
    out[idx] = acc;
}

torch::Tensor relu_causal_attn_forward(torch::Tensor qkv, int64_t n_head) {
    TORCH_CHECK(qkv.is_cuda(), "qkv must be a CUDA tensor");
    TORCH_CHECK(qkv.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(qkv.dim() == 3, "qkv must have shape [B,T,3*C]");
    TORCH_CHECK(qkv.is_contiguous(), "qkv must be contiguous");
    const int B = (int)qkv.size(0);
    const int T = (int)qkv.size(1);
    const int C3 = (int)qkv.size(2);
    TORCH_CHECK(C3 % 3 == 0, "last dimension must be 3*C");
    const int C = C3 / 3;
    const int nH = (int)n_head;
    TORCH_CHECK(C % nH == 0, "C must be divisible by n_head");
    const int hs = C / nH;

    auto out = torch::empty({B, T, C}, qkv.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (hs == 64) {
        dim3 block(64);
        dim3 grid((unsigned int)(B * nH * T));
        relu_causal_attn_hs64_kernel<<<grid, block, 0, stream>>>(
            qkv.data_ptr<float>(), out.data_ptr<float>(), B, T, C, nH);
    } else {
        int threads = 128;
        int total = B * T * C;
        int blocks = (total + threads - 1) / threads;
        relu_causal_attn_generic_kernel<<<blocks, threads, 0, stream>>>(
            qkv.data_ptr<float>(), out.data_ptr<float>(), B, T, C, nH, hs);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_relu_attn_ext = load_inline(
    name="relu_causal_attn_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["relu_causal_attn_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, n_embd, n_head, 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.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):
        qkv = self.c_attn(x)
        return _relu_attn_ext.relu_causal_attn_forward(qkv.contiguous(), self.n_head)