import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a matrix multiplication (Gemm), Batch Normalization, scaling, and Softmax.
    """
    def __init__(self, in_features, out_features, bn_eps=1e-5, bn_momentum=0.1, scale_shape=(1,)):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(out_features, eps=bn_eps, momentum=bn_momentum)
        self.scale = nn.Parameter(torch.ones(scale_shape))
        self.softmax = nn.Softmax(dim=1)

    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.bn(x)
        x = self.scale * x
        x = self.softmax(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
bn_eps = 1e-5
bn_momentum = 0.1
scale_shape = (1,)

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

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

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

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

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

torch::Tensor bn_softmax_forward(torch::Tensor x,
                                  torch::Tensor gamma,
                                  torch::Tensor beta,
                                  torch::Tensor running_mean,
                                  torch::Tensor running_var,
                                  torch::Tensor scale,
                                  double eps,
                                  double momentum,
                                  bool training) {
    return bn_softmax_cuda(x, gamma, beta, running_mean, running_var, scale, eps, momentum, training);
}
'''

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

__global__ void bn_stats_kernel(const float* __restrict__ x,
                                float* __restrict__ mean,
                                float* __restrict__ invstd,
                                float* __restrict__ running_mean,
                                float* __restrict__ running_var,
                                int N, int C, float eps, float momentum) {
    int tx = threadIdx.x;   // column inside tile, 0..31
    int ty = threadIdx.y;   // reducer lane over rows, 0..7
    int col = blockIdx.x * 32 + tx;

    float s = 0.0f;
    float q = 0.0f;
    if (col < C) {
        for (int r = ty; r < N; r += 8) {
            float v = x[(long long)r * C + col];
            s += v;
            q += v * v;
        }
    }

    __shared__ float sh_s[8][32];
    __shared__ float sh_q[8][32];
    sh_s[ty][tx] = s;
    sh_q[ty][tx] = q;
    __syncthreads();

    if (ty == 0 && col < C) {
        float sum = 0.0f;
        float ssq = 0.0f;
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            sum += sh_s[k][tx];
            ssq += sh_q[k][tx];
        }
        float m = sum / (float)N;
        float var = ssq / (float)N - m * m;
        var = var < 0.0f ? 0.0f : var;
        mean[col] = m;
        invstd[col] = rsqrtf(var + eps);

        float unbiased = var * ((float)N / (float)(N - 1));
        running_mean[col] = running_mean[col] * (1.0f - momentum) + momentum * m;
        running_var[col]  = running_var[col]  * (1.0f - momentum) + momentum * unbiased;
    }
}

__device__ __forceinline__ float block_reduce_max(float v, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = v;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            float other = smem[tid + stride];
            smem[tid] = smem[tid] > other ? smem[tid] : other;
        }
        __syncthreads();
    }
    return smem[0];
}

__device__ __forceinline__ float block_reduce_sum(float v, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = v;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) smem[tid] += smem[tid + stride];
        __syncthreads();
    }
    return smem[0];
}

__global__ void softmax_train_kernel(const float* __restrict__ x,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     const float* __restrict__ mean,
                                     const float* __restrict__ invstd,
                                     const float* __restrict__ scale,
                                     float* __restrict__ out,
                                     int N, int C) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float sc = scale[0];
    extern __shared__ float smem[];

    float mx = -FLT_MAX;
    long long base = (long long)row * C;
    for (int c = tid; c < C; c += blockDim.x) {
        float z = ((x[base + c] - mean[c]) * invstd[c] * gamma[c] + beta[c]) * sc;
        mx = mx > z ? mx : z;
    }
    mx = block_reduce_max(mx, smem);

    float ss = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) {
        float z = ((x[base + c] - mean[c]) * invstd[c] * gamma[c] + beta[c]) * sc;
        ss += __expf(z - mx);
    }
    ss = block_reduce_sum(ss, smem);
    float inv = 1.0f / ss;

    for (int c = tid; c < C; c += blockDim.x) {
        float z = ((x[base + c] - mean[c]) * invstd[c] * gamma[c] + beta[c]) * sc;
        out[base + c] = __expf(z - mx) * inv;
    }
}

__global__ void softmax_eval_kernel(const float* __restrict__ x,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ beta,
                                    const float* __restrict__ running_mean,
                                    const float* __restrict__ running_var,
                                    const float* __restrict__ scale,
                                    float* __restrict__ out,
                                    int N, int C, float eps) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float sc = scale[0];
    extern __shared__ float smem[];

    float mx = -FLT_MAX;
    long long base = (long long)row * C;
    for (int c = tid; c < C; c += blockDim.x) {
        float z = ((x[base + c] - running_mean[c]) * rsqrtf(running_var[c] + eps) * gamma[c] + beta[c]) * sc;
        mx = mx > z ? mx : z;
    }
    mx = block_reduce_max(mx, smem);

    float ss = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) {
        float z = ((x[base + c] - running_mean[c]) * rsqrtf(running_var[c] + eps) * gamma[c] + beta[c]) * sc;
        ss += __expf(z - mx);
    }
    ss = block_reduce_sum(ss, smem);
    float inv = 1.0f / ss;

    for (int c = tid; c < C; c += blockDim.x) {
        float z = ((x[base + c] - running_mean[c]) * rsqrtf(running_var[c] + eps) * gamma[c] + beta[c]) * sc;
        out[base + c] = __expf(z - mx) * inv;
    }
}

torch::Tensor bn_softmax_cuda(torch::Tensor x,
                              torch::Tensor gamma,
                              torch::Tensor beta,
                              torch::Tensor running_mean,
                              torch::Tensor running_var,
                              torch::Tensor scale,
                              double eps,
                              double momentum,
                              bool training) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    auto out = torch::empty_like(x);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int threads = 256;
    const size_t shmem = threads * sizeof(float);

    if (training) {
        auto mean = torch::empty({C}, x.options());
        auto invstd = torch::empty({C}, x.options());
        dim3 block_stats(32, 8);
        dim3 grid_stats((C + 31) / 32);
        bn_stats_kernel<<<grid_stats, block_stats, 0, stream>>>(
            x.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            N, C, (float)eps, (float)momentum);
        softmax_train_kernel<<<N, threads, shmem, stream>>>(
            x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
            mean.data_ptr<float>(), invstd.data_ptr<float>(), scale.data_ptr<float>(),
            out.data_ptr<float>(), N, C);
    } else {
        softmax_eval_kernel<<<N, threads, shmem, stream>>>(
            x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), scale.data_ptr<float>(),
            out.data_ptr<float>(), N, C, (float)eps);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

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

class ModelNew(nn.Module):
    """Fused replacement for Linear -> BatchNorm1d -> scalar scale -> Softmax."""
    def __init__(self, in_features, out_features, bn_eps=1e-5, bn_momentum=0.1, scale_shape=(1,)):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(out_features, eps=bn_eps, momentum=bn_momentum)
        self.scale = nn.Parameter(torch.ones(scale_shape))

    def forward(self, x):
        y = torch.nn.functional.linear(x, self.gemm.weight, self.gemm.bias)
        training = bool(self.training and self.bn.training and self.bn.track_running_stats)
        if training:
            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 = float(self.bn.momentum if self.bn.momentum is not None else 0.0)
        return _ext.bn_softmax_forward(y, self.bn.weight, self.bn.bias,
                                       self.bn.running_mean, self.bn.running_var,
                                       self.scale, float(self.bn.eps), mom, training)