```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 kernel: fused per-(batch, channel) slice reduction.
#
# For a 5D tensor (N, C, D, H, W) it computes, in a single pass over memory:
#     sums[b, c]   = sum over (d, h, w) of x[b, c, d, h, w]
#     sumsqs[b,c]  = sum over (d, h, w) of x[b, c, d, h, w]^2   (optional)
#
# This single kernel replaces the read/write passes of:
#   (a) the scalar multiply          (x * scale_factor)
#   (b) the BatchNorm3d normalization
#   (c) the AdaptiveAvgPool3d((1,1,1)) global reduction
# because all three collapse into a per-channel affine applied to the slice
# sums (BatchNorm is a per-channel affine and the global average pool is a
# linear reduction, so they commute).
# ---------------------------------------------------------------------------
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <vector>

#define BLOCK_THREADS 256
#define NUM_WARPS (BLOCK_THREADS / 32)

template <bool COMPUTE_SQ>
__global__ void slice_reduce_kernel(const float *__restrict__ x,
                                    float *__restrict__ sums,
                                    float *__restrict__ sumsqs,
                                    const long long slice_size) {
    const long long slice = blockIdx.x;
    const float *base = x + slice * slice_size;

    float s = 0.0f;
    float q = 0.0f;

    // Vectorized path (float4) whenever the slice is 16-byte aligned and its
    // length is a multiple of 4.
    if ((((uintptr_t)base) & 0xF) == 0 && (slice_size & 3LL) == 0) {
        const float4 *p4 = reinterpret_cast<const float4 *>(base);
        const long long n4 = slice_size >> 2;
        for (long long i = threadIdx.x; i < n4; i += BLOCK_THREADS) {
            const float4 v = p4[i];
            s += v.x + v.y + v.z + v.w;
            if (COMPUTE_SQ) {
                q += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
            }
        }
    } else {
        for (long long i = threadIdx.x; i < slice_size; i += BLOCK_THREADS) {
            const float v = base[i];
            s += v;
            if (COMPUTE_SQ) {
                q += v * v;
            }
        }
    }

    __shared__ float sh_s[NUM_WARPS];
    __shared__ float sh_q[NUM_WARPS];

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s += __shfl_down_sync(0xFFFFFFFFu, s, off);
        if (COMPUTE_SQ) q += __shfl_down_sync(0xFFFFFFFFu, q, off);
    }

    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        sh_s[warp] = s;
        if (COMPUTE_SQ) sh_q[warp] = q;
    }
    __syncthreads();

    if (warp == 0) {
        s = (lane < NUM_WARPS) ? sh_s[lane] : 0.0f;
        q = (lane < NUM_WARPS) ? sh_q[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(0xFFFFFFFFu, s, off);
            if (COMPUTE_SQ) q += __shfl_down_sync(0xFFFFFFFFu, q, off);
        }
        if (lane == 0) {
            sums[slice] = s;
            if (COMPUTE_SQ) sumsqs[slice] = q;
        }
    }
}

std::vector<torch::Tensor> slice_reduce(torch::Tensor x, bool compute_sq) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 5, "x must be 5D (N, C, D, H, W)");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t slice_size = xc.size(2) * xc.size(3) * xc.size(4);

    auto sums = at::empty({N, C}, xc.options());
    torch::Tensor sumsqs;
    float *sumsq_ptr = nullptr;
    if (compute_sq) {
        sumsqs = at::empty({N, C}, xc.options());
        sumsq_ptr = sumsqs.data_ptr<float>();
    } else {
        sumsqs = at::empty({0}, xc.options());
    }

    const int blocks = (int)(N * C);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (compute_sq) {
        slice_reduce_kernel<true><<<blocks, BLOCK_THREADS, 0, stream>>>(
            xc.data_ptr<float>(), sums.data_ptr<float>(), sumsq_ptr,
            (long long)slice_size);
    } else {
        slice_reduce_kernel<false><<<blocks, BLOCK_THREADS, 0, stream>>>(
            xc.data_ptr<float>(), sums.data_ptr<float>(), sumsq_ptr,
            (long long)slice_size);
    }
    return {sums, sumsqs};
}
"""

_cpp_source = "std::vector<torch::Tensor> slice_reduce(torch::Tensor x, bool compute_sq);"

_ext = load_inline(
    name="fused_scale_bn_gap_reduce",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["slice_reduce"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    Key optimizations:
      1. Eval mode: BatchNorm folds into a per-channel affine which commutes
         with the global average pool.  Additionally, for a stride-1,
         zero-padding transposed convolution (a "full" convolution), the global
         sum of the output factorizes:
             sum_t y[oc, t] = sum_ic (sum_i x[ic, i]) * (sum_j w[ic, oc, j])
         so the transposed convolution never has to be materialized.  The whole
         network reduces to one custom reduction kernel over the input plus a
         tiny (N x Cin) @ (Cin x Cout) matmul.
      2. Training mode: batch statistics are required, so the convolution is
         computed with cuDNN, but the scale + BatchNorm + global-average-pool
         chain is replaced by a single fused pass that produces slice sums and
         sums of squares, from which everything else is derived with tiny ops.
      3. Falls back to the reference ops whenever autograd is required or the
         input is not a CUDA float32 tensor.
    """

    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))
        self._wsum = None
        self._wsum_version = -1

    def _reference(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
        w = ct.weight

        if (not x.is_cuda) or x.dtype != torch.float32:
            return self._reference(x)
        if torch.is_grad_enabled() and (
            x.requires_grad
            or w.requires_grad
            or (bn.weight is not None and bn.weight.requires_grad)
        ):
            return self._reference(x)

        N = x.shape[0]
        sf = self.scale_factor
        gamma = bn.weight
        beta = bn.bias

        if self.training or not bn.track_running_stats:
            # ------------------------- training mode -------------------------
            y = ct(x)
            V = y.shape[2] * y.shape[3] * y.shape[4]
            S, T = _ext.slice_reduce(y, True)
            n = float(N) * float(V)

            mu_y = S.sum(dim=0) / n
            var_y = torch.clamp_min(T.sum(dim=0) / n - mu_y * mu_y, 0.0)
            mu = mu_y * sf
            var = var_y * (sf * sf)

            if bn.track_running_stats:
                with torch.no_grad():
                    if bn.momentum is None:
                        bn.num_batches_tracked.add_(1)
                        f = 1.0 / float(bn.num_batches_tracked)
                    else:
                        f = float(bn.momentum)
                    bn.running_mean.mul_(1.0 - f).add_(mu, alpha=f)
                    bn.running_var.mul_(1.0 - f).add_(var * (n / (n - 1.0)), alpha=f)

            invstd = torch.rsqrt(var + bn.eps)
            if gamma is None:
                gamma = torch.ones_like(invstd)
                beta = torch.zeros_like(invstd)
            alpha = (sf / float(V)) * gamma * invstd
            shift = beta - gamma * invstd * mu
            out = torch.addcmul(shift, S, alpha)
            return out.view(N, -1, 1, 1, 1)

        # --------------------------- eval mode -------------------------------
        invstd = torch.rsqrt(bn.running_var + bn.eps)
        if gamma is None:
            gamma = torch.ones_like(invstd)
            beta = torch.zeros_like(invstd)

        if (
            ct.groups == 1
            and tuple(ct.stride) == (1, 1, 1)
            and tuple(ct.padding) == (0, 0, 0)
            and tuple(ct.dilation) == (1, 1, 1)
        ):
            # Factorized path: the conv output is never materialized.
            D, H, Wd = x.shape[2], x.shape[3], x.shape[4]
            kD, kH, kW = w.shape[2], w.shape[3], w.shape[4]
            V = (D + kD - 1) * (H + kH - 1) * (Wd + kW - 1)

            if (
                self._wsum is None
                or self._wsum_version != w._version
                or self._wsum.device != w.device
            ):
                self._wsum = w.detach().sum(dim=(2, 3, 4))
                self._wsum_version = w._version

            Xs, _ = _ext.slice_reduce(x, False)
            Sy = Xs.matmul(self._wsum)
        else:
            # Generic conv configuration: run cuDNN, then fuse the rest.
            y = ct(x)
            V = y.shape[2] * y.shape[3] * y.shape[4]
            Sy, _ = _ext.slice_reduce(y, False)

        alpha = (sf / float(V)) * gamma * invstd
        shift = beta - gamma * invstd * bn.running_mean
        if ct.bias is not None:
            # The conv bias contributes bias * V to every output slice sum.
            shift = shift + ct.bias * (float(V) * alpha)
        out = torch.addcmul(shift, Sy, alpha)
        return out.view(N, -1, 1, 1, 1)
```