import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a convolution, applies Batch Normalization, and scales the output.
    """
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels)
        self.scaling_factor = scaling_factor

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = x * self.scaling_factor
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128
kernel_size = 3
scaling_factor = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scaling_factor]

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>
#include <c10/cuda/CUDAGuard.h>

void bn_scale_cuda(torch::Tensor y, torch::Tensor gamma, torch::Tensor beta,
                   torch::Tensor running_mean, torch::Tensor running_var,
                   torch::Tensor out, torch::Tensor a, torch::Tensor b,
                   torch::Tensor psum, torch::Tensor psq,
                   bool training, double momentum, double eps, double scaling);

torch::Tensor bn_scale_forward(torch::Tensor y, torch::Tensor gamma, torch::Tensor beta,
                               torch::Tensor running_mean, torch::Tensor running_var,
                               bool training, double momentum, double eps, double scaling) {
    TORCH_CHECK(y.is_cuda(), "y must be CUDA");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "only float32 supported");
    c10::cuda::CUDAGuard device_guard(y.device());
    y = y.contiguous();
    int64_t C = y.size(1);
    auto out = torch::empty_like(y);
    auto a = torch::empty({C}, y.options());
    auto b = torch::empty({C}, y.options());
    auto dopts = y.options().dtype(torch::kFloat64);
    torch::Tensor psum, psq;
    if (training) {
        psum = torch::empty({C, 32}, dopts);
        psq  = torch::empty({C, 32}, dopts);
    } else {
        psum = torch::empty({0}, dopts);
        psq  = torch::empty({0}, dopts);
    }
    bn_scale_cuda(y, gamma.contiguous(), beta.contiguous(), running_mean, running_var,
                  out, a, b, psum, psq, training, momentum, eps, scaling);
    return out;
}
'''

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

static constexpr int THREADS = 256;
static constexpr int PARTS = 32;

__global__ void partial_stats_kernel(const float* __restrict__ y,
                                     double* __restrict__ psum,
                                     double* __restrict__ psq,
                                     int N, int C, int H, int W) {
    int c = blockIdx.x;
    int part = blockIdx.y;
    int tid = threadIdx.x;
    long HW = (long)H * W;
    long M = (long)N * HW;
    double s = 0.0;
    double q = 0.0;
    for (long m = (long)part * blockDim.x + tid; m < M; m += (long)blockDim.x * gridDim.y) {
        long n = m / HW;
        long r = m - n * HW;
        float v = y[(n * C + c) * HW + r];
        s += (double)v;
        q += (double)v * (double)v;
    }
    extern __shared__ double sh[];
    double* ss = sh;
    double* qq = sh + blockDim.x;
    ss[tid] = s;
    qq[tid] = q;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ss[tid] += ss[tid + stride];
            qq[tid] += qq[tid + stride];
        }
        __syncthreads();
    }
    if (tid == 0) {
        psum[c * PARTS + part] = ss[0];
        psq [c * PARTS + part] = qq[0];
    }
}

__global__ void finalize_train_kernel(const double* __restrict__ psum,
                                      const double* __restrict__ psq,
                                      const float* __restrict__ gamma,
                                      const float* __restrict__ beta,
                                      float* __restrict__ running_mean,
                                      float* __restrict__ running_var,
                                      float* __restrict__ a,
                                      float* __restrict__ b,
                                      int C, long M, double momentum,
                                      double eps, double scaling) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    double s = 0.0, q = 0.0;
    #pragma unroll
    for (int p = 0; p < PARTS; ++p) {
        s += psum[c * PARTS + p];
        q += psq [c * PARTS + p];
    }
    double mean = s / (double)M;
    double var = q / (double)M - mean * mean;
    if (var < 0.0) var = 0.0;
    double inv = rsqrt(var + eps);
    double g = (double)gamma[c];
    double be = (double)beta[c];
    double sc = scaling;
    a[c] = (float)(g * inv * sc);
    b[c] = (float)((be - mean * g * inv) * sc);

    // Match BatchNorm's running-stat update (unbiased variance for running_var).
    double unb = (M > 1) ? var * ((double)M / (double)(M - 1)) : var;
    running_mean[c] = (float)(((1.0 - momentum) * (double)running_mean[c]) + momentum * mean);
    running_var[c]  = (float)(((1.0 - momentum) * (double)running_var[c])  + momentum * unb);
}

__global__ void coeff_eval_kernel(const float* __restrict__ gamma,
                                  const float* __restrict__ beta,
                                  const float* __restrict__ running_mean,
                                  const float* __restrict__ running_var,
                                  float* __restrict__ a,
                                  float* __restrict__ b,
                                  int C, double eps, double scaling) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    double inv = rsqrt((double)running_var[c] + eps);
    double g = (double)gamma[c];
    double sc = scaling;
    a[c] = (float)(g * inv * sc);
    b[c] = (float)((double)beta[c] * sc - (double)running_mean[c] * g * inv * sc);
}

__global__ void apply_affine_kernel(const float* __restrict__ y,
                                    const float* __restrict__ a,
                                    const float* __restrict__ b,
                                    float* __restrict__ out,
                                    long total, int C, int H, int W) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    long HW = (long)H * W;
    for (; idx < total; idx += stride) {
        int c = (int)((idx / HW) % C);
        out[idx] = y[idx] * a[c] + b[c];
    }
}

void bn_scale_cuda(torch::Tensor y, torch::Tensor gamma, torch::Tensor beta,
                   torch::Tensor running_mean, torch::Tensor running_var,
                   torch::Tensor out, torch::Tensor a, torch::Tensor b,
                   torch::Tensor psum, torch::Tensor psq,
                   bool training, double momentum, double eps, double scaling) {
    int N = (int)y.size(0);
    int C = (int)y.size(1);
    int H = (int)y.size(2);
    int W = (int)y.size(3);
    long total = (long)N * C * H * W;
    long M = (long)N * H * W;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (training) {
        dim3 grid(C, PARTS);
        partial_stats_kernel<<<grid, THREADS, 2 * THREADS * sizeof(double), stream>>>(
            y.data_ptr<float>(), psum.data_ptr<double>(), psq.data_ptr<double>(), N, C, H, W);
        int blocksC = (C + 127) / 128;
        finalize_train_kernel<<<blocksC, 128, 0, stream>>>(
            psum.data_ptr<double>(), psq.data_ptr<double>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), a.data_ptr<float>(), b.data_ptr<float>(),
            C, M, momentum, eps, scaling);
    } else {
        int blocksC = (C + 127) / 128;
        coeff_eval_kernel<<<blocksC, 128, 0, stream>>>(
            gamma.data_ptr<float>(), beta.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            a.data_ptr<float>(), b.data_ptr<float>(), C, eps, scaling);
    }
    int blocks = (int)((total + THREADS - 1) / THREADS);
    if (blocks > 65535) blocks = 65535;
    apply_affine_kernel<<<blocks, THREADS, 0, stream>>>(
        y.data_ptr<float>(), a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), total, C, H, W);
}
'''

_ext = load_inline(
    name="conv_bn_scale_fused_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["bn_scale_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    """Conv2d followed by BatchNorm2d and a scalar multiply; BN affine+scale is fused."""
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels)
        self.scaling_factor = scaling_factor

    def forward(self, x):
        y = self.conv(x)
        momentum = self.bn.momentum if self.bn.momentum is not None else 0.1
        if self.training and self.bn.track_running_stats:
            # Keep the observable BatchNorm buffer behavior consistent with PyTorch.
            self.bn.num_batches_tracked.add_(1)
            if self.bn.momentum is None:
                momentum = 1.0 / float(self.bn.num_batches_tracked.item())
        return _ext.bn_scale_forward(
            y, self.bn.weight, self.bn.bias, self.bn.running_mean, self.bn.running_var,
            bool(self.training), float(momentum), float(self.bn.eps), float(self.scaling_factor)
        )