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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor bn_gelu_relu_forward(torch::Tensor x,
                                   torch::Tensor gamma,
                                   torch::Tensor beta,
                                   torch::Tensor running_mean,
                                   torch::Tensor running_var,
                                   bool training,
                                   double eps,
                                   double momentum);
"""

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

#define COL_TILE 32
#define ROW_TILE 16
#define Y_THREADS 8

__global__ void partial_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ partial,
                                   float* __restrict__ partial2,
                                   int N, int C, int tiles) {
    int col0 = blockIdx.x * COL_TILE;
    int tile = blockIdx.y;
    int tx = threadIdx.x;  // 0..31, feature within tile
    int ty = threadIdx.y;  // 0..7, row reducer
    int c = col0 + tx;
    int row0 = tile * ROW_TILE;

    float s = 0.0f;
    float q = 0.0f;
    if (c < C) {
        #pragma unroll
        for (int r = ty; r < ROW_TILE; r += Y_THREADS) {
            int n = row0 + r;
            if (n < N) {
                float v = x[(long long)n * C + c];
                s += v;
                q += v * v;
            }
        }
    }

    __shared__ float sh[Y_THREADS][COL_TILE];
    __shared__ float sh2[Y_THREADS][COL_TILE];
    sh[ty][tx] = s;
    sh2[ty][tx] = q;
    __syncthreads();

    if (ty == 0 && c < C) {
        float ss = 0.0f;
        float qq = 0.0f;
        #pragma unroll
        for (int k = 0; k < Y_THREADS; ++k) {
            ss += sh[k][tx];
            qq += sh2[k][tx];
        }
        partial[(long long)tile * C + c] = ss;
        partial2[(long long)tile * C + c] = qq;
    }
}

__global__ void finalize_stats_kernel(const float* __restrict__ partial,
                                      const float* __restrict__ partial2,
                                      float* __restrict__ mean,
                                      float* __restrict__ invstd,
                                      float* __restrict__ running_mean,
                                      float* __restrict__ running_var,
                                      int N, int C, int tiles,
                                      float eps, float momentum) {
    int c = blockIdx.x;
    int tid = threadIdx.x;
    float s = 0.0f;
    float q = 0.0f;
    for (int t = tid; t < tiles; t += blockDim.x) {
        float a = partial[(long long)t * C + c];
        float b = partial2[(long long)t * C + c];
        s += a;
        q += b;
    }

    __shared__ float ss[256];
    __shared__ float qq[256];
    ss[tid] = s;
    qq[tid] = q;
    __syncthreads();

    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ss[tid] += ss[tid + stride];
            qq[tid] += qq[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        float n_inv = 1.0f / (float)N;
        float m = ss[0] * n_inv;
        float var = qq[0] * n_inv - m * m;
        var = var < 0.0f ? 0.0f : var;
        float inv = rsqrtf(var + eps);
        mean[c] = m;
        invstd[c] = inv;
        // PyTorch BatchNorm updates running_var with the unbiased variance.
        float unb = (N > 1) ? (var * ((float)N / (float)(N - 1))) : var;
        running_mean[c] = running_mean[c] * (1.0f - momentum) + m * momentum;
        running_var[c] = running_var[c] * (1.0f - momentum) + unb * momentum;
    }
}

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

__global__ void norm_act_train_kernel(float* __restrict__ x,
                                      const float* __restrict__ mean,
                                      const float* __restrict__ invstd,
                                      const float* __restrict__ gamma,
                                      const float* __restrict__ beta,
                                      long long total, int C) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = (int)(idx % C);
        float z = (x[idx] - mean[c]) * invstd[c];
        z = z * gamma[c] + beta[c];
        x[idx] = gelu_relu_exact(z);
    }
}

__global__ void norm_act_eval_kernel(float* __restrict__ x,
                                     const float* __restrict__ running_mean,
                                     const float* __restrict__ running_var,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     long long total, int C, float eps) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = (int)(idx % C);
        float z = (x[idx] - running_mean[c]) * rsqrtf(running_var[c] + eps);
        z = z * gamma[c] + beta[c];
        x[idx] = gelu_relu_exact(z);
    }
}

torch::Tensor bn_gelu_relu_forward(torch::Tensor x,
                                   torch::Tensor gamma,
                                   torch::Tensor beta,
                                   torch::Tensor running_mean,
                                   torch::Tensor running_var,
                                   bool training,
                                   double eps,
                                   double momentum) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const long long total = (long long)N * C;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;

    if (training) {
        int tiles = (N + ROW_TILE - 1) / ROW_TILE;
        auto opts = x.options();
        torch::Tensor partial = torch::empty({tiles, C}, opts);
        torch::Tensor partial2 = torch::empty({tiles, C}, opts);
        torch::Tensor mean = torch::empty({C}, opts);
        torch::Tensor invstd = torch::empty({C}, opts);

        dim3 block(COL_TILE, Y_THREADS);
        dim3 grid((C + COL_TILE - 1) / COL_TILE, tiles);
        partial_sum_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), partial.data_ptr<float>(), partial2.data_ptr<float>(), N, C, tiles);
        finalize_stats_kernel<<<C, 256, 0, stream>>>(
            partial.data_ptr<float>(), partial2.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), N, C, tiles, (float)eps, (float)momentum);
        norm_act_train_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), total, C);
    } else {
        norm_act_eval_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), total, C, (float)eps);
    }
    return x;
}
"""

_ext = load_inline(
    name="bn_gelu_relu_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["bn_gelu_relu_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """GEMM + BatchNorm1d + GELU + ReLU, with fused CUDA post-processing."""
    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):
        y = F.linear(x, self.gemm.weight, self.gemm.bias)
        bn = self.batch_norm
        if self.training and bn.track_running_stats:
            # Match BatchNorm's bookkeeping for fixed momentum (the default used here).
            bn.num_batches_tracked.add_(1)
        return _ext.bn_gelu_relu_forward(
            y,
            bn.weight,
            bn.bias,
            bn.running_mean,
            bn.running_var,
            bool(self.training),
            float(bn.eps),
            float(bn.momentum if bn.momentum is not None else 0.0),
        )