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>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>

void gelu_softmax_launcher(torch::Tensor y, torch::Tensor bias, int M, int N, cudaStream_t stream);

torch::Tensor linear_gelu_softmax(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat && bias.scalar_type() == at::kFloat, "float32 only");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "contiguous required");

    const int64_t M64 = x.size(0);
    const int64_t K64 = x.size(1);
    const int64_t N64 = weight.size(0);
    auto y = torch::empty({M64, N64}, x.options());

    const int M = (int)M64;
    const int K = (int)K64;
    const int N = (int)N64;
    const float alpha = 1.0f;
    const float beta = 0.0f;

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    // Row-major C[M,N] is viewed as column-major C^T[N,M].  The row-major
    // weight[N,K] is a column-major KxN matrix, so op(weight)=T gives NxK.
    cublasStatus_t stat = cublasGemmEx(handle,
                                       CUBLAS_OP_T, CUBLAS_OP_N,
                                       N, M, K,
                                       &alpha,
                                       weight.data_ptr<float>(), CUDA_R_32F, K,
                                       x.data_ptr<float>(), CUDA_R_32F, K,
                                       &beta,
                                       y.data_ptr<float>(), CUDA_R_32F, N,
                                       CUBLAS_COMPUTE_32F_FAST_TF32,
                                       CUBLAS_GEMM_DEFAULT_TENSOR_OP);
    TORCH_CHECK(stat == CUBLAS_STATUS_SUCCESS, "cublasGemmEx failed");

    gelu_softmax_launcher(y, bias, M, N, at::cuda::getCurrentCUDAStream());
    return y;
}
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ __forceinline__ float gelu_exact(float x) {
    return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440084436210485f));
}

template<int BLOCK>
__global__ void gelu_softmax_kernel(float* __restrict__ y,
                                    const float* __restrict__ bias,
                                    int N) {
    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    float* __restrict__ rowp = y + (size_t)row * (size_t)N;

    // First pass: add bias, apply exact GELU, store it, and find row max.
    float maxv = -INFINITY;
    for (int j = tid; j < N; j += BLOCK) {
        float v = gelu_exact(rowp[j] + bias[j]);
        rowp[j] = v;
        maxv = fmaxf(maxv, v);
    }

    __shared__ float smem[BLOCK];
    smem[tid] = maxv;
    __syncthreads();
    for (int s = BLOCK >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    maxv = smem[0];

    // Second pass: exponentiate shifted values and reduce the sum.
    float sum = 0.0f;
    for (int j = tid; j < N; j += BLOCK) {
        float e = expf(rowp[j] - maxv);
        rowp[j] = e;
        sum += e;
    }
    smem[tid] = sum;
    __syncthreads();
    for (int s = BLOCK >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    const float inv_sum = 1.0f / smem[0];

    // Third pass: normalize.
    for (int j = tid; j < N; j += BLOCK) {
        rowp[j] *= inv_sum;
    }
}

void gelu_softmax_launcher(torch::Tensor y, torch::Tensor bias, int M, int N, cudaStream_t stream) {
    constexpr int BLOCK = 256;
    gelu_softmax_kernel<BLOCK><<<M, BLOCK, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), N);
}
'''

ext = load_inline(
    name='lin_gelu_softmax_ext_submitted',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['linear_gelu_softmax'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)

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