import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, scales the output, applies batch normalization, 
    and then performs global average pooling. 
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 16, 32, 32
kernel_size = 5
scale_factor = 2.0

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

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

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

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

torch::Tensor fused_forward(torch::Tensor x,
                            torch::Tensor weight,
                            torch::Tensor bias,
                            torch::Tensor bn_weight,
                            torch::Tensor bn_bias,
                            torch::Tensor running_mean,
                            torch::Tensor running_var,
                            double scale,
                            double eps,
                            bool training);
"""

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

__global__ void input_sum_kernel(const float* __restrict__ x, float* __restrict__ in_sum,
                                 int N, int Cin, int D, int H, int W) {
    extern __shared__ float sdata[];
    int idx = blockIdx.x;
    int n = idx / Cin;
    int c = idx - n * Cin;
    int S = D * H * W;
    float acc = 0.0f;
    const float* base = x + ((n * Cin + c) * S);
    for (int i = threadIdx.x; i < S; i += blockDim.x) acc += base[i];
    sdata[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) sdata[threadIdx.x] += sdata[threadIdx.x + stride];
        __syncthreads();
    }
    if (threadIdx.x == 0) in_sum[idx] = sdata[0];
}

__global__ void weight_sum_kernel(const float* __restrict__ w, float* __restrict__ w_sum,
                                  int Cin, int Cout, int Kd, int Kh, int Kw) {
    extern __shared__ float sdata[];
    int idx = blockIdx.x;
    int ic = idx / Cout;
    int oc = idx - ic * Cout;
    int K = Kd * Kh * Kw;
    const float* base = w + ((ic * Cout + oc) * K);
    float acc = 0.0f;
    for (int i = threadIdx.x; i < K; i += blockDim.x) acc += base[i];
    sdata[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) sdata[threadIdx.x] += sdata[threadIdx.x + stride];
        __syncthreads();
    }
    if (threadIdx.x == 0) w_sum[idx] = sdata[0];
}

__global__ void eval_out_kernel(const float* __restrict__ in_sum,
                                const float* __restrict__ w_sum,
                                const float* __restrict__ bias,
                                const float* __restrict__ gamma,
                                const float* __restrict__ beta,
                                const float* __restrict__ rmean,
                                const float* __restrict__ rvar,
                                float* __restrict__ out,
                                int N, int Cin, int Cout, int Vout,
                                float scale, float eps) {
    extern __shared__ float sdata[];
    int bid = blockIdx.x;
    int n = bid / Cout;
    int oc = bid - n * Cout;
    float acc = 0.0f;
    for (int ic = threadIdx.x; ic < Cin; ic += blockDim.x) {
        acc += in_sum[n * Cin + ic] * w_sum[ic * Cout + oc];
    }
    sdata[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) sdata[threadIdx.x] += sdata[threadIdx.x + stride];
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        float conv_avg = (sdata[0] + bias[oc] * (float)Vout) / (float)Vout;
        float v = conv_avg * scale;
        float y = (v - rmean[oc]) * rsqrtf(rvar[oc] + eps) * gamma[oc] + beta[oc];
        out[n * Cout + oc] = y;
    }
}

__global__ void conv_sum_kernel(const float* __restrict__ in_sum,
                                const float* __restrict__ w_sum,
                                const float* __restrict__ bias,
                                float* __restrict__ conv_sum,
                                int N, int Cin, int Cout, int Vout) {
    extern __shared__ float sdata[];
    int bid = blockIdx.x;
    int n = bid / Cout;
    int oc = bid - n * Cout;
    float acc = 0.0f;
    for (int ic = threadIdx.x; ic < Cin; ic += blockDim.x) {
        acc += in_sum[n * Cin + ic] * w_sum[ic * Cout + oc];
    }
    sdata[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) sdata[threadIdx.x] += sdata[threadIdx.x + stride];
        __syncthreads();
    }
    if (threadIdx.x == 0) conv_sum[n * Cout + oc] = sdata[0] + bias[oc] * (float)Vout;
}

__global__ void sumsq_kernel(const float* __restrict__ x,
                             const float* __restrict__ w,
                             const float* __restrict__ bias,
                             float* __restrict__ sumsq,
                             int N, int Cin, int Cout,
                             int D, int H, int Wd,
                             int Kd, int Kh, int Kw,
                             int OD, int OH, int OW,
                             float scale, int blocks_per_oc) {
    int global_bid = blockIdx.x;
    int oc = global_bid / blocks_per_oc;
    int local_bid = global_bid - oc * blocks_per_oc;
    int M = N * OD * OH * OW;
    float acc2 = 0.0f;
    int stride = blockDim.x * blocks_per_oc;
    for (int m = local_bid * blockDim.x + threadIdx.x; m < M; m += stride) {
        int t = m;
        int ow = t % OW; t /= OW;
        int oh = t % OH; t /= OH;
        int od = t % OD; t /= OD;
        int n = t;
        float val = bias[oc];
        int kd0 = max(0, od - (D - 1));
        int kd1 = min(Kd - 1, od);
        int kh0 = max(0, oh - (H - 1));
        int kh1 = min(Kh - 1, oh);
        int kw0 = max(0, ow - (Wd - 1));
        int kw1 = min(Kw - 1, ow);
        for (int ic = 0; ic < Cin; ++ic) {
            for (int kd = kd0; kd <= kd1; ++kd) {
                int id = od - kd;
                for (int kh = kh0; kh <= kh1; ++kh) {
                    int ih = oh - kh;
                    const float* xb = x + (((n * Cin + ic) * D + id) * H + ih) * Wd;
                    const float* wb = w + (((ic * Cout + oc) * Kd + kd) * Kh + kh) * Kw;
                    for (int kw = kw0; kw <= kw1; ++kw) {
                        val += xb[ow - kw] * wb[kw];
                    }
                }
            }
        }
        float sv = val * scale;
        acc2 += sv * sv;
    }
    extern __shared__ float sdata[];
    sdata[threadIdx.x] = acc2;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) atomicAdd(sumsq + oc, sdata[0]);
}

__global__ void train_out_kernel(const float* __restrict__ conv_sum,
                                 const float* __restrict__ sumsq,
                                 const float* __restrict__ gamma,
                                 const float* __restrict__ beta,
                                 float* __restrict__ out,
                                 int N, int Cout, int Vout,
                                 float scale, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * Cout;
    if (idx >= total) return;
    int n = idx / Cout;
    int oc = idx - n * Cout;
    float tsum = 0.0f;
    for (int i = 0; i < N; ++i) tsum += conv_sum[i * Cout + oc];
    float invM = 1.0f / (float)(N * Vout);
    float mean = tsum * scale * invM;
    float var = sumsq[oc] * invM - mean * mean;
    if (var < 0.0f) var = 0.0f;
    float avg = conv_sum[n * Cout + oc] * scale / (float)Vout;
    out[idx] = (avg - mean) * rsqrtf(var + eps) * gamma[oc] + beta[oc];
}

torch::Tensor fused_forward(torch::Tensor x,
                            torch::Tensor weight,
                            torch::Tensor bias,
                            torch::Tensor bn_weight,
                            torch::Tensor bn_bias,
                            torch::Tensor running_mean,
                            torch::Tensor running_var,
                            double scale_d,
                            double eps_d,
                            bool training) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    int N = (int)x.size(0);
    int Cin = (int)x.size(1);
    int D = (int)x.size(2);
    int H = (int)x.size(3);
    int Wd = (int)x.size(4);
    int Cout = (int)weight.size(1);
    int Kd = (int)weight.size(2);
    int Kh = (int)weight.size(3);
    int Kw = (int)weight.size(4);
    int OD = D + Kd - 1;
    int OH = H + Kh - 1;
    int OW = Wd + Kw - 1;
    int Vout = OD * OH * OW;
    auto opts = x.options();
    auto in_sum = torch::empty({N, Cin}, opts);
    auto w_sum = torch::empty({Cin, Cout}, opts);
    auto out = torch::empty({N, Cout, 1, 1, 1}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int T = 256;
    input_sum_kernel<<<N * Cin, T, T * sizeof(float), stream>>>(x.data_ptr<float>(), in_sum.data_ptr<float>(), N, Cin, D, H, Wd);
    weight_sum_kernel<<<Cin * Cout, T, T * sizeof(float), stream>>>(weight.data_ptr<float>(), w_sum.data_ptr<float>(), Cin, Cout, Kd, Kh, Kw);
    float scale = (float)scale_d;
    float eps = (float)eps_d;
    if (!training) {
        eval_out_kernel<<<N * Cout, T, T * sizeof(float), stream>>>(
            in_sum.data_ptr<float>(), w_sum.data_ptr<float>(), bias.data_ptr<float>(),
            bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            out.data_ptr<float>(), N, Cin, Cout, Vout, scale, eps);
    } else {
        auto conv_sum = torch::empty({N, Cout}, opts);
        conv_sum_kernel<<<N * Cout, T, T * sizeof(float), stream>>>(
            in_sum.data_ptr<float>(), w_sum.data_ptr<float>(), bias.data_ptr<float>(), conv_sum.data_ptr<float>(), N, Cin, Cout, Vout);
        auto sumsq = torch::zeros({Cout}, opts);
        int blocks_per_oc = 256;
        sumsq_kernel<<<Cout * blocks_per_oc, 128, 128 * sizeof(float), stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), sumsq.data_ptr<float>(),
            N, Cin, Cout, D, H, Wd, Kd, Kh, Kw, OD, OH, OW, scale, blocks_per_oc);
        int total = N * Cout;
        train_out_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
            conv_sum.data_ptr<float>(), sumsq.data_ptr<float>(), bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
            out.data_ptr<float>(), N, Cout, Vout, scale, eps);
    }
    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def forward(self, x):
        # The benchmark input and parameters are contiguous CUDA fp32 tensors.  Keep a safe
        # contiguous view for the custom address arithmetic.
        if not x.is_contiguous():
            x = x.contiguous()
        return _ext.fused_forward(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            self.batch_norm.weight,
            self.batch_norm.bias,
            self.batch_norm.running_mean,
            self.batch_norm.running_var,
            float(self.scale_factor),
            float(self.batch_norm.eps),
            bool(self.training),
        )