import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a GEMM (general matrix multiplication), applies scaling, 
    and then batch normalization.
    """
    def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.scale = nn.Parameter(torch.randn(scale_shape))
        self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum)

    def forward(self, x):
        x = self.gemm(x)
        x = x * self.scale
        x = self.bn(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
scale_shape = (out_features,)

def get_inputs():
    return [torch.rand(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features, scale_shape]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

void launch_bn_train_unscaled(torch::Tensor y, torch::Tensor bias, torch::Tensor scale,
                              torch::Tensor gamma, torch::Tensor beta,
                              torch::Tensor running_mean, torch::Tensor running_var,
                              double eps, double momentum);
void launch_bn_eval_unscaled(torch::Tensor y, torch::Tensor bias, torch::Tensor scale,
                             torch::Tensor gamma, torch::Tensor beta,
                             torch::Tensor running_mean, torch::Tensor running_var, double eps);

torch::Tensor forward_train(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor scale,
                            torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean,
                            torch::Tensor running_var, double eps, double momentum) {
    auto y = torch::mm(x, weight.t());
    launch_bn_train_unscaled(y, bias, scale, gamma, beta, running_mean, running_var, eps, momentum);
    return y;
}

torch::Tensor forward_eval(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor scale,
                           torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean,
                           torch::Tensor running_var, double eps) {
    auto y = torch::mm(x, weight.t());
    launch_bn_eval_unscaled(y, bias, scale, gamma, beta, running_mean, running_var, eps);
    return y;
}
"""

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

__global__ void bn_train_unscaled_kernel(float* __restrict__ y,
                                         const float* __restrict__ bias,
                                         const float* __restrict__ scale,
                                         const float* __restrict__ gamma,
                                         const float* __restrict__ beta,
                                         float* __restrict__ running_mean,
                                         float* __restrict__ running_var,
                                         int M, int N, float eps, float momentum) {
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int col = blockIdx.x * 32 + tx;
    __shared__ float red[8][32];
    __shared__ float mean_s[32];
    __shared__ float fac_s[32];
    __shared__ float beta_s[32];
    if (col >= N) return;

    float sum = 0.0f;
    for (int r = ty; r < M; r += 8) sum += y[(long long)r * N + col];
    red[ty][tx] = sum;
    __syncthreads();
    if (ty == 0) {
        float total = red[0][tx] + red[1][tx] + red[2][tx] + red[3][tx] + red[4][tx] + red[5][tx] + red[6][tx] + red[7][tx];
        mean_s[tx] = total / (float)M;
    }
    __syncthreads();
    const float mean = mean_s[tx];
    float vsum = 0.0f;
    for (int r = ty; r < M; r += 8) { float d = y[(long long)r * N + col] - mean; vsum += d * d; }
    red[ty][tx] = vsum;
    __syncthreads();
    if (ty == 0) {
        float vtotal = red[0][tx] + red[1][tx] + red[2][tx] + red[3][tx] + red[4][tx] + red[5][tx] + red[6][tx] + red[7][tx];
        float var = vtotal / (float)M;
        float sc = scale[col];
        float inv = 1.0f / sqrtf(sc * sc * var + eps);
        fac_s[tx] = sc * inv * gamma[col];
        beta_s[tx] = beta[col];
        float zmean = sc * (mean + bias[col]);
        float zvar_unbiased = (M > 1) ? (sc * sc * var * ((float)M / (float)(M - 1))) : 0.0f;
        running_mean[col] = running_mean[col] * (1.0f - momentum) + momentum * zmean;
        running_var[col] = running_var[col] * (1.0f - momentum) + momentum * zvar_unbiased;
    }
    __syncthreads();
    const float fac = fac_s[tx];
    const float be = beta_s[tx];
    for (int r = ty; r < M; r += 8) {
        long long idx = (long long)r * N + col;
        y[idx] = (y[idx] - mean) * fac + be;
    }
}

__global__ void bn_eval_unscaled_kernel(float* __restrict__ y,
                                        const float* __restrict__ bias,
                                        const float* __restrict__ scale,
                                        const float* __restrict__ gamma,
                                        const float* __restrict__ beta,
                                        const float* __restrict__ running_mean,
                                        const float* __restrict__ running_var,
                                        long long total, int N, float eps) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; i < total; i += stride) {
        int col = (int)(i % N);
        float v = (y[i] + bias[col]) * scale[col];
        y[i] = (v - running_mean[col]) * (1.0f / sqrtf(running_var[col] + eps)) * gamma[col] + beta[col];
    }
}

void launch_bn_train_unscaled(torch::Tensor y, torch::Tensor bias, torch::Tensor scale,
                              torch::Tensor gamma, torch::Tensor beta,
                              torch::Tensor running_mean, torch::Tensor running_var,
                              double eps, double momentum) {
    int M = (int)y.size(0), N = (int)y.size(1);
    dim3 block(32, 8), grid((N + 31) / 32);
    bn_train_unscaled_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(), M, N, (float)eps, (float)momentum);
}

void launch_bn_eval_unscaled(torch::Tensor y, torch::Tensor bias, torch::Tensor scale,
                             torch::Tensor gamma, torch::Tensor beta,
                             torch::Tensor running_mean, torch::Tensor running_var, double eps) {
    long long total = y.numel(); int N = (int)y.size(1); int threads = 256;
    int blocks = (int)((total + threads - 1) / threads); if (blocks > 65535) blocks = 65535;
    bn_eval_unscaled_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(), total, N, (float)eps);
}
"""

_ext = load_inline(name="gemm_scale_bn_ext_v3", cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=["forward_train", "forward_eval"],
                   extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.scale = nn.Parameter(torch.randn(scale_shape))
        self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum)
    def forward(self, x):
        if self.training:
            if self.bn.num_batches_tracked is not None:
                self.bn.num_batches_tracked.add_(1)
                mom = (1.0 / float(self.bn.num_batches_tracked.item())) if self.bn.momentum is None else float(self.bn.momentum)
            else:
                mom = float(self.bn.momentum if self.bn.momentum is not None else 0.0)
            return _ext.forward_train(x, self.gemm.weight, self.gemm.bias, self.scale, self.bn.weight, self.bn.bias,
                                      self.bn.running_mean, self.bn.running_var, float(self.bn.eps), mom)
        else:
            return _ext.forward_eval(x, self.gemm.weight, self.gemm.bias, self.scale, self.bn.weight, self.bn.bias,
                                     self.bn.running_mean, self.bn.running_var, float(self.bn.eps))