import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a GEMM, BatchNorm, GELU, and ReLU in sequence.
    """
    def __init__(self, in_features, out_features):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.batch_norm = nn.BatchNorm1d(out_features)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.gemm(x)
        x = self.batch_norm(x)
        x = torch.nn.functional.gelu(x)
        x = torch.relu(x)
        return x

batch_size = 16384
in_features = 4096
out_features = 4096

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_forward(torch::Tensor x, torch::Tensor w, torch::Tensor linear_bias,
                            torch::Tensor bn_w, torch::Tensor bn_b,
                            torch::Tensor running_mean, torch::Tensor running_var,
                            bool training, double momentum, double eps);
'''

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

__device__ __forceinline__ float gelu_relu_val(float z) {
    if (z <= 0.0f) return 0.0f;
    return 0.5f * z * (1.0f + erff(z * 0.70710678118654752440f));
}

__global__ void stats_warp32_kernel(const float* __restrict__ y,
                                    float* __restrict__ psum,
                                    float* __restrict__ psq,
                                    int B, int C, int splits, int rows_per_split) {
    int tile = blockIdx.x;
    int split = blockIdx.y;
    int lane = threadIdx.x;
    int c = tile * 32 + lane;
    if (c >= C) return;
    int r0 = split * rows_per_split;
    int r1 = min(B, r0 + rows_per_split);
    float s = 0.0f;
    float q = 0.0f;
    const float* ptr = y + (int64_t)r0 * C + c;
    for (int r = r0; r < r1; ++r) {
        float v = *ptr;
        s += v;
        q = fmaf(v, v, q);
        ptr += C;
    }
    int out = split * C + c;
    psum[out] = s;
    psq[out] = q;
}

__global__ void finalize_stats_kernel(const float* __restrict__ psum,
                                      const float* __restrict__ psq,
                                      float* __restrict__ mean,
                                      float* __restrict__ invstd,
                                      float* __restrict__ running_mean,
                                      float* __restrict__ running_var,
                                      int B, int C, int splits, float momentum, float eps) {
    int tile = blockIdx.x;
    int lane = threadIdx.x;
    int c = tile * 32 + lane;
    if (c >= C) return;
    float s = 0.0f, q = 0.0f;
    for (int sp = 0; sp < splits; ++sp) {
        int idx = sp * C + c;
        s += psum[idx];
        q += psq[idx];
    }
    float m = s / (float)B;
    float var = q / (float)B - m * m;
    var = fmaxf(var, 0.0f);
    float inv = rsqrtf(var + eps);
    mean[c] = m;
    invstd[c] = inv;
    running_mean[c] = running_mean[c] * (1.0f - momentum) + momentum * m;
    float unbiased = var * ((float)B / (float)(B - 1));
    running_var[c] = running_var[c] * (1.0f - momentum) + momentum * unbiased;
}

__global__ void apply_train_bn_gelu_kernel(float* __restrict__ y,
                                           const float* __restrict__ gamma,
                                           const float* __restrict__ beta,
                                           const float* __restrict__ mean,
                                           const float* __restrict__ invstd,
                                           int B, int C, int rows_per_chunk) {
    int tile = blockIdx.x;
    int chunk = blockIdx.y;
    int lane = threadIdx.x;
    int c = tile * 32 + lane;
    if (c >= C) return;
    float g = gamma[c];
    float be = beta[c];
    float m = mean[c];
    float inv = invstd[c];
    int r0 = chunk * rows_per_chunk;
    int r1 = min(B, r0 + rows_per_chunk);
    float* ptr = y + (int64_t)r0 * C + c;
    for (int r = r0; r < r1; ++r) {
        float z = ((*ptr - m) * inv) * g + be;
        *ptr = gelu_relu_val(z);
        ptr += C;
    }
}

__global__ void bn_gelu_relu_kernel(float* __restrict__ y,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ beta,
                                    const float* __restrict__ mean,
                                    const float* __restrict__ var,
                                    int64_t total, int C, float eps) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = (C == 4096) ? (int)(idx & 4095) : (int)(idx % C);
        float z = (y[idx] - mean[c]) * rsqrtf(var[c] + eps) * gamma[c] + beta[c];
        y[idx] = gelu_relu_val(z);
    }
}

__global__ void gelu_relu_kernel(float* __restrict__ y, int64_t total) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) y[idx] = gelu_relu_val(y[idx]);
}

void launch_eval(torch::Tensor y, torch::Tensor gamma, torch::Tensor beta,
                    torch::Tensor mean, torch::Tensor var, double eps) {
    int64_t total = y.numel(); int C = (int)y.size(1);
    int threads = 256;
    int blocks = (int)std::min<int64_t>((total + threads - 1) / threads, 65535);
    bn_gelu_relu_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        mean.data_ptr<float>(), var.data_ptr<float>(), total, C, (float)eps);
}

void launch_gelu(torch::Tensor y) {
    int64_t total = y.numel(); int threads = 256;
    int blocks = (int)std::min<int64_t>((total + threads - 1) / threads, 65535);
    gelu_relu_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), total);
}

void launch_train_custom(torch::Tensor y, torch::Tensor gamma, torch::Tensor beta,
                         torch::Tensor running_mean, torch::Tensor running_var,
                         double momentum, double eps) {
    int B = (int)y.size(0);
    int C = (int)y.size(1);
    const int splits = 16;
    const int rows_per_split = (B + splits - 1) / splits;
    auto opts = y.options();
    auto psum = torch::empty({splits, C}, opts);
    auto psq = torch::empty({splits, C}, opts);
    auto mean = torch::empty({C}, opts);
    auto invstd = torch::empty({C}, opts);
    dim3 grid_stats((C + 31) / 32, splits);
    stats_warp32_kernel<<<grid_stats, 32, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), psum.data_ptr<float>(), psq.data_ptr<float>(), B, C, splits, rows_per_split);
    finalize_stats_kernel<<<(C + 31) / 32, 32, 0, at::cuda::getCurrentCUDAStream()>>>(
        psum.data_ptr<float>(), psq.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(), B, C, splits, (float)momentum, (float)eps);
    const int rows_per_chunk = 128;
    dim3 grid_apply((C + 31) / 32, (B + rows_per_chunk - 1) / rows_per_chunk);
    apply_train_bn_gelu_kernel<<<grid_apply, 32, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        B, C, rows_per_chunk);
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor linear_bias,
                            torch::Tensor bn_w, torch::Tensor bn_b,
                            torch::Tensor running_mean, torch::Tensor running_var,
                            bool training, double momentum, double eps) {
    auto y = at::linear(x, w, linear_bias);
    if (training && y.scalar_type() == at::kFloat && y.dim() == 2 && ((int)y.size(1) % 32 == 0)) {
        launch_train_custom(y, bn_w, bn_b, running_mean, running_var, momentum, eps);
    } else if (training) {
        y = at::batch_norm(y, bn_w, bn_b, running_mean, running_var, true, momentum, eps, true);
        launch_gelu(y);
    } else {
        launch_eval(y, bn_w, bn_b, running_mean, running_var, eps);
    }
    return y;
}
'''

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

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

    def forward(self, x):
        mom = self.batch_norm.momentum
        if mom is None:
            mom = 0.1
        return _ext.fused_forward(x, self.gemm.weight, self.gemm.bias,
                                  self.batch_norm.weight, self.batch_norm.bias,
                                  self.batch_norm.running_mean, self.batch_norm.running_var,
                                  self.training, float(mom), float(self.batch_norm.eps))