import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, scales the result, and applies 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 = 16384
in_features = 4096
out_features = 4096
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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

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

#define COL_TILE 32
#define TY 8
#define ROWS_PER_BLOCK 256

__global__ void partial_sums_kernel(const float* __restrict__ y,
                                    float* __restrict__ psum,
                                    float* __restrict__ psq,
                                    int N, int C) {
    __shared__ float sh_sum[TY][COL_TILE];
    __shared__ float sh_sq[TY][COL_TILE];
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int c = blockIdx.x * COL_TILE + tx;
    int r0 = blockIdx.y * ROWS_PER_BLOCK;
    int rend = r0 + ROWS_PER_BLOCK;
    if (rend > N) rend = N;

    float s = 0.0f;
    float q = 0.0f;
    if (c < C) {
        for (int r = r0 + ty; r < rend; r += TY) {
            float v = y[(long long)r * C + c];
            s += v;
            q += v * v;
        }
    }
    sh_sum[ty][tx] = s;
    sh_sq[ty][tx] = q;
    __syncthreads();

    if (ty < 4) { sh_sum[ty][tx] += sh_sum[ty+4][tx]; sh_sq[ty][tx] += sh_sq[ty+4][tx]; }
    __syncthreads();
    if (ty < 2) { sh_sum[ty][tx] += sh_sum[ty+2][tx]; sh_sq[ty][tx] += sh_sq[ty+2][tx]; }
    __syncthreads();
    if (ty == 0) {
        s = sh_sum[0][tx] + sh_sum[1][tx];
        q = sh_sq[0][tx] + sh_sq[1][tx];
        if (c < C) {
            psum[(long long)blockIdx.y * C + c] = s;
            psq[(long long)blockIdx.y * C + c] = q;
        }
    }
}

__global__ void final_stats_kernel(const float* __restrict__ psum,
                                   const float* __restrict__ psq,
                                   const float* __restrict__ scale,
                                   float* __restrict__ running_mean,
                                   float* __restrict__ running_var,
                                   float* __restrict__ mean,
                                   float* __restrict__ invstd,
                                   int N, int C, int NB,
                                   float momentum, float eps) {
    int c = blockIdx.x;
    int t = threadIdx.x;
    __shared__ float ss[256];
    __shared__ float qq[256];
    float s = 0.0f;
    float q = 0.0f;
    for (int b = t; b < NB; b += blockDim.x) {
        s += psum[(long long)b * C + c];
        q += psq[(long long)b * C + c];
    }
    ss[t] = s;
    qq[t] = q;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (t < stride) {
            ss[t] += ss[t + stride];
            qq[t] += qq[t + stride];
        }
        __syncthreads();
    }
    if (t == 0) {
        float invN = 1.0f / (float)N;
        float mz = ss[0] * invN;
        float vz = qq[0] * invN - mz * mz;
        if (vz < 0.0f) vz = 0.0f;
        float sc = scale[c];
        float m = mz * sc;
        float v = vz * sc * sc;
        mean[c] = m;
        invstd[c] = 1.0f / sqrtf(v + eps);
        float unbiased = v;
        if (N > 1) unbiased = v * ((float)N / (float)(N - 1));
        running_mean[c] = running_mean[c] * (1.0f - momentum) + momentum * m;
        running_var[c] = running_var[c] * (1.0f - momentum) + momentum * unbiased;
    }
}

__global__ void train_apply_kernel(float* __restrict__ y,
                                   const float* __restrict__ scale,
                                   const float* __restrict__ gamma,
                                   const float* __restrict__ beta,
                                   const float* __restrict__ mean,
                                   const float* __restrict__ invstd,
                                   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 v = y[idx] * scale[c];
        y[idx] = (v - mean[c]) * invstd[c] * gamma[c] + beta[c];
    }
}

__global__ void eval_apply_kernel(float* __restrict__ y,
                                  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 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 v = y[idx] * scale[c];
        float inv = 1.0f / sqrtf(running_var[c] + eps);
        y[idx] = (v - running_mean[c]) * inv * gamma[c] + beta[c];
    }
}

torch::Tensor bn_train_inplace(torch::Tensor y,
                               torch::Tensor scale,
                               torch::Tensor gamma,
                               torch::Tensor beta,
                               torch::Tensor running_mean,
                               torch::Tensor running_var,
                               double momentum,
                               double eps) {
    const int N = (int)y.size(0);
    const int C = (int)y.size(1);
    const int NB = (N + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK;
    auto opts = y.options();
    torch::Tensor psum = torch::empty({NB, C}, opts);
    torch::Tensor psq = torch::empty({NB, C}, opts);
    torch::Tensor mean = torch::empty({C}, opts);
    torch::Tensor invstd = torch::empty({C}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 block(COL_TILE, TY);
    dim3 grid((C + COL_TILE - 1) / COL_TILE, NB);
    partial_sums_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), psum.data_ptr<float>(), psq.data_ptr<float>(), N, C);
    final_stats_kernel<<<C, 256, 0, stream>>>(psum.data_ptr<float>(), psq.data_ptr<float>(), scale.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        N, C, NB, (float)momentum, (float)eps);
    long long total = (long long)N * C;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    train_apply_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), total, C);
    return y;
}

torch::Tensor bn_eval_inplace(torch::Tensor y,
                              torch::Tensor scale,
                              torch::Tensor gamma,
                              torch::Tensor beta,
                              torch::Tensor running_mean,
                              torch::Tensor running_var,
                              double eps) {
    const int N = (int)y.size(0);
    const int C = (int)y.size(1);
    long long total = (long long)N * C;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    eval_apply_kernel<<<blocks, threads, 0, stream>>>(y.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, C, (float)eps);
    return y;
}
'''

_ext = load_inline(
    name='fused_scale_bn_inline_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['bn_train_inplace', 'bn_eval_inplace'],
    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):
        y = F.linear(x, self.gemm.weight, self.gemm.bias)
        if self.bn.training:
            if self.bn.track_running_stats:
                self.bn.num_batches_tracked.add_(1)
                if self.bn.momentum is None:
                    mom = 1.0 / float(self.bn.num_batches_tracked.item())
                else:
                    mom = float(self.bn.momentum)
            else:
                mom = 0.0 if self.bn.momentum is None else float(self.bn.momentum)
            return _ext.bn_train_inplace(y, self.scale, self.bn.weight, self.bn.bias,
                                         self.bn.running_mean, self.bn.running_var,
                                         mom, float(self.bn.eps))
        else:
            return _ext.bn_eval_inplace(y, self.scale, self.bn.weight, self.bn.bias,
                                        self.bn.running_mean, self.bn.running_var,
                                        float(self.bn.eps))