```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
from torch.utils.cpp_extension import load_inline

# Custom fused CUDA kernels.
#
# Key insight: global average pooling commutes with the per-channel affine
# transforms (scale + BatchNorm). For a stride-1, zero-padded ConvTranspose3d,
# the spatial sum of the output factorizes exactly:
#     sum_p out[n,oc,p] = bias[oc]*SPAT + sum_ic (sum_k w[ic,oc,k]) * (sum_p x[n,ic,p])
# so in eval mode (BN uses running stats) the whole model collapses to two
# cheap reductions plus a tiny matmul, never materializing the huge conv output.
# In training mode (batch stats) we additionally need the per-channel sum of
# squares of the conv output, so we run the conv and fuse
# scale+BN+global-pool into a single reduction + finalize kernel.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TPB 256

// One block per (n, c) slice: sum over the spatial dimensions.
__global__ void spatial_sum_kernel(const float4* __restrict__ x4,
                                   float* __restrict__ out,
                                   int s4) {
    const long base = (long)blockIdx.x * (long)s4;
    float s = 0.0f;
    for (int i = threadIdx.x; i < s4; i += TPB) {
        float4 v = x4[base + i];
        s += v.x + v.y + v.z + v.w;
    }
    __shared__ float sh[TPB];
    sh[threadIdx.x] = s;
    __syncthreads();
    for (int o = TPB >> 1; o > 0; o >>= 1) {
        if (threadIdx.x < o) sh[threadIdx.x] += sh[threadIdx.x + o];
        __syncthreads();
    }
    if (threadIdx.x == 0) out[blockIdx.x] = sh[0];
}

// One thread per (ic, oc) pair: sum over the kernel dimensions.
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ ws,
                                  int K, int total) {
    int idx = blockIdx.x * TPB + threadIdx.x;
    if (idx >= total) return;
    const float* p = w + (long)idx * (long)K;
    float s = 0.0f;
    for (int k = 0; k < K; ++k) s += p[k];
    ws[idx] = s;
}

// S[n,oc] = sum_ic xsum[n,ic] * wsum[ic,oc], then apply the fused
// scale + batchnorm (running stats) + global-average-pool affine transform.
__global__ void matmul_affine_kernel(const float* __restrict__ xsum,
                                     const float* __restrict__ wsum,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     const float* __restrict__ rm,
                                     const float* __restrict__ rv,
                                     const float* __restrict__ cbias,
                                     float* __restrict__ out,
                                     int IC, int OC, int has_bias,
                                     float scale, float inv_spat, float eps) {
    int n = blockIdx.y;
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= OC) return;
    const float* xr = xsum + (long)n * (long)IC;
    float acc = 0.0f;
    for (int i = 0; i < IC; ++i) acc += xr[i] * wsum[(long)i * (long)OC + c];
    float mean = acc * inv_spat;
    if (has_bias) mean += cbias[c];
    float inv = rsqrtf(rv[c] + eps);
    out[(long)n * (long)OC + c] = (scale * mean - rm[c]) * inv * gamma[c] + beta[c];
}

// One block per (n, c) slice of the conv output: sum and sum of squares.
__global__ void sum_sumsq_kernel(const float4* __restrict__ y4,
                                 float* __restrict__ S,
                                 float* __restrict__ T,
                                 int s4, int OC) {
    const long base = (long)blockIdx.x * (long)s4;
    float s = 0.0f, q = 0.0f;
    for (int i = threadIdx.x; i < s4; i += TPB) {
        float4 v = y4[base + i];
        s += v.x + v.y + v.z + v.w;
        q += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    __shared__ float shs[TPB];
    __shared__ float shq[TPB];
    shs[threadIdx.x] = s;
    shq[threadIdx.x] = q;
    __syncthreads();
    for (int o = TPB >> 1; o > 0; o >>= 1) {
        if (threadIdx.x < o) {
            shs[threadIdx.x] += shs[threadIdx.x + o];
            shq[threadIdx.x] += shq[threadIdx.x + o];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        S[blockIdx.x] = shs[0];
        atomicAdd(&T[blockIdx.x % OC], shq[0]);
    }
}

// Single block: compute batch statistics, update running stats exactly like
// BatchNorm3d, and emit the fused scale+BN+global-pool output.
__global__ void bn_finalize_kernel(const float* __restrict__ S,
                                   const float* __restrict__ T,
                                   float* __restrict__ rm,
                                   float* __restrict__ rv,
                                   const float* __restrict__ gamma,
                                   const float* __restrict__ beta,
                                   float* __restrict__ out,
                                   int N, int OC,
                                   double scale, double M,
                                   float inv_spat, double eps, double momentum) {
    for (int c = threadIdx.x; c < OC; c += blockDim.x) {
        double U = 0.0;
        for (int n = 0; n < N; ++n) U += (double)S[(long)n * (long)OC + c];
        double Uz  = scale * U;                       // sum of z = scale * conv
        double Tz  = scale * scale * (double)T[c];    // sum of z^2
        double mean = Uz / M;
        double ssq = Tz - Uz * Uz / M;                // sum (z - mean)^2
        if (ssq < 0.0) ssq = 0.0;
        double var_b = ssq / M;                       // biased (normalization)
        double var_u = ssq / (M - 1.0);               // unbiased (running avg)
        rm[c] = (float)((1.0 - momentum) * (double)rm[c] + momentum * mean);
        rv[c] = (float)((1.0 - momentum) * (double)rv[c] + momentum * var_u);
        float inv = rsqrtf((float)(var_b + eps));
        float g = gamma[c], be = beta[c];
        float mf = (float)mean;
        for (int n = 0; n < N; ++n) {
            float m = S[(long)n * (long)OC + c] * inv_spat;
            out[(long)n * (long)OC + c] = (scale * m - mf) * inv * g + be;
        }
    }
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor weight,
                         torch::Tensor gamma, torch::Tensor beta,
                         torch::Tensor rm, torch::Tensor rv,
                         torch::Tensor cbias, int64_t has_bias,
                         double scale, double eps, int64_t out_spat) {
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    const int  N  = xc.size(0);
    const int  IC = xc.size(1);
    const long in_spat = (long)xc.size(2) * (long)xc.size(3) * (long)xc.size(4);
    const int  K  = wc.size(2) * wc.size(3) * wc.size(4);
    const int  OC = wc.size(1);  // ConvTranspose3d weight: (IC, OC, kD, kH, kW)

    auto opts = xc.options();
    auto xsum = torch::empty({N, IC}, opts);
    auto wsum = torch::empty({IC, OC}, opts);
    auto out  = torch::empty({N, OC, 1, 1, 1}, opts);

    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sum_kernel<<<N * IC, TPB, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        xsum.data_ptr<float>(), (int)(in_spat / 4));
    const int total = IC * OC;
    weight_sum_kernel<<<(total + TPB - 1) / TPB, TPB, 0, stream>>>(
        wc.data_ptr<float>(), wsum.data_ptr<float>(), K, total);
    dim3 grid((OC + 127) / 128, N);
    matmul_affine_kernel<<<grid, 128, 0, stream>>>(
        xsum.data_ptr<float>(), wsum.data_ptr<float>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        rm.data_ptr<float>(), rv.data_ptr<float>(),
        has_bias ? cbias.data_ptr<float>() : nullptr,
        out.data_ptr<float>(), IC, OC, (int)has_bias,
        (float)scale, (float)(1.0 / (double)out_spat), (float)eps);
    return out;
}

torch::Tensor fused_train(torch::Tensor y,
                          torch::Tensor gamma, torch::Tensor beta,
                          torch::Tensor rm, torch::Tensor rv,
                          double scale, double eps, double momentum) {
    auto yc = y.contiguous();
    const int  N  = yc.size(0);
    const int  OC = yc.size(1);
    const long spat = (long)yc.size(2) * (long)yc.size(3) * (long)yc.size(4);

    auto opts = yc.options();
    auto S   = torch::empty({N, OC}, opts);
    auto T   = torch::zeros({OC}, opts);
    auto out = torch::empty({N, OC, 1, 1, 1}, opts);

    auto stream = at::cuda::getCurrentCUDAStream();
    sum_sumsq_kernel<<<N * OC, TPB, 0, stream>>>(
        reinterpret_cast<const float4*>(yc.data_ptr<float>()),
        S.data_ptr<float>(), T.data_ptr<float>(), (int)(spat / 4), OC);
    bn_finalize_kernel<<<1, TPB, 0, stream>>>(
        S.data_ptr<float>(), T.data_ptr<float>(),
        rm.data_ptr<float>(), rv.data_ptr<float>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        out.data_ptr<float>(), N, OC,
        scale, (double)N * (double)spat,
        (float)(1.0 / (double)spat), eps, momentum);
    return out;
}
"""

cpp_source = """
torch::Tensor fused_eval(torch::Tensor x, torch::Tensor weight,
                         torch::Tensor gamma, torch::Tensor beta,
                         torch::Tensor rm, torch::Tensor rv,
                         torch::Tensor cbias, int64_t has_bias,
                         double scale, double eps, int64_t out_spat);
torch::Tensor fused_train(torch::Tensor y,
                          torch::Tensor gamma, torch::Tensor beta,
                          torch::Tensor rm, torch::Tensor rv,
                          double scale, double eps, double momentum);
"""

_ext = load_inline(
    name="fused_convt_scale_bn_gap",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_eval", "fused_train"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Eval mode: collapses ConvTranspose3d + scale + BatchNorm3d + global average
    pooling into two tiny reductions plus a small fused matmul (the conv output
    is never materialized) by exploiting the exact factorization of the spatial
    sum of a stride-1, zero-padded transposed convolution.

    Train mode: keeps the cuDNN transposed convolution and fuses the scaling,
    batch-statistics computation, running-stat update, batch normalization and
    global average pooling into a single reduction + finalize kernel pair.
    """
    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 _fallback(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

    def forward(self, x):
        bn = self.batch_norm
        conv = self.conv_transpose

        if not (x.is_cuda and x.dtype == torch.float32):
            return self._fallback(x)

        affine = bn.weight is not None and bn.bias is not None

        if bn.training:
            y = conv(x)
            spat = y.shape[2] * y.shape[3] * y.shape[4]
            if (affine and bn.track_running_stats and bn.momentum is not None
                    and spat % 4 == 0):
                bn.num_batches_tracked.add_(1)
                return _ext.fused_train(
                    y, bn.weight, bn.bias, bn.running_mean, bn.running_var,
                    float(self.scale_factor), float(bn.eps), float(bn.momentum))
            y = y * self.scale_factor
            y = bn(y)
            return self.global_avg_pool(y)

        # Eval mode: try the fully reduced fast path.
        in_spat = x.shape[2] * x.shape[3] * x.shape[4]
        fast = (affine and bn.track_running_stats and conv.groups == 1
                and in_spat % 4 == 0
                and all(int(s) == 1 for s in conv.stride)
                and all(int(p) == 0 for p in conv.padding)
                and all(int(op) == 0 for op in conv.output_padding))
        if not fast:
            return self._fallback(x)

        D, H, W = x.shape[2], x.shape[3], x.shape[4]
        od = ((D - 1) * conv.stride[0] - 2 * conv.padding[0]
              + conv.dilation[0] * (conv.kernel_size[0] - 1)
              + conv.output_padding[0] + 1)
        oh = ((H - 1) * conv.stride[1] - 2 * conv.padding[1]
              + conv.dilation[1] * (conv.kernel_size[1] - 1)
              + conv.output_padding[1] + 1)
        ow = ((W - 1) * conv.stride[2] - 2 * conv.padding[2]
              + conv.dilation[2] * (conv.kernel_size[2] - 1)
              + conv.output_padding[2] + 1)
        out_spat = od * oh * ow

        if conv.bias is not None:
            cbias = conv.bias
            has_bias = 1
        else:
            cbias = torch.empty(0, device=x.device, dtype=torch.float32)
            has_bias = 0

        return _ext.fused_eval(
            x, conv.weight, bn.weight, bn.bias, bn.running_mean, bn.running_var,
            cbias, has_bias, float(self.scale_factor), float(bn.eps), int(out_spat))
```