import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a series of operations: Gemm, Subtract, GlobalAvgPool, LogSumExp, GELU, and ResidualAdd.
    """
    def __init__(self, in_features, out_features, bias=True):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=bias)
        self.subtract = nn.Parameter(torch.randn(out_features))

    def forward(self, x):
        original_x = x.clone().detach()
        # Gemm
        x = self.gemm(x)

        # Subtract
        x = x - self.subtract

        # GlobalAvgPool
        x = torch.mean(x, dim=1, keepdim=True)

        # LogSumExp
        x = torch.logsumexp(x, dim=1, keepdim=True)

        # GELU
        x = torch.nn.functional.gelu(x)

        # ResidualAdd
        x = x + original_x

        return x

batch_size = 2048
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 <vector>

std::vector<torch::Tensor> precompute_cuda(torch::Tensor weight, torch::Tensor bias, torch::Tensor subtract);
torch::Tensor apply_cuda(torch::Tensor x, torch::Tensor colsum, torch::Tensor csum, int64_t O);

std::vector<torch::Tensor> precompute(torch::Tensor weight, torch::Tensor bias, torch::Tensor subtract) {
    return precompute_cuda(weight, bias, subtract);
}

torch::Tensor apply(torch::Tensor x, torch::Tensor colsum, torch::Tensor csum, int64_t O) {
    return apply_cuda(x, colsum, csum, O);
}
'''

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

__global__ void partial_colsum_kernel(const float* __restrict__ w, float* __restrict__ partial,
                                      int O, int I) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int rb = blockIdx.y;
    if (col >= I) return;
    int start = rb * 256;
    int end = min(start + 256, O);
    float s = 0.0f;
    #pragma unroll 4
    for (int o = start; o < end; ++o) {
        s += w[(long long)o * I + col];
    }
    partial[(long long)rb * I + col] = s;
}

__global__ void reduce_colsum_kernel(const float* __restrict__ partial, float* __restrict__ colsum,
                                     int I, int nrb) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= I) return;
    float s = 0.0f;
    #pragma unroll
    for (int rb = 0; rb < 32; ++rb) {
        if (rb < nrb) s += partial[(long long)rb * I + col];
    }
    colsum[col] = s;
}

__inline__ __device__ float warp_sum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

__global__ void reduce_const_kernel(const float* __restrict__ bias, const float* __restrict__ sub,
                                    float* __restrict__ csum, int O, int has_bias) {
    __shared__ float sh[8];
    int tid = threadIdx.x;
    float s = 0.0f;
    for (int i = tid; i < O; i += 256) {
        s += (has_bias ? bias[i] : 0.0f) - sub[i];
    }
    s = warp_sum(s);
    if ((tid & 31) == 0) sh[tid >> 5] = s;
    __syncthreads();
    if (tid < 32) {
        s = (tid < 8) ? sh[tid] : 0.0f;
        s = warp_sum(s);
        if (tid == 0) *csum = s;
    }
}

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

__global__ void apply_kernel(const float* __restrict__ x, const float* __restrict__ colsum,
                             const float* __restrict__ csum, float* __restrict__ out,
                             int I, int O) {
    __shared__ float sh[8];
    int b = blockIdx.x;
    int tid = threadIdx.x;
    const float* row = x + (long long)b * I;
    float s = 0.0f;
    #pragma unroll 4
    for (int i = tid; i < I; i += 256) {
        s += row[i] * colsum[i];
    }
    s = warp_sum(s);
    if ((tid & 31) == 0) sh[tid >> 5] = s;
    __syncthreads();
    if (tid < 32) {
        s = (tid < 8) ? sh[tid] : 0.0f;
        s = warp_sum(s);
        if (tid == 0) {
            float m = (s + *csum) / (float)O;
            sh[0] = gelu_exact(m);
        }
    }
    __syncthreads();
    float g = sh[0];
    float* orow = out + (long long)b * I;
    #pragma unroll 4
    for (int i = tid; i < I; i += 256) {
        orow[i] = row[i] + g;
    }
}

std::vector<torch::Tensor> precompute_cuda(torch::Tensor weight, torch::Tensor bias, torch::Tensor subtract) {
    CHECK_CUDA(weight); CHECK_CUDA(subtract);
    CHECK_FLOAT(weight); CHECK_FLOAT(subtract);
    int O = (int)weight.size(0);
    int I = (int)weight.size(1);
    auto opts = weight.options();
    int nrb = (O + 255) / 256;
    auto partial = torch::empty({nrb, I}, opts);
    auto colsum = torch::empty({I}, opts);
    auto csum = torch::empty({1}, opts);
    dim3 block(256);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 grid1((I + 255) / 256, nrb);
    partial_colsum_kernel<<<grid1, block, 0, stream>>>(weight.data_ptr<float>(), partial.data_ptr<float>(), O, I);
    dim3 grid2((I + 255) / 256);
    reduce_colsum_kernel<<<grid2, block, 0, stream>>>(partial.data_ptr<float>(), colsum.data_ptr<float>(), I, nrb);
    int has_bias = (bias.defined() && bias.numel() == O) ? 1 : 0;
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    reduce_const_kernel<<<1, block, 0, stream>>>(bptr, subtract.data_ptr<float>(), csum.data_ptr<float>(), O, has_bias);
    return {colsum, csum};
}

torch::Tensor apply_cuda(torch::Tensor x, torch::Tensor colsum, torch::Tensor csum, int64_t O64) {
    CHECK_CUDA(x); CHECK_CUDA(colsum); CHECK_CUDA(csum);
    CHECK_FLOAT(x); CHECK_FLOAT(colsum); CHECK_FLOAT(csum);
    int B = (int)x.size(0);
    int I = (int)x.size(1);
    int O = (int)O64;
    auto out = torch::empty_like(x);
    dim3 block(256);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    apply_kernel<<<B, block, 0, stream>>>(x.data_ptr<float>(), colsum.data_ptr<float>(), csum.data_ptr<float>(), out.data_ptr<float>(), I, O);
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=bias)
        self.subtract = nn.Parameter(torch.randn(out_features))
        self._cache_colsum = None
        self._cache_csum = None
        self._cache_versions = None

    def forward(self, x):
        if self.gemm.bias is None:
            bias = torch.empty(0, device=x.device, dtype=x.dtype)
            bv = -1
        else:
            bias = self.gemm.bias
            bv = bias._version
        versions = (self.gemm.weight._version, bv, self.subtract._version)
        if self._cache_colsum is None or self._cache_versions != versions:
            vals = _ext.precompute(self.gemm.weight, bias, self.subtract)
            self._cache_colsum = vals[0]
            self._cache_csum = vals[1]
            self._cache_versions = versions
        return _ext.apply(x, self._cache_colsum, self._cache_csum, self.gemm.weight.shape[0])