import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a matrix multiplication (Gemm), followed by LogSumExp, LeakyReLU, 
    LeakyReLU, GELU, and GELU activations.
    """
    def __init__(self, in_features, out_features, bias=True):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias=bias)

    def forward(self, x):
        # Gemm
        x = self.linear(x)
        # LogSumExp
        x = torch.logsumexp(x, dim=1, keepdim=True)
        # LeakyReLU
        x = torch.nn.functional.leaky_relu(x, negative_slope=0.01)
        # LeakyReLU
        x = torch.nn.functional.leaky_relu(x, negative_slope=0.01)
        # GELU
        x = torch.nn.functional.gelu(x)
        # GELU
        x = torch.nn.functional.gelu(x)
        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 fused_lse_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""

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

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

__global__ void lse_act_kernel(const float* __restrict__ y, float* __restrict__ out, int rows, int cols) {
    int r = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float s[];

    const float* row = y + (long long)r * cols;
    float m = -INFINITY;
    for (int c = tid; c < cols; c += blockDim.x) {
        float v = row[c];
        m = fmaxf(m, v);
    }
    s[tid] = m;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) s[tid] = fmaxf(s[tid], s[tid + stride]);
        __syncthreads();
    }
    m = s[0];

    float sum = 0.0f;
    for (int c = tid; c < cols; c += blockDim.x) {
        sum += expf(row[c] - m);
    }
    s[tid] = sum;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) s[tid] += s[tid + stride];
        __syncthreads();
    }

    if (tid == 0) {
        float v = m + logf(s[0]);
        // two LeakyReLUs with negative_slope=0.01
        if (v < 0.0f) v *= 0.0001f;
        // two exact GELUs (PyTorch default)
        v = gelu_exact_f(v);
        v = gelu_exact_f(v);
        out[r] = v;
    }
}

torch::Tensor fused_lse_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    // Compute the GEMM using PyTorch/cuBLAS, then fuse logsumexp and all following activations.
    auto y = at::matmul(x, weight.t());
    if (bias.numel() != 0) {
        y.add_(bias);
    }
    const int rows = (int)y.size(0);
    const int cols = (int)y.size(1);
    auto out = torch::empty({rows, 1}, y.options());
    const int threads = 256;
    lse_act_kernel<<<rows, threads, threads * sizeof(float), at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), out.data_ptr<float>(), rows, cols);
    return out;
}
"""

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

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

    def forward(self, x):
        b = self.linear.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.fused_lse_forward(x, self.linear.weight, b)