```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 CUDA kernels.
#
# Key insight: the network ends in a *global* average pool, and BatchNorm is a
# per-channel affine transform. Therefore the final result only depends on the
# per-channel spatial sums (and sums of squares, for train-mode batch stats)
# of the transposed-convolution output -- the full output tensor never needs
# to be materialized or re-read.
#
# In eval mode (fixed running stats) we can go one step further: the spatial
# sum of a stride-1/pad-0 ConvTranspose3d output decomposes exactly as
#     sum(out[b,o]) = sum_i sum(x[b,i]) * sum(w[i,o]) + bias[o] * S_out
# so the entire 500+ GFLOP convolution collapses into two cheap reductions
# plus a tiny 16x64x128 contraction, fused with scale + BN + pooling.
#
# In training mode we run the (unavoidable) convolution and fuse
# scale + batch-stat computation + running-stat update + BN + global pooling
# into a pair of lightweight reduction kernels.
# ---------------------------------------------------------------------------

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

#define BLOCK_SIZE 256

// Per-(batch, channel) spatial sums (and optionally sums of squares) of a
// contiguous NCDHW tensor. One block per (b, c) pair.
__global__ void spatial_sums_kernel(const float* __restrict__ x,
                                    float* __restrict__ sums,
                                    float* __restrict__ sumsq,
                                    int S, int compute_sq) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    long base = ((long)b * (long)gridDim.x + (long)c) * (long)S;
    const float* xp = x + base;

    float acc = 0.f, acc2 = 0.f;
    if ((S & 3) == 0) {
        const float4* xp4 = reinterpret_cast<const float4*>(xp);
        int S4 = S >> 2;
        for (int s = threadIdx.x; s < S4; s += BLOCK_SIZE) {
            float4 v = xp4[s];
            acc += (v.x + v.y) + (v.z + v.w);
            if (compute_sq) acc2 += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
        }
    } else {
        for (int s = threadIdx.x; s < S; s += BLOCK_SIZE) {
            float v = xp[s];
            acc += v;
            if (compute_sq) acc2 += v * v;
        }
    }

    // warp reduction
    for (int off = 16; off > 0; off >>= 1) {
        acc += __shfl_down_sync(0xffffffffu, acc, off);
        if (compute_sq) acc2 += __shfl_down_sync(0xffffffffu, acc2, off);
    }
    __shared__ float sh[BLOCK_SIZE / 32];
    __shared__ float sh2[BLOCK_SIZE / 32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh[wid] = acc;
        if (compute_sq) sh2[wid] = acc2;
    }
    __syncthreads();
    if (wid == 0) {
        int nw = BLOCK_SIZE / 32;
        float v  = (lane < nw) ? sh[lane] : 0.f;
        float v2 = (lane < nw && compute_sq) ? sh2[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
            if (compute_sq) v2 += __shfl_down_sync(0xffffffffu, v2, off);
        }
        if (lane == 0) {
            long idx = (long)b * (long)gridDim.x + (long)c;
            sums[idx] = v;
            if (compute_sq) sumsq[idx] = v2;
        }
    }
}

std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    int64_t B = xc.size(0);
    int64_t C = xc.size(1);
    int64_t S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    TORCH_CHECK(B <= 65535, "batch too large");
    auto sums  = torch::empty({B, C}, xc.options());
    auto sumsq = compute_sq ? torch::empty({B, C}, xc.options())
                            : torch::empty({0}, xc.options());
    if (B * C == 0) return {sums, sumsq};
    dim3 grid((unsigned)C, (unsigned)B);
    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sums_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), sums.data_ptr<float>(),
        compute_sq ? sumsq.data_ptr<float>() : nullptr,
        (int)S, compute_sq ? 1 : 0);
    return {sums, sumsq};
}

// Eval-mode fused kernel: contracts input sums with kernel sums, applies
// conv bias, scale, BatchNorm (running stats) and emits the globally pooled
// result. Sw is [In, Out] (ConvTranspose weight layout is (In, Out, kD,kH,kW)).
__global__ void eval_finalize_kernel(const float* __restrict__ Sx,
                                     const float* __restrict__ Sw,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ rmean,
                                     const float* __restrict__ rvar,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     float* __restrict__ out,
                                     int B, int In, int Out,
                                     float scale, float inv_Sout, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * Out) return;
    int b = idx / Out;
    int o = idx - b * Out;
    const float* sx = Sx + (long)b * In;
    float acc = 0.f;
    for (int i = 0; i < In; ++i) acc += sx[i] * Sw[(long)i * Out + o];
    float mean = acc * inv_Sout;
    if (bias != nullptr) mean += bias[o];
    mean *= scale;
    float g  = (gamma != nullptr) ? gamma[o] : 1.f;
    float be = (beta  != nullptr) ? beta[o]  : 0.f;
    out[idx] = (mean - rmean[o]) * rsqrtf(rvar[o] + eps) * g + be;
}

torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta,
                            double scale, double inv_Sout, double eps) {
    int B = Sx.size(0), In = Sx.size(1), Out = Sw.size(1);
    auto out = torch::empty({B, Out}, Sx.options());
    int total = B * Out;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    eval_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, In, Out,
        (float)scale, (float)inv_Sout, (float)eps);
    return out;
}

// Train-mode (or generic) fused kernel: given per-(b,c) sums and sums of
// squares of the convolution output, compute batch statistics, update the
// running statistics in place (exactly like PyTorch: biased variance for
// normalization, unbiased for the running average), then apply the affine
// transform and emit the globally pooled result. With training=0 the running
// statistics are used directly (eval fallback for non-unit conv params).
__global__ void train_finalize_kernel(const float* __restrict__ sums,
                                      const float* __restrict__ sumsq,
                                      float* __restrict__ rmean,
                                      float* __restrict__ rvar,
                                      const float* __restrict__ gamma,
                                      const float* __restrict__ beta,
                                      float* __restrict__ out,
                                      int B, int C, float S,
                                      float scale, float momentum,
                                      float eps, int training) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * C) return;
    int b = idx / C;
    int c = idx - b * C;
    float mean, invstd;
    float N = S * (float)B;
    if (training) {
        float s = 0.f, s2 = 0.f;
        for (int bb = 0; bb < B; ++bb) {
            s  += sums[(long)bb * C + c];
            s2 += sumsq[(long)bb * C + c];
        }
        s  *= scale;
        s2 *= scale * scale;
        mean = s / N;
        float var = s2 / N - mean * mean;
        var = fmaxf(var, 0.f);
        invstd = rsqrtf(var + eps);
        if (b == 0 && rmean != nullptr && momentum >= 0.f) {
            float unbiased = var * (N / (N - 1.f));
            rmean[c] = (1.f - momentum) * rmean[c] + momentum * mean;
            rvar[c]  = (1.f - momentum) * rvar[c]  + momentum * unbiased;
        }
    } else {
        mean = rmean[c];
        invstd = rsqrtf(rvar[c] + eps);
    }
    float g  = (gamma != nullptr) ? gamma[c] : 1.f;
    float be = (beta  != nullptr) ? beta[c]  : 0.f;
    float mbc = sums[(long)b * C + c] * (scale / S);
    out[idx] = (mbc - mean) * invstd * g + be;
}

torch::Tensor train_finalize(torch::Tensor sums, torch::Tensor sumsq,
                             torch::Tensor rmean, torch::Tensor rvar,
                             torch::Tensor gamma, torch::Tensor beta,
                             double scale, double S, double momentum,
                             double eps, bool training) {
    int B = sums.size(0), C = sums.size(1);
    auto out = torch::empty({B, C}, sums.options());
    int total = B * C;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    train_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        rmean.numel() ? rmean.data_ptr<float>() : nullptr,
        rvar.numel()  ? rvar.data_ptr<float>()  : nullptr,
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, C, (float)S, (float)scale,
        (float)momentum, (float)eps, training ? 1 : 0);
    return out;
}
"""

_cpp_src = r"""
#include <vector>
std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq);
torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta,
                            double scale, double inv_Sout, double eps);
torch::Tensor train_finalize(torch::Tensor sums, torch::Tensor sumsq,
                             torch::Tensor rmean, torch::Tensor rvar,
                             torch::Tensor gamma, torch::Tensor beta,
                             double scale, double S, double momentum,
                             double eps, bool training);
"""

_ext = load_inline(
    name="fused_convtranspose3d_bn_globalpool",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["spatial_sums", "eval_finalize", "train_finalize"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Keeps identical submodules (so state_dicts
    match) but replaces the operator chain with fused custom CUDA kernels.

    Eval mode: the ConvTranspose3d never runs -- its spatial sums are computed
    in closed form from per-channel input sums and kernel sums, then fused
    with scale + BatchNorm + global average pooling in a single tiny kernel.

    Train mode: the convolution runs (cuDNN) and scale + batch statistics +
    running-stat update + BatchNorm + global average pooling are fused into
    one reduction pass plus one tiny finalize kernel (no extra passes over the
    ~130 MB intermediate, which is also never re-read).
    """

    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 = float(scale_factor)
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        ct = self.conv_transpose
        # The closed-form sum decomposition requires the transposed conv to be
        # a "full" convolution (stride 1, no padding, unit dilation/groups).
        self._simple_conv = (
            ct.groups == 1
            and tuple(ct.stride) == (1, 1, 1)
            and tuple(ct.padding) == (0, 0, 0)
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
        )

    def _reference_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

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

        # Fall back to the reference pipeline for unsupported configurations.
        if not (x.is_cuda and x.dtype == torch.float32):
            return self._reference_forward(x)
        if bn.training and bn.momentum is None:
            return self._reference_forward(x)  # cumulative moving average
        if (not bn.training) and bn.track_running_stats and (
            bn.running_mean is None or not bn.running_mean.is_cuda
        ):
            return self._reference_forward(x)

        B = x.shape[0]
        gamma = bn.weight if bn.affine else x.new_empty(0)
        beta = bn.bias if bn.affine else x.new_empty(0)

        if (not bn.training) and bn.track_running_stats and self._simple_conv:
            # ---- closed-form path: no convolution at all -------------------
            In, Out = ct.in_channels, ct.out_channels
            Sx = _ext.spatial_sums(x, False)[0]            # [B, In]
            Sw = _ext.spatial_sums(ct.weight, False)[0]    # [In, Out]
            D, H, W = x.shape[2], x.shape[3], x.shape[4]
            k = ct.kernel_size
            S_out = (D + k[0] - 1) * (H + k[1] - 1) * (W + k[2] - 1)
            bias = ct.bias if ct.bias is not None else x.new_empty(0)
            out = _ext.eval_finalize(
                Sx, Sw, bias, bn.running_mean, bn.running_var, gamma, beta,
                self.scale_factor, 1.0 / float(S_out), bn.eps,
            )
            return out.view(B, Out, 1, 1, 1)

        # ---- general path: conv + fused reduction --------------------------
        y = ct(x)
        C = y.shape[1]
        S = y.shape[2] * y.shape[3] * y.shape[4]
        sums, sumsq = _ext.spatial_sums(y, True)
        if bn.track_running_stats:
            rmean, rvar = bn.running_mean, bn.running_var
        else:
            rmean = rvar = x.new_empty(0)
        momentum = bn.momentum if bn.momentum is not None else -1.0
        if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
            bn.num_batches_tracked += 1
        out = _ext.train_finalize(
            sums, sumsq, rmean, rvar, gamma, beta,
            self.scale_factor, float(S), momentum, bn.eps, bn.training,
        )
        return out.view(B, C, 1, 1, 1)
```