import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Batch Normalization.
    """
    def __init__(self, num_features: int):
        """
        Initializes the BatchNorm layer.

        Args:
            num_features (int): Number of features in the input tensor.
        """
        super(Model, self).__init__()
        self.bn = nn.BatchNorm2d(num_features=num_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Batch Normalization to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, *).

        Returns:
            torch.Tensor: Output tensor with Batch Normalization applied, same shape as input.
        """
        return self.bn(x)

batch_size = 64
features = 64
dim1 = 512
dim2 = 512

def get_inputs():
    x = torch.rand(batch_size, features, dim1, dim2)
    return [x]

def get_init_inputs():
    return [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 bn_forward(torch::Tensor x,
                         torch::Tensor weight,
                         torch::Tensor bias,
                         torch::Tensor running_mean,
                         torch::Tensor running_var,
                         bool training);
"""

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

#ifndef __CUDACC_RTC__
#include <cstdio>
#endif

static __device__ __forceinline__ void warp_reduce_pair(float &a, float &b) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        a += __shfl_down_sync(0xffffffff, a, offset);
        b += __shfl_down_sync(0xffffffff, b, offset);
    }
}

__global__ void reduce_bn_kernel(const float* __restrict__ x,
                                 float* __restrict__ partial,
                                 long long N, int C, long long HW,
                                 int chunks_per_n, long long chunk_size,
                                 int chunks_per_c) {
    int bid = blockIdx.x;
    int c = bid / chunks_per_c;
    int k = bid - c * chunks_per_c;
    long long n = k / chunks_per_n;
    int seg = k - (int)(n * chunks_per_n);
    long long start = (long long)seg * chunk_size;
    long long len = chunk_size;
    if (start + len > HW) len = HW - start;
    const float* ptr = x + ((n * C + c) * HW + start);

    float sum = 0.0f;
    float sum2 = 0.0f;
    for (long long i = threadIdx.x; i < len; i += blockDim.x) {
        float v = ptr[i];
        sum += v;
        sum2 = fmaf(v, v, sum2);
    }

    warp_reduce_pair(sum, sum2);

    __shared__ float smem_sum[32];
    __shared__ float smem_sum2[32];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if (lane == 0) {
        smem_sum[warp] = sum;
        smem_sum2[warp] = sum2;
    }
    __syncthreads();

    if (warp == 0) {
        sum = (lane < (blockDim.x >> 5)) ? smem_sum[lane] : 0.0f;
        sum2 = (lane < (blockDim.x >> 5)) ? smem_sum2[lane] : 0.0f;
        warp_reduce_pair(sum, sum2);
        if (lane == 0) {
            long long out = ((long long)c * chunks_per_c + k) * 2LL;
            partial[out] = sum;
            partial[out + 1] = sum2;
        }
    }
}

__global__ void finalize_bn_kernel(const float* __restrict__ partial,
                                   float* __restrict__ stats,
                                   float* __restrict__ running_mean,
                                   float* __restrict__ running_var,
                                   int C, int chunks_per_c, long long M,
                                   float eps, float momentum) {
    int c = blockIdx.x;
    float sum = 0.0f;
    float sum2 = 0.0f;
    for (int k = threadIdx.x; k < chunks_per_c; k += blockDim.x) {
        long long idx = ((long long)c * chunks_per_c + k) * 2LL;
        sum += partial[idx];
        sum2 += partial[idx + 1];
    }

    warp_reduce_pair(sum, sum2);

    __shared__ float smem_sum[32];
    __shared__ float smem_sum2[32];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if (lane == 0) {
        smem_sum[warp] = sum;
        smem_sum2[warp] = sum2;
    }
    __syncthreads();

    if (warp == 0) {
        sum = (lane < (blockDim.x >> 5)) ? smem_sum[lane] : 0.0f;
        sum2 = (lane < (blockDim.x >> 5)) ? smem_sum2[lane] : 0.0f;
        warp_reduce_pair(sum, sum2);
        if (lane == 0) {
            float invM = 1.0f / (float)M;
            float mean = sum * invM;
            float var = sum2 * invM - mean * mean;
            var = fmaxf(var, 0.0f);
            float invstd = rsqrtf(var + eps);
            stats[2 * c] = mean;
            stats[2 * c + 1] = invstd;

            float old_m = running_mean[c];
            float old_v = running_var[c];
            float unbiased = var * ((float)M / (float)(M - 1));
            running_mean[c] = old_m * (1.0f - momentum) + mean * momentum;
            running_var[c] = old_v * (1.0f - momentum) + unbiased * momentum;
        }
    }
}

__global__ void norm_train_kernel(const float* __restrict__ x,
                                  const float* __restrict__ weight,
                                  const float* __restrict__ bias,
                                  const float* __restrict__ stats,
                                  float* __restrict__ y,
                                  long long N, int C, long long HW) {
    int bc = blockIdx.x;
    int c = bc % C;
    long long n = bc / C;
    long long base = (n * C + c) * HW;
    float mean = stats[2 * c];
    float invstd = stats[2 * c + 1];
    float gamma = weight[c];
    float beta = bias[c];
    float scale = gamma * invstd;
    float shift = beta - mean * scale;
    for (long long s = threadIdx.x; s < HW; s += blockDim.x) {
        y[base + s] = fmaf(x[base + s], scale, shift);
    }
}

__global__ void norm_eval_kernel(const float* __restrict__ x,
                                 const float* __restrict__ weight,
                                 const float* __restrict__ bias,
                                 const float* __restrict__ running_mean,
                                 const float* __restrict__ running_var,
                                 float* __restrict__ y,
                                 long long N, int C, long long HW, float eps) {
    int bc = blockIdx.x;
    int c = bc % C;
    long long n = bc / C;
    long long base = (n * C + c) * HW;
    float invstd = rsqrtf(running_var[c] + eps);
    float scale = weight[c] * invstd;
    float shift = bias[c] - running_mean[c] * scale;
    for (long long s = threadIdx.x; s < HW; s += blockDim.x) {
        y[base + s] = fmaf(x[base + s], scale, shift);
    }
}

torch::Tensor bn_forward(torch::Tensor x,
                         torch::Tensor weight,
                         torch::Tensor bias,
                         torch::Tensor running_mean,
                         torch::Tensor running_var,
                         bool training) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.dim() == 4, "expected NCHW tensor");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const long long N = x.size(0);
    const int C = (int)x.size(1);
    const long long H = x.size(2);
    const long long W = x.size(3);
    const long long HW = H * W;
    auto y = torch::empty_like(x);

    const int threads = 256;
    const float eps = 1.0e-5f;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (!training) {
        dim3 grid((unsigned int)(N * C));
        norm_eval_kernel<<<grid, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            y.data_ptr<float>(), N, C, HW, eps);
        return y;
    }

    // 64k spatial elements per partial block.  For the benchmark shape this is
    // exactly four chunks per image/channel, giving 16k independent reductions.
    const long long chunk_size = 65536LL;
    const int chunks_per_n = (int)((HW + chunk_size - 1) / chunk_size);
    const int chunks_per_c = (int)(N * (long long)chunks_per_n);

    auto partial = torch::empty({C, chunks_per_c, 2}, x.options());
    auto stats = torch::empty({C, 2}, x.options());

    dim3 red_grid((unsigned int)(C * chunks_per_c));
    reduce_bn_kernel<<<red_grid, threads, 0, stream>>>(
        x.data_ptr<float>(), partial.data_ptr<float>(), N, C, HW,
        chunks_per_n, chunk_size, chunks_per_c);

    finalize_bn_kernel<<<C, threads, 0, stream>>>(
        partial.data_ptr<float>(), stats.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        C, chunks_per_c, N * HW, eps, 0.1f);

    dim3 norm_grid((unsigned int)(N * C));
    norm_train_kernel<<<norm_grid, threads, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        stats.data_ptr<float>(), y.data_ptr<float>(), N, C, HW);

    return y;
}
"""

_bn_ext = load_inline(
    name="bn2d_forward_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["bn_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_features: int):
        super(ModelNew, self).__init__()
        self.bn = nn.BatchNorm2d(num_features=num_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.bn.training:
            # Match BatchNorm2d's bookkeeping for the default track_running_stats=True.
            self.bn.num_batches_tracked.add_(1)
        return _bn_ext.bn_forward(x, self.bn.weight, self.bn.bias,
                                  self.bn.running_mean, self.bn.running_var,
                                  bool(self.bn.training))