import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies GELU, and then applies Softmax.
    """
    def __init__(self, in_features, out_features):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        x = self.linear(x)
        x = torch.nn.functional.gelu(x)
        x = torch.nn.functional.softmax(x, dim=1)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192

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

def get_init_inputs():
    return [in_features, out_features]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor linear_gelu_softmax(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
'''

_cuda_src = r'''
#include <torch/types.h>
#include <ATen/ATen.h>
#include <ATen/ops/linear.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.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;
}

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

__inline__ __device__ float block_reduce_sum(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    int nwarps = (blockDim.x + 31) >> 5;
    v = (threadIdx.x < nwarps) ? shared[lane] : 0.0f;
    if (wid == 0) v = warp_reduce_sum(v);
    if (threadIdx.x == 0) shared[0] = v;
    __syncthreads();
    return shared[0];
}

__inline__ __device__ float block_reduce_max(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_max(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    int nwarps = (blockDim.x + 31) >> 5;
    v = (threadIdx.x < nwarps) ? shared[lane] : -CUDART_INF_F;
    if (wid == 0) v = warp_reduce_max(v);
    if (threadIdx.x == 0) shared[0] = v;
    __syncthreads();
    return shared[0];
}

__global__ void gelu_softmax_kernel(float* __restrict__ y, int rows, int cols) {
    int r = blockIdx.x;
    if (r >= rows) return;
    float* row = y + (size_t)r * cols;

    const float inv_sqrt2 = 0.70710678118654752440f;
    float local_max = -CUDART_INF_F;

    // GELU in-place, while finding the row maximum for softmax.
    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
        float x = row[c];
        float g = 0.5f * x * (1.0f + erff(x * inv_sqrt2));
        row[c] = g;
        local_max = fmaxf(local_max, g);
    }

    float m = block_reduce_max(local_max);

    float local_sum = 0.0f;
    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
        float e = expf(row[c] - m);
        row[c] = e;
        local_sum += e;
    }

    float s = block_reduce_sum(local_sum);
    float inv_s = 1.0f / s;

    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
        row[c] *= inv_s;
    }
}

torch::Tensor linear_gelu_softmax(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    c10::cuda::CUDAGuard device_guard(x.device());
    torch::Tensor y = at::linear(x, weight, bias);

    const int64_t cols64 = y.size(-1);
    const int64_t rows64 = y.numel() / cols64;
    const int cols = (int)cols64;
    const int rows = (int)rows64;

    const int threads = 256;
    gelu_softmax_kernel<<<rows, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), rows, cols);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='linear_gelu_softmax_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['linear_gelu_softmax'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    """
    Drop-in replacement: Linear -> exact GELU -> row-wise Softmax.
    The GEMM is delegated to ATen/cuBLAS, and GELU+Softmax are fused into one CUDA kernel.
    """
    def __init__(self, in_features, out_features):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        return _ext.linear_gelu_softmax(x, self.linear.weight, self.linear.bias)