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

# ---------------------------------------------------------------------------
# Fused kernel: BatchNorm3d + spatial mean subtraction.
#
# Math note: for the composition  y = BN(x);  out = y - mean_spatial(y)
#   BN(x)_i      = g * (x_i - mu) / sqrt(var + eps) + b
#   mean_s(BN(x)) = g * (mean_s(x) - mu) / sqrt(var + eps) + b
#   out_i        = (g / sqrt(var + eps)) * (x_i - mean_s(x))
# The batch mean (mu) and bias (b) cancel exactly, so the whole post-conv
# pipeline collapses to a single fused kernel given per-channel scale
#   scale[c] = weight[c] / sqrt(var[c] + eps)
#
# The kernel uses one thread block per (batch, channel) slice: it first
# performs a block-wide reduction (double precision accumulation) to obtain
# the spatial mean, then writes scale * (x - mean) in a second pass.
# ---------------------------------------------------------------------------

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

__global__ void fused_scale_sub_spatial_mean_kernel(
    const float* __restrict__ x,
    const float* __restrict__ scale,
    float* __restrict__ out,
    const int C,
    const long long spatial)
{
    const long long base = (long long)blockIdx.x * spatial;
    const float* xs = x + base;
    float* os = out + base;
    const float s = scale[blockIdx.x % C];

    const long long stride = blockDim.x;

    // Pass 1: block reduction of the spatial sum (4x unrolled, double accum).
    double local = 0.0;
    long long i = threadIdx.x;
    for (; i + 3 * stride < spatial; i += 4 * stride) {
        const float v0 = xs[i];
        const float v1 = xs[i + stride];
        const float v2 = xs[i + 2 * stride];
        const float v3 = xs[i + 3 * stride];
        local += ((double)v0 + (double)v1) + ((double)v2 + (double)v3);
    }
    for (; i < spatial; i += stride) {
        local += (double)xs[i];
    }

    __shared__ double sh[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        local += __shfl_down_sync(0xffffffffu, local, off);
    }
    if (lane == 0) { sh[wid] = local; }
    __syncthreads();
    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        double v = (lane < nwarps) ? sh[lane] : 0.0;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
        }
        if (lane == 0) { sh[0] = v; }
    }
    __syncthreads();
    const float mean = (float)(sh[0] / (double)spatial);

    // Pass 2: write scale * (x - mean), 4x unrolled.
    i = threadIdx.x;
    for (; i + 3 * stride < spatial; i += 4 * stride) {
        os[i]             = s * (xs[i]             - mean);
        os[i + stride]     = s * (xs[i + stride]     - mean);
        os[i + 2 * stride] = s * (xs[i + 2 * stride] - mean);
        os[i + 3 * stride] = s * (xs[i + 3 * stride] - mean);
    }
    for (; i < spatial; i += stride) {
        os[i] = s * (xs[i] - mean);
    }
}

torch::Tensor fused_bn_spatial_demean(torch::Tensor x, 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_CHECK(x.dim() >= 3, "x must have at least 3 dims");
    TORCH_CHECK(scale.is_contiguous(), "scale must be contiguous");

    auto out = torch::empty_like(x);
    const int C = (int)x.size(1);
    long long spatial = 1;
    for (int d = 2; d < x.dim(); ++d) { spatial *= x.size(d); }
    const long long nslices = x.numel() / spatial;   // N * C

    const int threads = 256;
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_scale_sub_spatial_mean_kernel<<<(int)nslices, threads, 0, stream>>>(
        x.data_ptr<float>(),
        scale.data_ptr<float>(),
        out.data_ptr<float>(),
        C,
        spatial);
    return out;
}
"""

_fused_cpp_source = (
    "torch::Tensor fused_bn_spatial_demean(torch::Tensor x, torch::Tensor scale);"
)

_fused_ext = load_inline(
    name="fused_bn_spatial_demean_ext",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_source,
    functions=["fused_bn_spatial_demean"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class _FusedScaleSubSpatialMean(torch.autograd.Function):
    """Autograd wrapper around the fused kernel.

    forward : out[n,c,i] = scale[c] * (x[n,c,i] - mean_spatial(x[n,c,:]))
    backward: grad_x     = scale * (g - mean_spatial(g))
              grad_scale = sum_{n,spatial} g * (x - mean_spatial(x))
    """

    @staticmethod
    def forward(ctx, x, scale):
        out = _fused_ext.fused_bn_spatial_demean(x, scale)
        ctx.save_for_backward(x, scale)
        return out

    @staticmethod
    def backward(ctx, grad_out):
        x, scale = ctx.saved_tensors
        g = grad_out.contiguous()
        mean_x = x.mean(dim=(2, 3, 4), keepdim=True)
        grad_scale = (g * (x - mean_x)).sum(dim=(0, 2, 3, 4))
        mean_g = g.mean(dim=(2, 3, 4), keepdim=True)
        grad_x = scale.view(1, -1, 1, 1, 1) * (g - mean_g)
        return grad_x, grad_scale


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cuDNN) + fused BatchNorm3d/spatial-mean-subtraction
    implemented with a custom CUDA kernel.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        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 forward(self, x):
        x = self.conv_transpose(x)
        bn = self.batch_norm

        if x.is_cuda and x.dtype == torch.float32 and bn.track_running_stats:
            x = x.contiguous()
            if bn.training:
                # Batch statistics over (N, D, H, W) per channel.
                M = x.shape[0]
                for d in (2, 3, 4):
                    M *= x.shape[d]
                var_unbiased, mean = torch.var_mean(x, dim=(0, 2, 3, 4))
                # Normalization uses the biased variance.
                var = var_unbiased * ((M - 1.0) / float(M))
                # Mirror nn.BatchNorm3d running-stat updates.
                with torch.no_grad():
                    bn._num_batches_tracked.add_(1)
                    if bn.momentum is None:
                        mom = 1.0 / float(bn._num_batches_tracked.item())
                    else:
                        mom = bn.momentum
                    bn.running_mean.lerp_(mean, mom)
                    bn.running_var.lerp_(var_unbiased, mom)
            else:
                var = bn.running_var

            if bn.affine:
                scale = bn.weight / torch.sqrt(var + bn.eps)
            else:
                scale = torch.rsqrt(var + bn.eps)
            scale = scale.contiguous()

            return _FusedScaleSubSpatialMean.apply(x, scale)

        # Reference fallback for non-CUDA / non-float32 / no-running-stats.
        x = self.batch_norm(x)
        return x - torch.mean(x, dim=(2, 3, 4), keepdim=True)
```