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

# Key idea:
#   The network ends with ConvTranspose3d -> scale -> BatchNorm3d -> global avg pool.
#   Global average pooling only needs the *spatial sum* of the conv-transpose output.
#   For a stride-1, pad-0 transposed convolution, every (input element, kernel tap)
#   pair contributes exactly once to the output, so:
#       sum_spatial(out[b,o]) = sum_ci sum_spatial(x[b,ci]) * sum_k(W[ci,o,k])
#                               + bias[o] * S_out
#   Hence in eval mode (batch-norm uses running stats) the whole pipeline collapses
#   to: one reduction over the input + a tiny (C_in x C_out) matmul + a tiny affine.
#   In training mode we still run the cuDNN conv-transpose, but fuse
#   scale+BN-stats+BN-apply+pool into two small kernels.

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

#define TPB 256

// ---------------------------------------------------------------------------
// Per-(batch, channel) spatial sum of a contiguous (B, C, D, H, W) tensor.
// One thread block per (b, c) pair.
// ---------------------------------------------------------------------------
__global__ void spatial_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   long S) {
    const float* base = x + (long)blockIdx.x * S;
    float acc = 0.0f;
    if ((S & 3) == 0) {
        long n4 = S >> 2;
        const float4* b4 = reinterpret_cast<const float4*>(base);
        for (long i = threadIdx.x; i < n4; i += TPB) {
            float4 v = b4[i];
            acc += v.x + v.y + v.z + v.w;
        }
    } else {
        for (long i = threadIdx.x; i < S; i += TPB) acc += base[i];
    }
    __shared__ float sh[TPB / 32];
    for (int o = 16; o > 0; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o);
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = acc;
    __syncthreads();
    if (wid == 0) {
        acc = (lane < TPB / 32) ? sh[lane] : 0.0f;
        for (int o = 16; o > 0; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o);
        if (lane == 0) out[blockIdx.x] = acc;
    }
}

// ---------------------------------------------------------------------------
// Per (C_in, C_out) sum over the K^3 kernel taps of the conv-transpose weight.
// ---------------------------------------------------------------------------
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ wsum,
                                  long total, int K3) {
    long idx = (long)blockIdx.x * TPB + threadIdx.x;
    if (idx >= total) return;
    const float* base = w + idx * (long)K3;
    float acc = 0.0f;
    for (int k = 0; k < K3; ++k) acc += base[k];
    wsum[idx] = acc;
}

// ---------------------------------------------------------------------------
// Eval-mode fused kernel:
//   outsum[b,o] = sum_ci insum[b,ci] * wsum[ci,o]
//   mean        = scale * (outsum / S_out + bias)
//   out         = (mean - running_mean) / sqrt(running_var + eps) * gamma + beta
// ---------------------------------------------------------------------------
__global__ void eval_output_kernel(const float* __restrict__ insum,
                                   const float* __restrict__ wsum,
                                   const float* __restrict__ bias,
                                   const float* __restrict__ rmean,
                                   const float* __restrict__ rvar,
                                   const float* __restrict__ gamma,
                                   const float* __restrict__ beta,
                                   float* __restrict__ out,
                                   float scale, float inv_spatial, float eps,
                                   int B, int C_in, int C_out) {
    int idx = blockIdx.x * TPB + threadIdx.x;
    if (idx >= B * C_out) return;
    int b = idx / C_out;
    int o = idx - b * C_out;
    const float* in_row = insum + (long)b * C_in;
    float acc = 0.0f;
    for (int ci = 0; ci < C_in; ++ci) acc += in_row[ci] * wsum[(long)ci * C_out + o];
    float mean = acc * inv_spatial;
    if (bias != nullptr) mean += bias[o];
    mean *= scale;
    float invstd = rsqrtf(rvar[o] + eps);
    out[idx] = (mean - rmean[o]) * invstd * gamma[o] + beta[o];
}

// ---------------------------------------------------------------------------
// Training-mode reduction over the (already computed) conv-transpose output:
// per-(b,c) spatial sums plus per-channel sum / sum-of-squares (for batch stats).
// ---------------------------------------------------------------------------
__global__ void train_reduce_kernel(const float* __restrict__ x,
                                    float* __restrict__ bcsum,
                                    double* __restrict__ chsum,
                                    double* __restrict__ chsumsq,
                                    long S, int C) {
    long bc = blockIdx.x;
    const float* base = x + bc * S;
    double s = 0.0, q = 0.0;
    if ((S & 3) == 0) {
        long n4 = S >> 2;
        const float4* b4 = reinterpret_cast<const float4*>(base);
        for (long i = threadIdx.x; i < n4; i += TPB) {
            float4 v = b4[i];
            s += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
            q += (double)v.x * v.x + (double)v.y * v.y + (double)v.z * v.z + (double)v.w * v.w;
        }
    } else {
        for (long i = threadIdx.x; i < S; i += TPB) {
            double v = (double)base[i];
            s += v;
            q += v * v;
        }
    }
    __shared__ double shs[TPB / 32];
    __shared__ double shq[TPB / 32];
    for (int o = 16; o > 0; o >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, o);
        q += __shfl_down_sync(0xffffffffu, q, o);
    }
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) { shs[wid] = s; shq[wid] = q; }
    __syncthreads();
    if (wid == 0) {
        s = (lane < TPB / 32) ? shs[lane] : 0.0;
        q = (lane < TPB / 32) ? shq[lane] : 0.0;
        for (int o = 16; o > 0; o >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, o);
            q += __shfl_down_sync(0xffffffffu, q, o);
        }
        if (lane == 0) {
            bcsum[bc] = (float)s;
            int c = (int)(bc % (long)C);
            atomicAdd(&chsum[c], s);
            atomicAdd(&chsumsq[c], q);
        }
    }
}

// ---------------------------------------------------------------------------
// Training-mode finalize: batch stats of y = scale * conv, BN affine on the
// pooled means, and running-stat update (once per channel).
// ---------------------------------------------------------------------------
__global__ void train_output_kernel(const float* __restrict__ bcsum,
                                    const double* __restrict__ chsum,
                                    const double* __restrict__ chsumsq,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ beta,
                                    float* __restrict__ rmean,
                                    float* __restrict__ rvar,
                                    float* __restrict__ out,
                                    float scale, float inv_spatial, float eps,
                                    float momentum, long N, int B, int C) {
    int idx = blockIdx.x * TPB + threadIdx.x;
    if (idx >= B * C) return;
    int c = idx % C;
    double mean = chsum[c] / (double)N;
    double var = chsumsq[c] / (double)N - mean * mean;
    if (var < 0.0) var = 0.0;
    double mean_y = (double)scale * mean;
    double var_y = (double)scale * (double)scale * var;
    float invstd = rsqrtf((float)var_y + eps);
    if (rmean != nullptr && idx < C) {
        double unb = (N > 1) ? var_y * (double)N / (double)(N - 1) : var_y;
        rmean[c] = rmean[c] * (1.0f - momentum) + (float)mean_y * momentum;
        rvar[c] = rvar[c] * (1.0f - momentum) + (float)unb * momentum;
    }
    float ymean = scale * bcsum[idx] * inv_spatial;
    out[idx] = (ymean - (float)mean_y) * invstd * gamma[c] + beta[c];
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
torch::Tensor fused_eval_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 torch::Tensor running_mean, torch::Tensor running_var,
                                 torch::Tensor bn_weight, torch::Tensor bn_bias,
                                 double scale, double eps) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    auto xc = x.contiguous();
    auto w = weight.contiguous();
    auto g = bn_weight.contiguous();
    auto be = bn_bias.contiguous();
    auto rm = running_mean.contiguous();
    auto rv = running_var.contiguous();
    torch::Tensor bs = bias;
    const float* bptr = nullptr;
    if (bs.defined() && bs.numel() > 0) {
        bs = bs.contiguous();
        bptr = bs.data_ptr<float>();
    }

    long B = xc.size(0), C_in = xc.size(1);
    long S = xc.size(2) * xc.size(3) * xc.size(4);
    long C_out = w.size(1);
    long K3 = w.size(2) * w.size(3) * w.size(4);
    long S_out = (xc.size(2) + w.size(2) - 1) *
                 (xc.size(3) + w.size(3) - 1) *
                 (xc.size(4) + w.size(4) - 1);

    auto opts = xc.options();
    auto insum = torch::empty({B, C_in}, opts);
    auto wsum = torch::empty({C_in, C_out}, opts);
    auto out = torch::empty({B, C_out, 1, 1, 1}, opts);

    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sum_kernel<<<(int)(B * C_in), TPB, 0, stream>>>(
        xc.data_ptr<float>(), insum.data_ptr<float>(), S);
    long tot = C_in * C_out;
    weight_sum_kernel<<<(int)((tot + TPB - 1) / TPB), TPB, 0, stream>>>(
        w.data_ptr<float>(), wsum.data_ptr<float>(), tot, (int)K3);
    long tot2 = B * C_out;
    eval_output_kernel<<<(int)((tot2 + TPB - 1) / TPB), TPB, 0, stream>>>(
        insum.data_ptr<float>(), wsum.data_ptr<float>(), bptr,
        rm.data_ptr<float>(), rv.data_ptr<float>(),
        g.data_ptr<float>(), be.data_ptr<float>(),
        out.data_ptr<float>(), (float)scale,
        (float)(1.0 / (double)S_out), (float)eps,
        (int)B, (int)C_in, (int)C_out);
    return out;
}

torch::Tensor fused_train_forward(torch::Tensor conv, torch::Tensor running_mean, torch::Tensor running_var,
                                  torch::Tensor bn_weight, torch::Tensor bn_bias,
                                  double scale, double eps, double momentum, bool update_stats) {
    TORCH_CHECK(conv.is_cuda() && conv.scalar_type() == at::kFloat, "conv must be CUDA float32");
    auto c = conv.contiguous();
    auto g = bn_weight.contiguous();
    auto be = bn_bias.contiguous();

    long B = c.size(0), C = c.size(1);
    long S = c.size(2) * c.size(3) * c.size(4);
    long N = B * S;

    auto opts = c.options();
    auto dopts = opts.dtype(at::kDouble);
    auto bcsum = torch::empty({B, C}, opts);
    auto chsum = torch::zeros({C}, dopts);
    auto chsumsq = torch::zeros({C}, dopts);
    auto out = torch::empty({B, C, 1, 1, 1}, opts);

    float* rmp = nullptr;
    float* rvp = nullptr;
    torch::Tensor rm = running_mean, rv = running_var;
    if (update_stats && rm.defined() && rm.numel() > 0) {
        rm = rm.contiguous();
        rv = rv.contiguous();
        rmp = rm.data_ptr<float>();
        rvp = rv.data_ptr<float>();
    }

    auto stream = at::cuda::getCurrentCUDAStream();
    train_reduce_kernel<<<(int)(B * C), TPB, 0, stream>>>(
        c.data_ptr<float>(), bcsum.data_ptr<float>(),
        chsum.data_ptr<double>(), chsumsq.data_ptr<double>(), S, (int)C);
    long tot = B * C;
    train_output_kernel<<<(int)((tot + TPB - 1) / TPB), TPB, 0, stream>>>(
        bcsum.data_ptr<float>(), chsum.data_ptr<double>(), chsumsq.data_ptr<double>(),
        g.data_ptr<float>(), be.data_ptr<float>(), rmp, rvp,
        out.data_ptr<float>(), (float)scale,
        (float)(1.0 / (double)S), (float)eps, (float)momentum,
        N, (int)B, (int)C);
    return out;
}
"""

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor fused_eval_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 torch::Tensor running_mean, torch::Tensor running_var,
                                 torch::Tensor bn_weight, torch::Tensor bn_bias,
                                 double scale, double eps);
torch::Tensor fused_train_forward(torch::Tensor conv, torch::Tensor running_mean, torch::Tensor running_var,
                                  torch::Tensor bn_weight, torch::Tensor bn_bias,
                                  double scale, double eps, double momentum, bool update_stats);
"""

_ext = load_inline(
    name="fused_convtranspose_bn_gap_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_eval_forward", "fused_train_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    Eval mode (batch-norm uses running statistics):
        ConvTranspose3d + scale + BatchNorm3d + global avg pool is collapsed
        analytically into a single input reduction, a tiny (C_in x C_out)
        contraction, and a tiny affine kernel - the transposed convolution
        itself is never materialized.

    Training mode (batch-norm uses batch statistics):
        The cuDNN conv-transpose is kept, but scale + batch-stat computation +
        normalization + global average pooling are fused into two custom kernels
        (a single full read pass over the conv output instead of ~5 passes).
    """

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

    def _fast_ok(self, x):
        ct = self.conv_transpose
        bn = self.batch_norm
        if not (x.is_cuda and x.dtype == torch.float32):
            return False
        if (ct.groups != 1 or ct.stride != (1, 1, 1) or ct.padding != (0, 0, 0)
                or ct.dilation != (1, 1, 1) or ct.output_padding != (0, 0, 0)):
            return False
        if not bn.affine:
            return False
        if bn.training and bn.track_running_stats and bn.momentum is None:
            return False
        return True

    def forward(self, x):
        ct = self.conv_transpose
        bn = self.batch_norm
        if self._fast_ok(x):
            if bn.training or not bn.track_running_stats:
                conv = ct(x)
                if bn.track_running_stats:
                    rm, rv = bn.running_mean, bn.running_var
                else:
                    rm = rv = x.new_empty(0)
                return _ext.fused_train_forward(
                    conv, rm, rv, bn.weight, bn.bias,
                    self.scale_factor, float(bn.eps),
                    float(bn.momentum) if bn.momentum is not None else 0.0,
                    bool(bn.training and bn.track_running_stats))
            else:
                bias = ct.bias if ct.bias is not None else x.new_empty(0)
                return _ext.fused_eval_forward(
                    x, ct.weight, bias, bn.running_mean, bn.running_var,
                    bn.weight, bn.bias, self.scale_factor, float(bn.eps))
        # Generic fallback (matches the reference implementation op-for-op).
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x
```