```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Key insight: the network ends with a global average pool over all spatial
# positions. With zero padding, every input voxel of a transposed convolution
# contributes to the output exactly once per kernel tap, so:
#
#   sum_spatial(convT(x))[n,oc] = sum_ic ( sum_spatial(x)[n,ic] * sum_k(w[ic,oc]) )
#                                  + V_out * bias[oc]
#
# i.e. the entire conv-transpose -> scale -> batchnorm -> global-avg-pool chain
# only needs per-channel sums of the input plus a tiny per-channel affine.
# This file implements that with custom CUDA kernels (eval mode), and a
# fused conv+reduction path for training mode.

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

#define BLOCK_THREADS 256

__inline__ __device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

__inline__ __device__ float blockReduceSum(float val) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    val = warpReduceSum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    int nwarps = (blockDim.x + 31) >> 5;
    val = (threadIdx.x < nwarps) ? shared[threadIdx.x] : 0.0f;
    if (wid == 0) val = warpReduceSum(val);
    __syncthreads();
    return val;
}

// One block per (n, c); reduces V contiguous spatial elements.
__global__ void spatial_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   long long V) {
    long long base = (long long)blockIdx.x * V;
    float sum = 0.0f;
    if ((V & 3LL) == 0LL) {
        const float4* x4 = reinterpret_cast<const float4*>(x + base);
        long long V4 = V >> 2;
        for (long long i = (long long)threadIdx.x; i < V4; i += blockDim.x) {
            float4 v = x4[i];
            sum += v.x + v.y + v.z + v.w;
        }
    } else {
        for (long long i = (long long)threadIdx.x; i < V; i += blockDim.x) {
            sum += x[base + i];
        }
    }
    sum = blockReduceSum(sum);
    if (threadIdx.x == 0) out[blockIdx.x] = sum;
}

// sw[ic,oc] = sum over kernel taps of w[ic,oc]
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ sw,
                                  long long K, long long total) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    long long base = idx * K;
    float s = 0.0f;
    for (long long k = 0; k < K; ++k) s += w[base + k];
    sw[idx] = s;
}

// out[n,oc] = gamma * (scale * (sum_ic sx[n,ic]*sw[ic,oc] / V + bias) - rmean)
//             * rsqrt(rvar + eps) + beta
__global__ void combine_eval_kernel(const float* __restrict__ sx,
                                    const float* __restrict__ sw,
                                    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,
                                    float invV, float scale, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * Cout) return;
    int n = idx / Cout;
    int oc = idx - n * Cout;
    const float* sxr = sx + (long long)n * Cin;
    const float* swc = sw + oc;
    float acc = 0.0f;
    for (int ic = 0; ic < Cin; ++ic) {
        acc += sxr[ic] * swc[(long long)ic * Cout];
    }
    float mean = acc * invV + bias[oc];
    out[idx] = gamma[oc] * (scale * mean - rmean[oc]) * rsqrtf(rvar[oc] + eps) + beta[oc];
}

// Per (n,c) sum and sum-of-squares of the conv output.
__global__ void reduce_train_kernel(const float* __restrict__ y,
                                    float* __restrict__ s1,
                                    float* __restrict__ s2,
                                    long long V) {
    long long base = (long long)blockIdx.x * V;
    float sum = 0.0f, sq = 0.0f;
    if ((V & 3LL) == 0LL) {
        const float4* y4 = reinterpret_cast<const float4*>(y + base);
        long long V4 = V >> 2;
        for (long long i = (long long)threadIdx.x; i < V4; i += blockDim.x) {
            float4 v = y4[i];
            sum += v.x + v.y + v.z + v.w;
            sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
        }
    } else {
        for (long long i = (long long)threadIdx.x; i < V; i += blockDim.x) {
            float v = y[base + i];
            sum += v;
            sq += v * v;
        }
    }
    sum = blockReduceSum(sum);
    sq = blockReduceSum(sq);
    if (threadIdx.x == 0) {
        s1[blockIdx.x] = sum;
        s2[blockIdx.x] = sq;
    }
}

// Batch statistics of the BN input (scale * conv) + running-stat update.
__global__ void train_stats_kernel(const float* __restrict__ s1,
                                   const float* __restrict__ s2,
                                   float* __restrict__ rmean,
                                   float* __restrict__ rvar,
                                   float* __restrict__ mean_y,
                                   float* __restrict__ invstd_y,
                                   int N, int C, float V, float scale,
                                   float eps, float momentum, int update) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    float sum = 0.0f, sq = 0.0f;
    for (int n = 0; n < N; ++n) {
        sum += s1[(long long)n * C + c];
        sq  += s2[(long long)n * C + c];
    }
    float M = (float)N * V;
    float mean_c = sum / M;
    float var_c = sq / M - mean_c * mean_c;
    var_c = fmaxf(var_c, 0.0f);
    float my = scale * mean_c;
    float vy = scale * scale * var_c;
    float istd = rsqrtf(vy + eps);
    if (update) {
        float m = momentum;
        rmean[c] = (1.0f - m) * rmean[c] + m * my;
        float var_unb = vy * (M / (M - 1.0f));
        rvar[c] = (1.0f - m) * rvar[c] + m * var_unb;
    }
    mean_y[c] = my;
    invstd_y[c] = istd;
}

__global__ void train_output_kernel(const float* __restrict__ s1,
                                    const float* __restrict__ mean_y,
                                    const float* __restrict__ invstd_y,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ beta,
                                    float* __restrict__ out,
                                    int N, int C, float invV, float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * C) return;
    int c = idx % C;
    float spatial_mean = scale * s1[idx] * invV;
    out[idx] = gamma[c] * (spatial_mean - mean_y[c]) * invstd_y[c] + beta[c];
}

torch::Tensor spatial_sum(torch::Tensor x) {
    auto xc = x.contiguous();
    long long numel = xc.numel();
    long long NC = (long long)xc.size(0) * xc.size(1);
    long long V = numel / NC;
    auto out = torch::empty({NC}, xc.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sum_kernel<<<(int)NC, BLOCK_THREADS, 0, stream>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(), V);
    return out;
}

torch::Tensor weight_sum(torch::Tensor w) {
    auto wc = w.contiguous();
    long long Cin = wc.size(0), Cout = wc.size(1);
    long long total = Cin * Cout;
    long long K = wc.numel() / total;
    auto out = torch::empty({total}, wc.options());
    int blocks = (int)((total + BLOCK_THREADS - 1) / BLOCK_THREADS);
    auto stream = at::cuda::getCurrentCUDAStream();
    weight_sum_kernel<<<blocks, BLOCK_THREADS, 0, stream>>>(
        wc.data_ptr<float>(), out.data_ptr<float>(), K, total);
    return out;
}

torch::Tensor combine_eval(torch::Tensor sx, torch::Tensor sw, torch::Tensor bias,
                           torch::Tensor gamma, torch::Tensor beta,
                           torch::Tensor rmean, torch::Tensor rvar,
                           int64_t N, int64_t Cin, int64_t Cout,
                           double invV, double scale, double eps) {
    auto out = torch::empty({N * Cout}, sx.options());
    int total = (int)(N * Cout);
    int blocks = (total + BLOCK_THREADS - 1) / BLOCK_THREADS;
    auto stream = at::cuda::getCurrentCUDAStream();
    combine_eval_kernel<<<blocks, BLOCK_THREADS, 0, stream>>>(
        sx.data_ptr<float>(), sw.data_ptr<float>(), bias.data_ptr<float>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        out.data_ptr<float>(), (int)N, (int)Cin, (int)Cout,
        (float)invV, (float)scale, (float)eps);
    return out;
}

std::vector<torch::Tensor> reduce_train(torch::Tensor y) {
    auto yc = y.contiguous();
    long long N = yc.size(0), C = yc.size(1);
    long long NC = N * C;
    long long V = yc.numel() / NC;
    auto s1 = torch::empty({NC}, yc.options());
    auto s2 = torch::empty({NC}, yc.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    reduce_train_kernel<<<(int)NC, BLOCK_THREADS, 0, stream>>>(
        yc.data_ptr<float>(), s1.data_ptr<float>(), s2.data_ptr<float>(), V);
    return {s1, s2};
}

torch::Tensor combine_train(torch::Tensor s1, torch::Tensor s2,
                            torch::Tensor rmean, torch::Tensor rvar,
                            torch::Tensor gamma, torch::Tensor beta,
                            int64_t N, int64_t C, double V, double scale,
                            double eps, double momentum, int64_t update) {
    auto opts = s1.options();
    auto mean_y = torch::empty({C}, opts);
    auto invstd = torch::empty({C}, opts);
    auto out = torch::empty({N * C}, opts);
    auto stream = at::cuda::getCurrentCUDAStream();
    int cblocks = (int)((C + BLOCK_THREADS - 1) / BLOCK_THREADS);
    train_stats_kernel<<<cblocks, BLOCK_THREADS, 0, stream>>>(
        s1.data_ptr<float>(), s2.data_ptr<float>(),
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        mean_y.data_ptr<float>(), invstd.data_ptr<float>(),
        (int)N, (int)C, (float)V, (float)scale, (float)eps,
        (float)momentum, (int)update);
    int total = (int)(N * C);
    int blocks = (total + BLOCK_THREADS - 1) / BLOCK_THREADS;
    train_output_kernel<<<blocks, BLOCK_THREADS, 0, stream>>>(
        s1.data_ptr<float>(), mean_y.data_ptr<float>(), invstd.data_ptr<float>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)C, (float)(1.0 / V), (float)scale);
    return out;
}
"""

cpp_source = r"""
#include <vector>
torch::Tensor spatial_sum(torch::Tensor x);
torch::Tensor weight_sum(torch::Tensor w);
torch::Tensor combine_eval(torch::Tensor sx, torch::Tensor sw, torch::Tensor bias,
                           torch::Tensor gamma, torch::Tensor beta,
                           torch::Tensor rmean, torch::Tensor rvar,
                           int64_t N, int64_t Cin, int64_t Cout,
                           double invV, double scale, double eps);
std::vector<torch::Tensor> reduce_train(torch::Tensor y);
torch::Tensor combine_train(torch::Tensor s1, torch::Tensor s2,
                            torch::Tensor rmean, torch::Tensor rvar,
                            torch::Tensor gamma, torch::Tensor beta,
                            int64_t N, int64_t C, double V, double scale,
                            double eps, double momentum, int64_t update);
"""

_ext = load_inline(
    name="fused_convt_bn_gap",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["spatial_sum", "weight_sum", "combine_eval", "reduce_train", "combine_train"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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))
        self._sw_cache = None

    def _output_sizes(self, x, w):
        ct = self.conv_transpose
        sizes = []
        for i in range(3):
            L = x.shape[2 + i]
            k = w.shape[2 + i]
            sizes.append((L - 1) * ct.stride[i] - 2 * ct.padding[i]
                         + ct.dilation[i] * (k - 1) + ct.output_padding[i] + 1)
        return sizes

    def _bn_affine(self, Cout, device):
        bn = self.batch_norm
        gamma = bn.weight if bn.weight is not None else torch.ones(Cout, device=device)
        beta = bn.bias if bn.bias is not None else torch.zeros(Cout, device=device)
        return gamma.contiguous(), beta.contiguous()

    def _forward_eval_fast(self, x):
        ct = self.conv_transpose
        bn = self.batch_norm
        w = ct.weight
        N, Cin = x.shape[0], x.shape[1]
        Cout = w.shape[1]
        s = self._output_sizes(x, w)
        V = float(s[0] * s[1] * s[2])

        key = (w.data_ptr(), w._version)
        if self._sw_cache is not None and self._sw_cache[0] == key:
            sw = self._sw_cache[1]
        else:
            sw = _ext.weight_sum(w)
            self._sw_cache = (key, sw)

        sx = _ext.spatial_sum(x)
        bias = ct.bias if ct.bias is not None else torch.zeros(Cout, device=x.device)
        gamma, beta = self._bn_affine(Cout, x.device)
        out = _ext.combine_eval(
            sx, sw, bias.contiguous(), gamma, beta,
            bn.running_mean, bn.running_var,
            N, Cin, Cout, 1.0 / V, float(self.scale_factor), float(bn.eps))
        return out.view(N, Cout, 1, 1, 1)

    def _forward_train_fast(self, x):
        ct = self.conv_transpose
        bn = self.batch_norm
        y = F.conv_transpose3d(x, ct.weight, ct.bias, ct.stride, ct.padding,
                               ct.output_padding, ct.groups, ct.dilation)
        N, Cout = y.shape[0], y.shape[1]
        V = float(y.numel() // (N * Cout))
        s1, s2 = _ext.reduce_train(y)
        gamma, beta = self._bn_affine(Cout, x.device)
        update = 1 if (self.training and bn.track_running_stats) else 0
        rmean = bn.running_mean if bn.running_mean is not None else gamma
        rvar = bn.running_var if bn.running_var is not None else beta
        if update:
            bn.num_batches_tracked += 1
        out = _ext.combine_train(
            s1, s2, rmean, rvar, gamma, beta,
            N, Cout, V, float(self.scale_factor), float(bn.eps),
            float(bn.momentum), update)
        return out.view(N, Cout, 1, 1, 1)

    def forward(self, x):
        ct = self.conv_transpose
        bn = self.batch_norm
        fast = (x.is_cuda and x.dtype == torch.float32 and ct.groups == 1
                and all(p == 0 for p in ct.padding) and bn.momentum is not None)
        if fast:
            if (not self.training) and bn.track_running_stats and bn.running_mean is not None:
                return self._forward_eval_fast(x)
            return self._forward_train_fast(x)
        y = ct(x)
        y = y * self.scale_factor
        y = bn(y)
        return self.global_avg_pool(y)
```