```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

# ---------------------------------------------------------------------------
# Algorithmic rewrite of: ConvTranspose3d -> scale -> BatchNorm3d -> global avg pool.
#
# The network ends in a *global* average pool and BatchNorm is a per-channel
# affine transform, so the result only depends on the per-(batch, channel)
# spatial sums and sums-of-squares of the transposed-convolution output.
#
# For stride-1 / pad-0 ConvTranspose3d, the output is the *full* linear
# convolution of x and w, so by the convolution theorem and Parseval:
#     sum_z y[c,z]  = Y[c,0]                     (DC Fourier component)
#     sum_z y[c,z]^2 = (1/N) * sum_f |Y[c,f]|^2  (Parseval)
# with Y[c,f] = sum_i X[i,f] * W[i,c,f] (a small complex GEMM per frequency).
# This avoids materializing the ~200MB / ~850 GFLOP convolution entirely:
#   cost ~ one rfftn of the input + one memory-bound complex GEMM.
# The pointwise |Y|^2 reduction is fused into the GEMM kernel, and a final
# tiny kernel applies conv bias, scale, BatchNorm (batch or running stats)
# and emits the globally pooled result.
# ---------------------------------------------------------------------------

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

#define BLOCK_SIZE 256

// ------------------------------ helpers -----------------------------------
__device__ __forceinline__ float warp_reduce_sum(float v) {
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;
}

// Per-(batch, channel) spatial sums of a contiguous NCDHW tensor.
// One block per (b, c).
__global__ void spatial_sums_kernel(const float* __restrict__ x,
                                    float* __restrict__ sums,
                                    int S) {
    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;
    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);
        }
    } else {
        for (int s = threadIdx.x; s < S; s += BLOCK_SIZE) acc += xp[s];
    }
    __shared__ float sh[BLOCK_SIZE / 32];
    acc = warp_reduce_sum(acc);
    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = acc;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < BLOCK_SIZE / 32) ? sh[lane] : 0.f;
        v = warp_reduce_sum(v);
        if (lane == 0) sums[(long)b * (long)gridDim.x + (long)c] = v;
    }
}

torch::Tensor spatial_sums(torch::Tensor x) {
    auto xc = x.contiguous();
    int64_t B = xc.size(0), C = xc.size(1), S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    auto sums = torch::empty({B, C}, xc.options());
    if (B * C == 0) return sums;
    TORCH_CHECK(B <= 65535, "batch too large");
    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>(), (int)S);
    return sums;
}

// -------------------- fused frequency GEMM + Parseval ---------------------
// X: [F, B, I] complex (float2), W: [F, I, O] complex.
// For every frequency f the block computes ACC[b,o] = sum_i X[f,b,i]*W[f,i,o]
// (complex), then folds the Parseval weight:
//   block f==0            -> sums[b,o]  = Re(ACC[b,o])
//   block with 0<f<dc_blk -> sumsq[b,o] += 2*|ACC|^2 (interior frequencies)
//   block f==dc_blk       -> sumsq[b,o] += w_nyq*|ACC|^2 (DC/Nyquist planes)
// Grid: (dc_blk+1, num_tiles), block: 256 threads, 8 (b,o) pairs each.
// Requires B*O <= 2048, I <= 1024, O % 32 == 0.
__global__ void freq_gemm_parseval_kernel(const float2* __restrict__ X,
                                          const float2* __restrict__ W,
                                          float* __restrict__ sums,
                                          float* __restrict__ sumsq,
                                          int F, int B, int I, int O,
                                          int dc_blk, float nyq_w,
                                          float inv_n) {
    int f = blockIdx.x;
    int tile = blockIdx.y;
    int pbase = tile * 2048;

    __shared__ float2 sX[16 * 64];        // [B][I] tile
    __shared__ float2 sW[64 * 32];        // [I][32] chunk
    __shared__ float sDC[2048];           // DC accumulators (f == 0 only)

    const int BI = B * I;
    for (int t = threadIdx.x; t < BI; t += BLOCK_SIZE)
        sX[t] = X[(long)f * BI + t];
    if (f == 0) {
        for (int t = threadIdx.x; t < 2048; t += BLOCK_SIZE) sDC[t] = 0.f;
    }
    __syncthreads();

    float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f;
    float acc4 = 0.f, acc5 = 0.f, acc6 = 0.f, acc7 = 0.f;
    float bw[8];
    #pragma unroll
    for (int k = 0; k < 8; ++k) {
        int p = pbase + k * BLOCK_SIZE + threadIdx.x;
        int b = (p < B * O) ? (p / O) : 0;
        bw[k] = (float)b * (float)I;
    }

    for (int oc = 0; oc < O; oc += 32) {
        for (int t = threadIdx.x; t < (I << 5); t += BLOCK_SIZE)
            sW[t] = W[(long)f * I * O + (long)(t >> 5) * O + oc + (t & 31)];
        __syncthreads();
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            int p = pbase + k * BLOCK_SIZE + threadIdx.x;
            if (p >= B * O) break;
            int o = p - (int)bw[k] / (float)I * 0;  // placeholder (unused)
            o = p % O;
            float2 a = make_float2(0.f, 0.f);
            const float2* xrow = sX + (int)bw[k];
            const float2* wcol = sW + (o - oc);
            for (int i = 0; i < I; ++i) {
                float2 xv = __ldg(xrow + i);
                float2 wv = wcol[i << 5];
                a.x += xv.x * wv.x - xv.y * wv.y;
                a.y += xv.x * wv.y + xv.y * wv.x;
            }
            float pw = a.x * a.x + a.y * a.y;
            if (f == 0) {
                atomicAdd(&sDC[k * BLOCK_SIZE + threadIdx.x], a.x);
                if (oc == 0) acc0 += 0.f;  // keep structure simple
            }
            switch (k) {
                case 0: acc0 += pw; break; case 1: acc1 += pw; break;
                case 2: acc2 += pw; break; case 3: acc3 += pw; break;
                case 4: acc4 += pw; break; case 5: acc5 += pw; break;
                case 6: acc6 += pw; break; default: acc7 += pw; break;
            }
        }
        __syncthreads();
    }

    if (f == 0) {
        __syncthreads();
        for (int t = threadIdx.x; t < 2048; t += BLOCK_SIZE) {
            if (pbase + t < B * O) sums[pbase + t] = sDC[t];
        }
    }

    float wgt;
    if (f == 0 || f == dc_blk) wgt = nyq_w * inv_n;
    else                       wgt = 2.0f * inv_n;
    #pragma unroll
    for (int k = 0; k < 8; ++k) {
        int p = pbase + k * BLOCK_SIZE + threadIdx.x;
        if (p >= B * O) break;
        float v;
        switch (k) {
            case 0: v = acc0; break; case 1: v = acc1; break;
            case 2: v = acc2; break; case 3: v = acc3; break;
            case 4: v = acc4; break; case 5: v = acc5; break;
            case 6: v = acc6; break; default: v = acc7; break;
        }
        atomicAdd(&sumsq[p], v * wgt);
    }
}

std::vector<torch::Tensor> freq_gemm_parseval(torch::Tensor X, torch::Tensor W,
                                              int64_t B, int64_t I, int64_t O,
                                              int64_t dc_blk, double nyq_w,
                                              double inv_n) {
    int F = X.size(0);
    auto opts = X.options().dtype(at::kFloat);
    auto sums  = torch::empty({B * O}, opts);
    auto sumsq = torch::zeros({B * O}, opts);
    int BO = B * (int)O;
    int tiles = (BO + 2047) / 2048;
    dim3 grid((unsigned)(dc_blk + 1), (unsigned)tiles);
    auto stream = at::cuda::getCurrentCUDAStream();
    freq_gemm_parseval_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        reinterpret_cast<const float2*>(X.data_ptr<float>()),
        reinterpret_cast<const float2*>(W.data_ptr<float>()),
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        F, (int)B, (int)I, (int)O, (int)dc_blk, (float)nyq_w, (float)inv_n);
    return {sums, sumsq};
}

// ------------------------------ finalize ----------------------------------
// Applies conv-bias correction to the moments, then scale + BatchNorm and
// emits the globally pooled result. Training mode uses batch statistics and
// updates running stats exactly like PyTorch (biased var for normalize,
// unbiased for running). sum0: raw spatial sums (DC), sumsq0: raw sums of
// squares (Parseval). y = conv(x) + bias; value = scale * y.
__global__ void finalize_kernel(const float* __restrict__ sum0,
                                const float* __restrict__ sumsq0,
                                const float* __restrict__ bias,
                                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 bc = (bias != nullptr) ? bias[c] : 0.f;
    float s_raw  = sum0[idx];
    float s2_raw = sumsq0[idx];
    float su  = s_raw + bc * S;
    float s2u = s2_raw + 2.f * bc * s_raw + bc * bc * S;
    float mean, invstd;
    float N = S * (float)B;
    if (training) {
        float s = 0.f, s2 = 0.f;
        for (int bb = 0; bb < B; ++bb) {
            float sr = sum0[(long)bb * C + c];
            float s2r = sumsq0[(long)bb * C + c];
            s  += sr + bc * S;
            s2 += s2r + 2.f * bc * sr + bc * bc * S;
        }
        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 = su * (scale / S);
    out[idx] = (mbc - mean) * invstd * g + be;
}

torch::Tensor finalize(torch::Tensor sums, torch::Tensor sumsq,
                       torch::Tensor bias, 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();
    finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        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;
}

// Eval closed form: per-channel sums of the conv output via
// sum(out[b,o]) = sum_i Sx[b,i] * Sw[i,o] + bias[o]*S_out, then scale+BN+pool.
__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;
}
"""

_cpp_src = r"""
#include <vector>
torch::Tensor spatial_sums(torch::Tensor x);
std::vector<torch::Tensor> freq_gemm_parseval(torch::Tensor X, torch::Tensor W,
                                              int64_t B, int64_t I, int64_t O,
                                              int64_t dc_blk, double nyq_w,
                                              double inv_n);
torch::Tensor finalize(torch::Tensor sums, torch::Tensor sumsq,
                       torch::Tensor bias, torch::Tensor rmean,
                       torch::Tensor rvar, torch::Tensor gamma,
                       torch::Tensor beta, double scale, double S,
                       double momentum, double eps, bool training);
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);
"""

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


class ModelNew(nn.Module):
    """
    Drop-in optimized replacement for Model (identical submodules/state dict).

    Train mode: the ConvTranspose3d is evaluated in the Fourier domain
    (full convolution == pointwise product of spectra).  Only the DC
    component (spatial sums) and the Parseval power (sums of squares) of the
    output spectrum are computed -- fused into one custom CUDA kernel -- and a
    tiny finalize kernel applies bias, scale, BatchNorm (with running-stat
    update) and global average pooling.  The huge conv output tensor and the
    ~850 GFLOP direct convolution are never materialized.

    Eval mode: spatial sums of the conv output are computed in closed form
    from per-channel input sums and kernel sums (also exact), so no
    convolution runs at all.
    """

    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
        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)
        )
        self._w_spec = None   # cached weight spectrum [F, In, Out] (complex64)
        self._w_key = None

    # ------------------------------------------------------------------ util
    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 _out_spatial(self, x):
        ct = self.conv_transpose
        D, H, W = x.shape[2], x.shape[3], x.shape[4]
        k = ct.kernel_size
        return (D + k[0] - 1, H + k[1] - 1, W + k[2] - 1)

    def _get_weight_spectrum(self, s):
        ct = self.conv_transpose
        w = ct.weight
        key = (w._version, w.data_ptr(), s)
        if self._w_key != key or self._w_spec is None:
            wp = w.detach().permute(2, 3, 4, 0, 1).contiguous()  # [kd,kh,kw,In,Out]
            spec = torch.fft.rfftn(wp, s=s, dim=(0, 1, 2))       # [D,H,Wh,In,Out]
            F = spec.shape[0] * spec.shape[1] * spec.shape[2]
            self._w_spec = spec.reshape(F, spec.shape[3], spec.shape[4]).contiguous()
            self._w_key = key
        return self._w_spec

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

        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]
        In, Out = ct.in_channels, ct.out_channels
        gamma = bn.weight if bn.affine else x.new_empty(0)
        beta = bn.bias if bn.affine else x.new_empty(0)
        bias = ct.bias if ct.bias is not None else x.new_empty(0)

        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 not self._simple_conv:
            # Generic fallback: run the conv, fuse the rest.
            y = ct(x)
            C = y.shape[1]
            S = y.shape[2] * y.shape[3] * y.shape[4]
            yv = (y * self.scale_factor).contiguous()
            sums = _ext.spatial_sums(yv)
            sumsq = _ext.spatial_sums(yv * yv)
            if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
                bn.num_batches_tracked += 1
            out = _ext.finalize(sums, sumsq, x.new_empty(0), rmean, rvar,
                                gamma, beta, 1.0, float(S), momentum,
                                bn.eps, bn.training)
            return out.view(B, C, 1, 1, 1)

        Do, Ho, Wo = self._out_spatial(x)
        S = Do * Ho * Wo
        Wh = Wo // 2 + 1
        F = Do * Ho * Wh

        if (not bn.training) and bn.track_running_stats:
            # ---- closed-form eval path: no convolution ---------------------
            Sx = _ext.spatial_sums(x)               # [B, In]
            Sw = _ext.spatial_sums(ct.weight)       # [In, Out]
            out = _ext.eval_finalize(Sx, Sw, bias, rmean, rvar, gamma, beta,
                                     self.scale_factor, 1.0 / float(S), bn.eps)
            return out.view(B, Out, 1, 1, 1)

        # -------------------- FFT path (train / no-track eval) --------------
        # Custom fused kernel constraints.
        if not (B * Out <= 2048 and In <= 64 and Out % 32 == 0):
            return self._reference_forward(x)

        s = (Do, Ho, Wo)
        xp = x.detach().permute(2, 3, 4, 0, 1)                     # [D,H,W,B,In]
        Xf = torch.fft.rfftn(xp, s=s, dim=(0, 1, 2))               # [D,H,Wh,B,In]
        Xf = Xf.reshape(F, B, In)
        Wf = self._get_weight_spectrum(s)                          # [F, In, Out]

        dc_blk = Ho * Wh                                           # d == 0 plane
        nyq_w = 2.0 if (Wo % 2 == 1) else 1.0
        sums, sumsq = _ext.freq_gemm_parseval(
            Xf, Wf, B, In, Out, dc_blk, nyq_w, 1.0 / float(S))

        if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
            bn.num_batches_tracked += 1
        out = _ext.finalize(sums.view(B, Out), sumsq.view(B, Out), bias,
                            rmean, rvar, gamma, beta, self.scale_factor,
                            float(S), momentum, bn.eps, bn.training)
        return out.view(B, Out, 1, 1, 1)
```