```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 algorithmic insight (eval mode): for a stride-1/pad-0 ConvTranspose3d,
# summing its output over the whole spatial domain makes every kernel tap
# multiply *every* input voxel exactly once:
#     sum_o out[co, o] = bias[co] * N_out + sum_ci Wsum[ci, co] * S[ci]
# where S[ci] is the spatial sum of the input and Wsum[ci, co] is the sum of
# the kernel taps. Since scale / BatchNorm / global-avg-pool are all affine
# per-channel ops, the whole network collapses to a tiny per-channel affine
# transform of the input's spatial sums - no need to materialize the huge
# (16, 128, 20, 36, 36) intermediate tensor at all.
#
# In training mode BatchNorm needs the true voxel-wise variance, so we keep
# the cuDNN transposed conv and fuse scale + batch-stat reduction + running
# stat update + normalization + pooling into two custom kernels.

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

#define TPB 256

__global__ void spatial_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out, int inner) {
    const float* base = x + (long long)blockIdx.x * inner;
    double acc = 0.0;
    if ((inner & 3) == 0) {
        int n4 = inner >> 2;
        const float4* b4 = reinterpret_cast<const float4*>(base);
        for (int i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = b4[i];
            acc += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        }
    } else {
        for (int i = threadIdx.x; i < inner; i += blockDim.x)
            acc += (double)base[i];
    }
    __shared__ double sh[32];
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    for (int o = 16; o > 0; o >>= 1) acc += __shfl_down_sync(0xffffffff, acc, o);
    if (lane == 0) sh[warp] = acc;
    __syncthreads();
    if (warp == 0) {
        int nw = (blockDim.x + 31) >> 5;
        double v = (lane < nw) ? sh[lane] : 0.0;
        for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
        if (lane == 0) out[blockIdx.x] = (float)v;
    }
}

__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ out, int kvol, int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    const float* base = w + (long long)idx * kvol;
    double acc = 0.0;
    for (int i = 0; i < kvol; ++i) acc += (double)base[i];
    out[idx] = (float)acc;
}

__global__ void eval_fused_kernel(const float* __restrict__ S,
                                  const float* __restrict__ Wsum,
                                  const float* __restrict__ cbias,
                                  double scale,
                                  const float* __restrict__ gamma,
                                  const float* __restrict__ beta,
                                  const float* __restrict__ rmean,
                                  const float* __restrict__ rvar,
                                  double eps, double inv_nout,
                                  float* __restrict__ out,
                                  int B, int CI, int CO) {
    int co = blockIdx.x;
    double g = (double)gamma[co], be = (double)beta[co], rm = (double)rmean[co];
    double invstd = 1.0 / sqrt((double)rvar[co] + eps);
    double cb = (cbias != nullptr) ? (double)cbias[co] : 0.0;
    for (int b = threadIdx.x; b < B; b += blockDim.x) {
        double acc = 0.0;
        for (int ci = 0; ci < CI; ++ci)
            acc += (double)S[b * CI + ci] * (double)Wsum[ci * CO + co];
        double y = scale * (acc * inv_nout + cb);
        out[b * CO + co] = (float)(g * (y - rm) * invstd + be);
    }
}

__global__ void train_reduce_kernel(const float* __restrict__ y,
                                    float* __restrict__ sum_out,
                                    double* __restrict__ sumsq_out, int inner) {
    const float* base = y + (long long)blockIdx.x * inner;
    double s1 = 0.0, s2 = 0.0;
    if ((inner & 3) == 0) {
        int n4 = inner >> 2;
        const float4* b4 = reinterpret_cast<const float4*>(base);
        for (int i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = b4[i];
            s1 += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
            s2 += (double)v.x * v.x + (double)v.y * v.y
                + (double)v.z * v.z + (double)v.w * v.w;
        }
    } else {
        for (int i = threadIdx.x; i < inner; i += blockDim.x) {
            double v = (double)base[i];
            s1 += v; s2 += v * v;
        }
    }
    __shared__ double sh1[32], sh2[32];
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    for (int o = 16; o > 0; o >>= 1) {
        s1 += __shfl_down_sync(0xffffffff, s1, o);
        s2 += __shfl_down_sync(0xffffffff, s2, o);
    }
    if (lane == 0) { sh1[warp] = s1; sh2[warp] = s2; }
    __syncthreads();
    if (warp == 0) {
        int nw = (blockDim.x + 31) >> 5;
        double a = (lane < nw) ? sh1[lane] : 0.0;
        double b = (lane < nw) ? sh2[lane] : 0.0;
        for (int o = 16; o > 0; o >>= 1) {
            a += __shfl_down_sync(0xffffffff, a, o);
            b += __shfl_down_sync(0xffffffff, b, o);
        }
        if (lane == 0) {
            sum_out[blockIdx.x] = (float)a;
            sumsq_out[blockIdx.x] = b;
        }
    }
}

__global__ void train_finalize_kernel(const float* __restrict__ sum_bc,
                                      const double* __restrict__ sumsq_bc,
                                      double scale,
                                      float* __restrict__ gamma,
                                      float* __restrict__ beta,
                                      float* __restrict__ rmean,
                                      float* __restrict__ rvar,
                                      double eps, double momentum,
                                      double inv_inner, double n_total,
                                      float* __restrict__ out, int B, int C) {
    int c = blockIdx.x;
    double total = 0.0, total2 = 0.0;
    for (int i = 0; i < B; ++i) {
        total += (double)sum_bc[i * C + c];
        total2 += sumsq_bc[i * C + c];
    }
    double mu = scale * total / n_total;
    double var = scale * scale * total2 / n_total - mu * mu;
    if (var < 0.0) var = 0.0;
    double invstd = 1.0 / sqrt(var + eps);
    if (threadIdx.x == 0) {
        rmean[c] = (float)((1.0 - momentum) * (double)rmean[c] + momentum * mu);
        double var_unb = var * n_total / (n_total - 1.0);
        rvar[c] = (float)((1.0 - momentum) * (double)rvar[c] + momentum * var_unb);
    }
    double g = (double)gamma[c], be = (double)beta[c];
    for (int b = threadIdx.x; b < B; b += blockDim.x) {
        double pooled = (double)sum_bc[b * C + c] * inv_inner;
        out[b * C + c] = (float)(g * (scale * pooled - mu) * invstd + be);
    }
}

torch::Tensor fast_eval(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                        double scale, torch::Tensor gamma, torch::Tensor beta,
                        torch::Tensor rmean, torch::Tensor rvar, double eps) {
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    int B = xc.size(0), CI = xc.size(1);
    long long D = xc.size(2), H = xc.size(3), W = xc.size(4);
    long long inner = D * H * W;
    int CO = wc.size(1);
    long long kvol = (long long)wc.size(2) * wc.size(3) * wc.size(4);
    long long nout = (D + wc.size(2) - 1) * (H + wc.size(3) - 1) * (W + wc.size(4) - 1);
    auto stream = at::cuda::getCurrentCUDAStream();

    auto S = torch::empty({B, CI}, xc.options());
    spatial_sum_kernel<<<B * CI, TPB, 0, stream>>>(
        xc.data_ptr<float>(), S.data_ptr<float>(), (int)inner);

    auto Wsum = torch::empty({CI, CO}, xc.options());
    int total = CI * CO;
    weight_sum_kernel<<<(total + TPB - 1) / TPB, TPB, 0, stream>>>(
        wc.data_ptr<float>(), Wsum.data_ptr<float>(), (int)kvol, total);

    auto out = torch::empty({B, CO, 1, 1, 1}, xc.options());
    const float* bias_ptr =
        (bias.defined() && bias.numel() == CO) ? bias.data_ptr<float>() : nullptr;
    int threads = B < TPB ? B : TPB;
    eval_fused_kernel<<<CO, threads, 0, stream>>>(
        S.data_ptr<float>(), Wsum.data_ptr<float>(), bias_ptr, scale,
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        eps, 1.0 / (double)nout, out.data_ptr<float>(), B, CI, CO);
    return out;
}

torch::Tensor fused_bn_pool_train(torch::Tensor y, double scale,
                                  torch::Tensor gamma, torch::Tensor beta,
                                  torch::Tensor rmean, torch::Tensor rvar,
                                  double eps, double momentum) {
    auto yc = y.contiguous();
    int B = yc.size(0), C = yc.size(1);
    long long inner = (long long)yc.size(2) * yc.size(3) * yc.size(4);
    auto stream = at::cuda::getCurrentCUDAStream();

    auto sum_bc = torch::empty({B, C}, yc.options());
    auto sumsq_bc = torch::empty({B, C}, yc.options().dtype(at::kDouble));
    train_reduce_kernel<<<B * C, TPB, 0, stream>>>(
        yc.data_ptr<float>(), sum_bc.data_ptr<float>(),
        sumsq_bc.data_ptr<double>(), (int)inner);

    auto out = torch::empty({B, C, 1, 1, 1}, yc.options());
    int threads = B < TPB ? B : TPB;
    double n_total = (double)B * (double)inner;
    train_finalize_kernel<<<C, threads, 0, stream>>>(
        sum_bc.data_ptr<float>(), sumsq_bc.data_ptr<double>(), scale,
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        eps, momentum, 1.0 / (double)inner, n_total,
        out.data_ptr<float>(), B, C);
    return out;
}
"""

cpp_source = """
torch::Tensor fast_eval(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                        double scale, torch::Tensor gamma, torch::Tensor beta,
                        torch::Tensor rmean, torch::Tensor rvar, double eps);
torch::Tensor fused_bn_pool_train(torch::Tensor y, double scale,
                                  torch::Tensor gamma, torch::Tensor beta,
                                  torch::Tensor rmean, torch::Tensor rvar,
                                  double eps, double momentum);
"""

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


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

    Eval mode: uses an exact algebraic collapse of
    ConvTranspose3d(stride=1,pad=0) -> scale -> BatchNorm -> global avg pool
    into per-channel affine transforms of the input's spatial sums, so the
    giant intermediate activation is never materialized.

    Training mode: keeps the cuDNN transposed convolution (needed for true
    voxel-wise batch statistics) and fuses scale + batch statistics +
    running-stat update + normalization + global average pooling into custom
    CUDA kernels.
    """

    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 _default_conv_config(self):
        ct = self.conv_transpose
        return (ct.stride == (1, 1, 1) and ct.padding == (0, 0, 0)
                and ct.dilation == (1, 1, 1) and ct.output_padding == (0, 0, 0)
                and ct.groups == 1)

    def forward(self, x):
        bn = self.batch_norm
        ct = self.conv_transpose
        fast_ok = (x.is_cuda and x.dtype == torch.float32
                   and bn.affine and bn.track_running_stats)

        if fast_ok and bn.training and bn.momentum is not None:
            y = ct(x)
            out = _ext.fused_bn_pool_train(
                y, float(self.scale_factor), bn.weight, bn.bias,
                bn.running_mean, bn.running_var, float(bn.eps), float(bn.momentum))
            if bn.num_batches_tracked is not None:
                bn.num_batches_tracked.add_(1)
            return out

        if fast_ok and (not bn.training) and self._default_conv_config():
            bias = ct.bias if ct.bias is not None else torch.empty(
                0, device=x.device, dtype=x.dtype)
            return _ext.fast_eval(
                x, ct.weight, bias, float(self.scale_factor), bn.weight, bn.bias,
                bn.running_mean, bn.running_var, float(bn.eps))

        # Safe fallback for any non-default configuration.
        x = ct(x)
        x = x * self.scale_factor
        x = bn(x)
        return self.global_avg_pool(x)
```