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

# Let cuDNN autotune the (fixed-shape) transposed convolution algorithm.
torch.backends.cudnn.benchmark = True

# ---------------------------------------------------------------------------
# Fused CUDA kernels for the post-convolution part of the network.
#
# Mathematical simplification used here:
#   BatchNorm is an affine per-channel transform: z = a_c * x + b_c
#   The final op subtracts the per-(batch, channel) spatial mean of z:
#       out = z - mean_s(z) = a_c * (x - mean_s(x))
#   i.e. the BN shift b_c (and the BN mean) cancel out completely!
#   We only need the per-channel scale a_c = gamma / sqrt(var + eps)
#   and the per-plane spatial mean of the raw conv output.
#
# Kernel 1 computes all (N*C) spatial means (and optionally second moments
# for training-mode batch statistics) with one block per plane.
# Kernel 2 applies out = a_c * (x - plane_mean) in a single fused pass,
# replacing PyTorch's separate BN + mean + subtract passes.
# ---------------------------------------------------------------------------

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

#define FULL_MASK 0xffffffffu

// One block per (n, c) plane: compute spatial mean (and mean of squares).
__global__ void plane_stats_kernel(const float* __restrict__ x,
                                   double* __restrict__ mean,
                                   double* __restrict__ sqmean,
                                   int S,
                                   int need_sumsq) {
    const int plane = blockIdx.x;
    const float* __restrict__ xp = x + (long long)plane * (long long)S;

    double sum = 0.0;
    double sq  = 0.0;
    const int stride = blockDim.x;
    for (int i = threadIdx.x; i < S; i += stride) {
        double v = (double)xp[i];
        sum += v;
        if (need_sumsq) sq += v * v;
    }

    // Warp-level reduction (doubles are supported by shuffle intrinsics).
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(FULL_MASK, sum, offset);
        sq  += __shfl_down_sync(FULL_MASK, sq,  offset);
    }

    __shared__ double s_sum[32];
    __shared__ double s_sq[32];
    const int lane   = threadIdx.x & 31;
    const int warp   = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;
    if (lane == 0) {
        s_sum[warp] = sum;
        s_sq[warp]  = sq;
    }
    __syncthreads();

    if (warp == 0) {
        sum = (lane < nwarps) ? s_sum[lane] : 0.0;
        sq  = (lane < nwarps) ? s_sq[lane]  : 0.0;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(FULL_MASK, sum, offset);
            sq  += __shfl_down_sync(FULL_MASK, sq,  offset);
        }
        if (lane == 0) {
            mean[plane] = sum / (double)S;
            if (need_sumsq) sqmean[plane] = sq / (double)S;
        }
    }
}

// Fused apply: out = scale_c * (x - plane_mean).
// grid.y indexes the (n, c) plane, grid.x tiles the spatial extent.
__global__ void normalize_kernel(const float* __restrict__ x,
                                 float* __restrict__ out,
                                 const double* __restrict__ mean,
                                 const float* __restrict__ scale,
                                 int C,
                                 int S) {
    const int plane = blockIdx.y;
    const int c     = plane % C;
    const float m   = (float)mean[plane];
    const float s   = scale[c];

    const long long base = (long long)plane * (long long)S;
    const float* __restrict__ xp = x + base;
    float* __restrict__ op       = out + base;

    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < S; i += stride) {
        op[i] = s * (xp[i] - m);
    }
}

std::vector<torch::Tensor> stats_cuda(torch::Tensor x, bool need_sumsq) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const int64_t C = x.size(1);
    int64_t S = 1;
    for (int i = 2; i < x.dim(); ++i) S *= x.size(i);
    const int64_t planes = x.size(0) * C;

    auto opts = torch::TensorOptions().dtype(torch::kDouble).device(x.device());
    torch::Tensor mean   = torch::empty({planes}, opts);
    torch::Tensor sqmean = torch::empty({need_sumsq ? planes : int64_t(0)}, opts);
    double* sq_ptr = need_sumsq ? sqmean.data_ptr<double>() : nullptr;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    plane_stats_kernel<<<(int)planes, threads, 0, stream>>>(
        x.data_ptr<float>(), mean.data_ptr<double>(), sq_ptr,
        (int)S, need_sumsq ? 1 : 0);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return {mean, sqmean};
}

torch::Tensor normalize_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor scale) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    torch::Tensor out = torch::empty_like(x);

    const int64_t C = x.size(1);
    int64_t S = 1;
    for (int i = 2; i < x.dim(); ++i) S *= x.size(i);
    const int64_t planes = x.size(0) * C;

    const int threads = 256;
    const int64_t blocks_x = (S + threads - 1) / threads;
    dim3 grid((unsigned int)blocks_x, (unsigned int)planes, 1);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    normalize_kernel<<<grid, threads, 0, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(),
        mean.data_ptr<double>(), scale.data_ptr<float>(),
        (int)C, (int)S);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_cpp_src = r"""
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> stats_cuda(torch::Tensor x, bool need_sumsq);
torch::Tensor normalize_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor scale);
"""

_ext = load_inline(
    name="fused_bn_spatial_mean_sub",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["stats_cuda", "normalize_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cuDNN) + fused BatchNorm/spatial-mean-subtraction.

    Key idea: BatchNorm is per-channel affine (z = a_c * x + b_c), and the
    subsequent spatial mean subtraction removes the affine shift entirely:
        z - mean_s(z) = a_c * (x - mean_s(x))
    so BN + mean + subtraction collapse into a single scale-and-center
    operation implemented with two fused custom CUDA kernels.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, bias=bias,
        )
        self.batch_norm = nn.BatchNorm3d(out_channels)

    def _reference_path(self, x):
        y = self.conv_transpose(x)
        y = self.batch_norm(y)
        return y - torch.mean(y, dim=(2, 3, 4), keepdim=True)

    def forward(self, x):
        bn = self.batch_norm
        # Fallback for anything other than the standard fp32 CUDA case.
        if (not x.is_cuda) or x.dtype != torch.float32 or not bn.track_running_stats:
            return self._reference_path(x)

        y = self.conv_transpose(x)
        y = y.contiguous()

        N, C = y.shape[0], y.shape[1]

        if self.training:
            # Need batch statistics (per channel, over N and spatial dims).
            mean, sqmean = _ext.stats_cuda(y, True)
            m_c = mean.view(N, C).mean(dim=0)                       # batch mean
            var_c = (sqmean.view(N, C).mean(dim=0) - m_c * m_c).clamp_min(0.0)
            scale = (bn.weight.double() * torch.rsqrt(var_c + bn.eps)).float()

            # Replicate BatchNorm's running-stat bookkeeping.
            with torch.no_grad():
                S = y.shape[2] * y.shape[3] * y.shape[4]
                count = float(N * S)
                mom = bn.momentum
                if mom is None:
                    mom = 1.0 / float(bn.num_batches_tracked.item() + 1)
                bn.running_mean.lerp_(m_c.float(), mom)
                unbiased = var_c * (count / (count - 1.0))
                bn.running_var.lerp_(unbiased.float(), mom)
                bn.num_batches_tracked.add_(1)
        else:
            mean, _ = _ext.stats_cuda(y, False)
            scale = (bn.weight.double()
                     * torch.rsqrt(bn.running_var.double() + bn.eps)).float()

        return _ext.normalize_cuda(y, mean, scale)
```