```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

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

#define TPB 256

// ---------------------------------------------------------------------------
// Fused ConvTranspose3d (kernel=3, stride=2, pad=1, dilation=1, groups=1)
//   + bias
//   + per-(n, oc) sum / sum-of-squares accumulation (for the epilogue stats)
//
// out[n, oc, od, oh, ow] = bias[oc]
//   + sum_{ic, kd, kh, kw} w[ic, oc, kd, kh, kw]
//        * x[n, ic, (od+1-kd)/2, (oh+1-kh)/2, (ow+1-kw)/2]
// where only terms with (o + 1 - k) even and in-range contribute.
//
// grid: (ceil(Sout/(TPB*4)), OC, N), block: TPB
// ---------------------------------------------------------------------------
template <int IC>
__global__ void convt3d_s2k3p1_kernel(
        const float* __restrict__ x,
        const float* __restrict__ w,
        const float* __restrict__ bias,
        float* __restrict__ out,
        float* __restrict__ sums,
        float* __restrict__ sumsq,
        int ic_in, int OC,
        int Di, int Hi, int Wi,
        int Do, int Ho, int Wo,
        int has_bias) {
    const int icn = (IC > 0) ? IC : ic_in;
    const int oc = blockIdx.y;
    const int n  = blockIdx.z;
    const int Sout = Do * Ho * Wo;

    // weights for this output channel -> shared memory: wsh[ic*27 + t]
    extern __shared__ float wsh[];
    for (int idx = threadIdx.x; idx < icn * 27; idx += TPB) {
        int ic = idx / 27;
        int t  = idx - ic * 27;
        wsh[idx] = w[((long)ic * OC + oc) * 27 + t];
    }
    __syncthreads();

    const long plane = (long)Di * Hi * Wi;
    const float* xn = x + (long)n * icn * plane;
    float* outn = out + ((long)n * OC + oc) * Sout;
    const float b = has_bias ? bias[oc] : 0.0f;

    float psum = 0.0f, psum2 = 0.0f;
    const int start = blockIdx.x * (TPB * 4) + threadIdx.x;

    #pragma unroll
    for (int r = 0; r < 4; ++r) {
        const int p = start + r * TPB;
        if (p < Sout) {
            const int ow = p % Wo;
            const int t1 = p / Wo;
            const int oh = t1 % Ho;
            const int od = t1 / Ho;

            // valid (input index, kernel tap) pairs per dimension
            int id_[2], kd_[2], nd_ = 0;
            if ((od & 1) == 0) { int t = od >> 1;       if (t < Di) { id_[nd_] = t; kd_[nd_] = 1; ++nd_; } }
            else { int t0 = (od + 1) >> 1;              if (t0 < Di) { id_[nd_] = t0; kd_[nd_] = 0; ++nd_; }
                   int t2 = (od - 1) >> 1;              if (t2 >= 0) { id_[nd_] = t2; kd_[nd_] = 2; ++nd_; } }
            int ih_[2], kh_[2], nh_ = 0;
            if ((oh & 1) == 0) { int t = oh >> 1;       if (t < Hi) { ih_[nh_] = t; kh_[nh_] = 1; ++nh_; } }
            else { int t0 = (oh + 1) >> 1;              if (t0 < Hi) { ih_[nh_] = t0; kh_[nh_] = 0; ++nh_; }
                   int t2 = (oh - 1) >> 1;              if (t2 >= 0) { ih_[nh_] = t2; kh_[nh_] = 2; ++nh_; } }
            int iw_[2], kw_[2], nw_ = 0;
            if ((ow & 1) == 0) { int t = ow >> 1;       if (t < Wi) { iw_[nw_] = t; kw_[nw_] = 1; ++nw_; } }
            else { int t0 = (ow + 1) >> 1;              if (t0 < Wi) { iw_[nw_] = t0; kw_[nw_] = 0; ++nw_; }
                   int t2 = (ow - 1) >> 1;              if (t2 >= 0) { iw_[nw_] = t2; kw_[nw_] = 2; ++nw_; } }

            float acc = b;
            for (int a = 0; a < nd_; ++a) {
                const int ida = id_[a], kda = kd_[a];
                for (int bb = 0; bb < nh_; ++bb) {
                    const int ihb = ih_[bb], khb = kh_[bb];
                    for (int g = 0; g < nw_; ++g) {
                        const long off = ((long)ida * Hi + ihb) * Wi + iw_[g];
                        const int tp = (kda * 3 + khb) * 3 + kw_[g];
                        const float* wt = wsh + tp;
                        const float* xt = xn + off;
                        #pragma unroll
                        for (int ic = 0; ic < icn; ++ic) {
                            acc += wt[ic * 27] * xt[ic * plane];
                        }
                    }
                }
            }
            outn[p] = acc;
            psum  += acc;
            psum2 += acc * acc;
        }
    }

    __shared__ float red[TPB];
    __shared__ float red2[TPB];
    red[threadIdx.x]  = psum;
    red2[threadIdx.x] = psum2;
    __syncthreads();
    for (int s = TPB / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            red[threadIdx.x]  += red[threadIdx.x + s];
            red2[threadIdx.x] += red2[threadIdx.x + s];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        atomicAdd(&sums[blockIdx.z * OC + oc], red[0]);
        atomicAdd(&sumsq[blockIdx.z * OC + oc], red2[0]);
    }
}

// ---------------------------------------------------------------------------
// Standalone partial sums / sum-of-squares per (n, c) plane (fallback path).
// grid: (chunks, N*C), block: TPB
// ---------------------------------------------------------------------------
__global__ void partial_sums_kernel(
        const float* __restrict__ x,
        float* __restrict__ sums,
        float* __restrict__ sumsq,
        long S) {
    const int nc = blockIdx.y;
    const long base = (long)nc * S;
    const long chunk = (S + gridDim.x - 1) / gridDim.x;
    const long start = (long)blockIdx.x * chunk;
    const long end = min(start + chunk, S);

    float acc = 0.0f, acc2 = 0.0f;
    for (long i = start + threadIdx.x; i < end; i += TPB) {
        float v = x[base + i];
        acc  += v;
        acc2 += v * v;
    }
    __shared__ float red[TPB];
    __shared__ float red2[TPB];
    red[threadIdx.x]  = acc;
    red2[threadIdx.x] = acc2;
    __syncthreads();
    for (int s = TPB / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            red[threadIdx.x]  += red[threadIdx.x + s];
            red2[threadIdx.x] += red2[threadIdx.x + s];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        atomicAdd(&sums[nc], red[0]);
        atomicAdd(&sumsq[nc], red2[0]);
    }
}

// ---------------------------------------------------------------------------
// Tiny per-channel stats kernel.
//   spmean[n*C+c] = spatial mean of conv output for (n, c)
//   scale[c]      = batch-norm scale (batch stats or running stats)
//   optionally updates running mean / variance in-place.
// ---------------------------------------------------------------------------
__global__ void stats_kernel(
        const float* __restrict__ sums,
        const float* __restrict__ sumsq,
        float* __restrict__ spmean,
        float* __restrict__ scale,
        const float* __restrict__ weight,
        float* __restrict__ running_mean,
        float* __restrict__ running_var,
        int N, int C, long S,
        float eps, float expo,
        int use_batch, int update_running, int affine) {
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;

    const float invS = 1.0f / (float)S;
    for (int n = 0; n < N; ++n) {
        spmean[n * C + c] = sums[n * C + c] * invS;
    }

    const float w = affine ? weight[c] : 1.0f;
    if (use_batch) {
        double total = 0.0, total2 = 0.0;
        for (int n = 0; n < N; ++n) {
            total  += (double)sums[n * C + c];
            total2 += (double)sumsq[n * C + c];
        }
        const double count = (double)N * (double)S;
        const double mean = total / count;
        double var = total2 / count - mean * mean;
        if (var < 0.0) var = 0.0;
        scale[c] = w / sqrtf((float)var + eps);
        if (update_running) {
            running_mean[c] = (1.0f - expo) * running_mean[c] + expo * (float)mean;
            const double unb = var * (count / (count - 1.0));
            running_var[c]  = (1.0f - expo) * running_var[c] + expo * (float)unb;
        }
    } else {
        scale[c] = w / sqrtf(running_var[c] + eps);
    }
}

// ---------------------------------------------------------------------------
// out[n, c, s] = scale[c] * (x[n, c, s] - spmean[n*C+c])
// (batch-norm shift and channel-mean terms cancel under spatial-mean
//  subtraction, so this is exactly bn(x) - mean_spatial(bn(x)))
// grid: (ceil(S/(TPB*4)), N*C), block: TPB
// ---------------------------------------------------------------------------
__global__ void apply_kernel(
        const float* __restrict__ x,
        const float* __restrict__ spmean,
        const float* __restrict__ scale,
        float* __restrict__ out,
        int C, long S) {
    const int nc = blockIdx.y;
    const int c  = nc % C;
    const long base = (long)nc * S;
    const float m = spmean[nc];
    const float sc = scale[c];
    const long start = (long)blockIdx.x * (TPB * 4) + threadIdx.x;
    #pragma unroll
    for (int r = 0; r < 4; ++r) {
        const long i = start + (long)r * TPB;
        if (i < S) {
            out[base + i] = sc * (x[base + i] - m);
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
void convt3d_fused(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                   torch::Tensor out, torch::Tensor sums, torch::Tensor sumsq,
                   long IC, long OC,
                   long Di, long Hi, long Wi,
                   long Do, long Ho, long Wo, long has_bias) {
    const long Sout = Do * Ho * Wo;
    const int N = x.size(0);
    dim3 grid((unsigned)((Sout + TPB * 4 - 1) / (TPB * 4)), (unsigned)OC, (unsigned)N);
    dim3 block(TPB);
    const size_t shmem = (size_t)IC * 27 * sizeof(float);
    const float* bias_ptr = has_bias ? bias.data_ptr<float>() : nullptr;
    if (IC == 16) {
        convt3d_s2k3p1_kernel<16><<<grid, block, shmem>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(),
            (int)IC, (int)OC, (int)Di, (int)Hi, (int)Wi, (int)Do, (int)Ho, (int)Wo,
            (int)has_bias);
    } else {
        convt3d_s2k3p1_kernel<0><<<grid, block, shmem>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(),
            (int)IC, (int)OC, (int)Di, (int)Hi, (int)Wi, (int)Do, (int)Ho, (int)Wo,
            (int)has_bias);
    }
}

void partial_sums(torch::Tensor x, torch::Tensor sums, torch::Tensor sumsq, long S) {
    const long NC = sums.numel();
    long chunks = (S + 8191) / 8192;
    if (chunks < 1) chunks = 1;
    if (chunks > 1024) chunks = 1024;
    dim3 grid((unsigned)chunks, (unsigned)NC);
    dim3 block(TPB);
    partial_sums_kernel<<<grid, block>>>(
        x.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(), S);
}

void stats_op(torch::Tensor sums, torch::Tensor sumsq,
              torch::Tensor spmean, torch::Tensor scale,
              torch::Tensor weight, torch::Tensor running_mean,
              torch::Tensor running_var,
              long N, long C, long S,
              double eps, double expo,
              long use_batch, long update_running, long affine) {
    const int threads = 128;
    const int blocks = (int)((C + threads - 1) / threads);
    stats_kernel<<<blocks, threads>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        spmean.data_ptr<float>(), scale.data_ptr<float>(),
        affine ? weight.data_ptr<float>() : nullptr,
        (update_running || !use_batch) ? running_mean.data_ptr<float>() : nullptr,
        (update_running || !use_batch) ? running_var.data_ptr<float>() : nullptr,
        (int)N, (int)C, S, (float)eps, (float)expo,
        (int)use_batch, (int)update_running, (int)affine);
}

void apply_op(torch::Tensor x, torch::Tensor spmean, torch::Tensor scale,
              torch::Tensor out, long C, long S) {
    const long NC = spmean.numel();
    dim3 grid((unsigned)((S + TPB * 4 - 1) / (TPB * 4)), (unsigned)NC);
    dim3 block(TPB);
    apply_kernel<<<grid, block>>>(
        x.data_ptr<float>(), spmean.data_ptr<float>(), scale.data_ptr<float>(),
        out.data_ptr<float>(), (int)C, S);
}
"""

_cpp_src = r"""
void convt3d_fused(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                   torch::Tensor out, torch::Tensor sums, torch::Tensor sumsq,
                   long IC, long OC,
                   long Di, long Hi, long Wi,
                   long Do, long Ho, long Wo, long has_bias);
void partial_sums(torch::Tensor x, torch::Tensor sums, torch::Tensor sumsq, long S);
void stats_op(torch::Tensor sums, torch::Tensor sumsq,
              torch::Tensor spmean, torch::Tensor scale,
              torch::Tensor weight, torch::Tensor running_mean,
              torch::Tensor running_var,
              long N, long C, long S,
              double eps, double expo,
              long use_batch, long update_running, long affine);
void apply_op(torch::Tensor x, torch::Tensor spmean, torch::Tensor scale,
              torch::Tensor out, long C, long S);
"""

_fused_ops = load_inline(
    name="fused_convt_bn_meansub",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["convt3d_fused", "partial_sums", "stats_op", "apply_op"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d + BatchNorm3d + spatial-mean subtraction, with:
      * a fused gather-based ConvTranspose3d CUDA kernel (k=3, s=2, p=1) that
        also produces per-(batch, channel) spatial sums / sums-of-squares
      * a tiny stats kernel computing batch-norm scale and spatial means
      * a fused epilogue kernel: out = scale[c] * (x - spatial_mean[n, c])
        (valid because bn(x) - mean(bn(x)) == scale * (x - mean(x)))
    Falls back to the reference ops whenever the configuration is unsupported.
    """

    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 _custom_ok(self, x, ct):
        return (
            x.is_cuda and x.dtype == torch.float32 and
            ct.groups == 1 and
            tuple(ct.kernel_size) == (3, 3, 3) and
            tuple(ct.stride) == (2, 2, 2) and
            tuple(ct.padding) == (1, 1, 1) and
            tuple(ct.dilation) == (1, 1, 1) and
            tuple(ct.output_padding) == (0, 0, 0)
        )

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

        if not self._custom_ok(x, ct):
            out = ct(x)
            out = bn(out)
            return out - torch.mean(out, dim=(2, 3, 4), keepdim=True)

        x = x.contiguous()
        N, IC, Di, Hi, Wi = x.shape
        OC = ct.out_channels
        Do = (Di - 1) * 2 - 2 + 3
        Ho = (Hi - 1) * 2 - 2 + 3
        Wo = (Wi - 1) * 2 - 2 + 3
        S = Do * Ho * Wo

        dev = x.device
        sums = torch.zeros(N * OC, device=dev, dtype=torch.float32)
        sumsq = torch.zeros(N * OC, device=dev, dtype=torch.float32)
        y = torch.empty((N, OC, Do, Ho, Wo), device=dev, dtype=torch.float32)

        has_bias = 1 if ct.bias is not None else 0
        bias_t = ct.bias if ct.bias is not None else sums  # dummy when unused
        _fused_ops.convt3d_fused(
            x, ct.weight, bias_t, y, sums, sumsq,
            IC, OC, Di, Hi, Wi, Do, Ho, Wo, has_bias)

        use_batch = 1 if (bn.training or not bn.track_running_stats) else 0
        update_running = 1 if (bn.training and bn.track_running_stats) else 0
        expo = float(bn.momentum) if bn.momentum is not None else 0.0
        if update_running:
            with torch.no_grad():
                bn.num_batches_tracked.add_(1)
            if bn.momentum is None:
                expo = 1.0 / float(bn.num_batches_tracked.item())

        affine = 1 if bn.weight is not None else 0
        weight_t = bn.weight if bn.weight is not None else sums  # dummy
        if bn.track_running_stats:
            rm, rv = bn.running_mean, bn.running_var
        else:
            rm = rv = sums  # dummies

        spmean = torch.empty(N * OC, device=dev, dtype=torch.float32)
        scale = torch.empty(OC, device=dev, dtype=torch.float32)
        _fused_ops.stats_op(
            sums, sumsq, spmean, scale, weight_t, rm, rv,
            N, OC, S, float(bn.eps), expo, use_batch, update_running, affine)

        out = torch.empty_like(y)
        _fused_ops.apply_op(y, spmean, scale, out, OC, S)
        return out
```