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):
    """
    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(Model, 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))))

batch_size = 8192
dim = 8192

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

def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

torch::Tensor gelu_forward_cuda(torch::Tensor x);
"""

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

static __device__ __forceinline__ float gelu_tanh_f32(float x) {
    const float kAlpha = 0.79788456080286535588f;  // sqrt(2/pi)
    const float kBeta  = 0.044715f;
    float x2 = x * x;
    float u = kAlpha * (x + kBeta * x * x2);
    return 0.5f * x * (1.0f + tanhf(u));
}

__global__ void gelu_float4_kernel(const float* __restrict__ x,
                                   float* __restrict__ y,
                                   long long n4) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
    for (; i < n4; i += stride) {
        float4 v = x4[i];
        v.x = gelu_tanh_f32(v.x);
        v.y = gelu_tanh_f32(v.y);
        v.z = gelu_tanh_f32(v.z);
        v.w = gelu_tanh_f32(v.w);
        y4[i] = v;
    }
}

__global__ void gelu_tail_kernel(const float* __restrict__ x,
                                 float* __restrict__ y,
                                 long long start,
                                 long long n) {
    long long idx = start + (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) y[idx] = gelu_tanh_f32(x[idx]);
}

torch::Tensor gelu_forward_cuda(torch::Tensor x) {
    auto y = torch::empty_like(x);
    const long long n = x.numel();
    if (n == 0) return y;

    const int threads = 256;
    const long long n4 = n >> 2;
    if (n4 > 0) {
        int blocks = (int)((n4 + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        gelu_float4_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), n4);
    }
    const long long start = n4 << 2;
    if (start < n) {
        gelu_tail_kernel<<<1, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), start, n);
    }
    return y;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, x):
        return _ext.gelu_forward_cuda(x)