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 torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
void mask_relu_scale_(torch::Tensor att, double scale);
"""

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

static __device__ __forceinline__ float apply_mask_relu_scale(float x, int col, int row, float scale) {
    float v = x * scale;
    return (col <= row && v > 0.0f) ? v : 0.0f;
}

__global__ void mask_relu_scale_1024_kernel(float4* __restrict__ a4, long long n4, float scale) {
    long long i4 = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; i4 < n4; i4 += stride) {
        long long base = i4 << 2;
        int col = (int)(base & 1023LL);
        int row = (int)((base >> 10) & 1023LL);
        float4 v = a4[i4];
        v.x = apply_mask_relu_scale(v.x, col,     row, scale);
        v.y = apply_mask_relu_scale(v.y, col + 1, row, scale);
        v.z = apply_mask_relu_scale(v.z, col + 2, row, scale);
        v.w = apply_mask_relu_scale(v.w, col + 3, row, scale);
        a4[i4] = v;
    }
}

__global__ void mask_relu_scale_generic_kernel(float* __restrict__ a, long long n, int T, float scale) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    long long TT = (long long)T * T;
    for (; idx < n; idx += stride) {
        int pos = (int)(idx % TT);
        int col = pos % T;
        int row = pos / T;
        float v = a[idx] * scale;
        if (col > row || v < 0.0f) v = 0.0f;
        a[idx] = v;
    }
}

void mask_relu_scale_(torch::Tensor att, double scale) {
    long long n = att.numel();
    int T = (int)att.size(-1);
    const int block = 256;
    if (T == 1024 && (n % 4 == 0)) {
        long long n4 = n >> 2;
        int grid = (int)((n4 + block - 1) / block);
        if (grid > 65535) grid = 65535;
        mask_relu_scale_1024_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(reinterpret_cast<float4*>(att.data_ptr<float>()), n4, (float)scale);
    } else {
        int grid = (int)((n + block - 1) / block);
        if (grid > 65535) grid = 65535;
        mask_relu_scale_generic_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(att.data_ptr<float>(), n, T, (float)scale);
    }
}
"""

_ext = load_inline(
    name="mask_relu_scale_ext_final",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["mask_relu_scale_"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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):
        B, T, C = x.size()
        hs = C // self.n_head
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, hs).transpose(1, 2)
        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
        v = v.view(B, T, self.n_head, hs).transpose(1, 2)

        att = q @ k.transpose(-2, -1)
        _ext.mask_relu_scale_(att, 1.0 / (hs ** 0.5))

        out = torch.empty((B, T, C), device=x.device, dtype=x.dtype)
        y_view = torch.as_strided(out, (B, self.n_head, T, hs), (T * C, hs, C, 1))
        torch.matmul(att, v, out=y_view)
        return out