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

# Optimizations vs. the previous version:
#  1) The transposed conv is run in channels_last_3d (NHWC), which is the layout
#     cudnn's dgrad engine uses natively. This eliminates the big NHWC->NCHW
#     transpose kernel (~396us) and the input transforms.
#  2) The conv bias-add kernel (~378us) is removed entirely: since
#     BN(y) - mean_spatial(BN(y)) == scale[c] * (g - mean_spatial(g)) with
#     g = conv_transpose_nobias(x) (a per-channel constant cancels exactly),
#     the bias never has to be materialized. The bias parameter is kept for
#     state_dict compatibility and still used on the pure-PyTorch fallback path.
#  3) New fused tail that consumes NHWC directly: a coalesced stats kernel
#     (per-(n,c) spatial means + per-channel BN statistics) and a tiled
#     shared-memory transpose kernel that applies scale*(g-mean) while writing
#     the final NCHW result.
#  4) The previous NCHW fused kernels are kept as a fallback for any layout.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256

// ---------------------------------------------------------------------------
// NCHW fallback kernels (from previous version, with alignment-safe indexing)
// ---------------------------------------------------------------------------
__global__ void bn_ms_stats_kernel(
    const float* __restrict__ y,
    double* __restrict__ slice_sum,
    double* __restrict__ chan_sum,
    double* __restrict__ chan_sumsq,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const long base = (long)slice * slice_elems;

    int head = (int)(base & 3L);
    head = head ? (4 - head) : 0;
    if (head > slice_elems) head = slice_elems;
    const int rem  = slice_elems - head;
    const int vec  = rem >> 2;
    const int tail = rem & 3;

    float s = 0.f, sq = 0.f;
    const float4* y4 = reinterpret_cast<const float4*>(y + base + head);
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < vec; i += gridDim.x * THREADS) {
        const float4 v = y4[i];
        s  += v.x + v.y + v.z + v.w;
        sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const float v = y[base + threadIdx.x];
            s += v; sq += v * v;
        }
        if (threadIdx.x < tail) {
            const float v = y[base + head + ((long)vec << 2) + threadIdx.x];
            s += v; sq += v * v;
        }
    }

    __shared__ double bsum[THREADS];
    __shared__ double bsq[THREADS];
    bsum[threadIdx.x] = (double)s;
    bsq[threadIdx.x]  = (double)sq;
    __syncthreads();
    #pragma unroll
    for (int st = THREADS / 2; st > 0; st >>= 1) {
        if (threadIdx.x < st) {
            bsum[threadIdx.x] += bsum[threadIdx.x + st];
            bsq[threadIdx.x]  += bsq[threadIdx.x + st];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        atomicAdd(&slice_sum[slice], bsum[0]);
        atomicAdd(&chan_sum[oc],     bsum[0]);
        atomicAdd(&chan_sumsq[oc],   bsq[0]);
    }
}

__global__ void bn_ms_apply_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,
    const float* __restrict__ slice_mean,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const float sc = scale[oc];
    const float m  = slice_mean[slice];
    const long base = (long)slice * slice_elems;

    int head = (int)(base & 3L);
    head = head ? (4 - head) : 0;
    if (head > slice_elems) head = slice_elems;
    const int rem  = slice_elems - head;
    const int vec  = rem >> 2;
    const int tail = rem & 3;

    const float4* y4 = reinterpret_cast<const float4*>(y + base + head);
    float4* o4 = reinterpret_cast<float4*>(out + base + head);
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < vec; i += gridDim.x * THREADS) {
        const float4 v = y4[i];
        float4 r;
        r.x = sc * (v.x - m);
        r.y = sc * (v.y - m);
        r.z = sc * (v.z - m);
        r.w = sc * (v.w - m);
        o4[i] = r;
    }
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const long idx = base + threadIdx.x;
            out[idx] = sc * (y[idx] - m);
        }
        if (threadIdx.x < tail) {
            const long idx = base + head + ((long)vec << 2) + threadIdx.x;
            out[idx] = sc * (y[idx] - m);
        }
    }
}

// ---------------------------------------------------------------------------
// Shared finalize kernel: per-slice means + per-channel scales (+running stats)
// ---------------------------------------------------------------------------
__global__ void bn_ms_finalize_kernel(
    const double* __restrict__ slice_sum,
    const double* __restrict__ chan_sum,
    const double* __restrict__ chan_sumsq,
    const float* __restrict__ gamma,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ slice_mean,
    float* __restrict__ scale,
    int NSLICE, int OC,
    double inv_spatial, double inv_total, double total,
    float eps, float momentum,
    int use_batch, int update_running)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < NSLICE) {
        slice_mean[i] = (float)(slice_sum[i] * inv_spatial);
    }
    if (i < OC) {
        float sc;
        if (use_batch) {
            const double mu = chan_sum[i] * inv_total;
            double var = chan_sumsq[i] * inv_total - mu * mu;
            if (var < 0.0) var = 0.0;
            sc = gamma[i] * (float)rsqrt(var + (double)eps);
            if (update_running) {
                running_mean[i] = (1.f - momentum) * running_mean[i] + momentum * (float)mu;
                const double unb = var * (total / (total - 1.0));
                running_var[i]  = (1.f - momentum) * running_var[i] + momentum * (float)unb;
            }
        } else {
            sc = gamma[i] * rsqrtf(running_var[i] + eps);
        }
        scale[i] = sc;
    }
}

// ---------------------------------------------------------------------------
// NHWC (channels_last_3d) kernels. y is logically [N, C, D, H, W] but stored
// as [N, S=D*H*W, C] with C == 32. Requires 16B alignment (C multiple of 4).
// ---------------------------------------------------------------------------

// Pass 1 (NHWC): per-(n,c) spatial sums, per-channel sum / sum-of-squares.
// Block: 256 threads = 32 spatial lanes x 8 float4 channel groups.
__global__ void stats_nhwc_kernel(
    const float* __restrict__ y,
    double* __restrict__ slice_sum,   // [N*C]
    double* __restrict__ chan_sum,    // [C]
    double* __restrict__ chan_sumsq,  // [C]
    int C, long S)
{
    const int n = blockIdx.y;
    const long base = (long)n * S * C;
    const int c0 = (threadIdx.x & 7) << 2;   // 0,4,...,28
    const int s_lane = threadIdx.x >> 3;     // 0..31

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;

    const int tiles = (int)((S + 31) >> 5);
    for (int t = blockIdx.x; t < tiles; t += gridDim.x) {
        const long s = ((long)t << 5) + s_lane;
        if (s < S) {
            const float4 v = *reinterpret_cast<const float4*>(y + base + s * C + c0);
            s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
            q0 = fmaf(v.x, v.x, q0);
            q1 = fmaf(v.y, v.y, q1);
            q2 = fmaf(v.z, v.z, q2);
            q3 = fmaf(v.w, v.w, q3);
        }
    }

    __shared__ double sh_sum[32];
    __shared__ double sh_sq[32];
    if (threadIdx.x < 32) {
        sh_sum[threadIdx.x] = 0.0;
        sh_sq[threadIdx.x]  = 0.0;
    }
    __syncthreads();
    atomicAdd(&sh_sum[c0 + 0], (double)s0);
    atomicAdd(&sh_sum[c0 + 1], (double)s1);
    atomicAdd(&sh_sum[c0 + 2], (double)s2);
    atomicAdd(&sh_sum[c0 + 3], (double)s3);
    atomicAdd(&sh_sq[c0 + 0], (double)q0);
    atomicAdd(&sh_sq[c0 + 1], (double)q1);
    atomicAdd(&sh_sq[c0 + 2], (double)q2);
    atomicAdd(&sh_sq[c0 + 3], (double)q3);
    __syncthreads();

    if (threadIdx.x < 32) {
        const int c = threadIdx.x;
        atomicAdd(&slice_sum[n * C + c], sh_sum[c]);
        atomicAdd(&chan_sum[c],          sh_sum[c]);
        atomicAdd(&chan_sumsq[c],        sh_sq[c]);
    }
}

// Pass 2 (NHWC -> NCHW transpose with fused affine):
//   out[n, c, s] = scale[c] * (y[n, s, c] - slice_mean[n, c])
// 32x32 shared-memory tile, conflict-free in both phases.
__global__ void apply_transpose_nhwc_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,      // [C]
    const float* __restrict__ slice_mean, // [N*C]
    int C, long S)
{
    __shared__ float tile[32][33];
    const int n = blockIdx.y;
    const long s0 = ((long)blockIdx.x << 5);

    const int c0 = (threadIdx.x & 7) << 2;
    const int s_lane = threadIdx.x >> 3;
    const long s = s0 + s_lane;
    if (s < S) {
        const float4 v  = *reinterpret_cast<const float4*>(y + ((long)n * S + s) * C + c0);
        const float4 sc = *reinterpret_cast<const float4*>(scale + c0);
        const float4 mn = *reinterpret_cast<const float4*>(slice_mean + n * C + c0);
        tile[c0 + 0][s_lane] = sc.x * (v.x - mn.x);
        tile[c0 + 1][s_lane] = sc.y * (v.y - mn.y);
        tile[c0 + 2][s_lane] = sc.z * (v.z - mn.z);
        tile[c0 + 3][s_lane] = sc.w * (v.w - mn.w);
    }
    __syncthreads();

    const int c  = threadIdx.x >> 3;         // 0..31
    const int s4 = (threadIdx.x & 7) << 2;   // 0,4,...,28
    float* orow = out + ((long)n * C + c) * S + s0;
    const long rem = S - s0;
    if (rem >= 32) {
        orow[s4 + 0] = tile[c][s4 + 0];
        orow[s4 + 1] = tile[c][s4 + 1];
        orow[s4 + 2] = tile[c][s4 + 2];
        orow[s4 + 3] = tile[c][s4 + 3];
    } else {
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            if (s4 + i < rem) orow[s4 + i] = tile[c][s4 + i];
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
static inline void launch_finalize(
    double* slice_sum_p, double* chan_sum_p, double* chan_sumsq_p,
    float* gamma_p, float* rm_p, float* rv_p,
    float* slice_mean_p, float* scale_p,
    int NSLICE, int OC, int slice_elems, int N,
    float eps, float momentum, bool use_batch, bool update_running,
    cudaStream_t stream)
{
    const double spatial = (double)slice_elems;
    const double total   = spatial * (double)N;
    const int work = NSLICE > OC ? NSLICE : OC;
    bn_ms_finalize_kernel<<<(work + 255) / 256, 256, 0, stream>>>(
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        gamma_p, rm_p, rv_p,
        slice_mean_p, scale_p,
        NSLICE, OC, 1.0 / spatial, 1.0 / total, total,
        eps, momentum,
        use_batch ? 1 : 0, update_running ? 1 : 0);
}

// NCHW path (fallback)
torch::Tensor fused_bn_meansub(
    torch::Tensor y,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine)
{
    TORCH_CHECK(y.is_cuda() && y.dtype() == torch::kFloat32,
                "fused_bn_meansub expects CUDA float32 input");
    auto yc = y.contiguous();
    const int N  = yc.size(0);
    const int OC = yc.size(1);
    const long slice_elems_l = yc.numel() / ((long)N * (long)OC);
    const int slice_elems = (int)slice_elems_l;
    const int NSLICE = N * OC;

    auto out = torch::empty_like(yc);
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({OC}, yc.options());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(yc.device());
    auto dbuf = torch::zeros({(long)NSLICE + 2L * OC}, opts_d);
    auto fbuf = torch::empty({(long)NSLICE + OC}, yc.options());
    double* slice_sum_p  = dbuf.data_ptr<double>();
    double* chan_sum_p   = slice_sum_p + NSLICE;
    double* chan_sumsq_p = chan_sum_p + OC;
    float*  slice_mean_p = fbuf.data_ptr<float>();
    float*  scale_p      = slice_mean_p + NSLICE;

    auto stream = at::cuda::getCurrentCUDAStream();

    const int max_vec = (slice_elems + 3) / 4;
    int chunks = (max_vec + THREADS * 4 - 1) / (THREADS * 4);
    if (chunks < 1) chunks = 1;
    if (chunks > 65535) chunks = 65535;
    dim3 grid(chunks, NSLICE);

    bn_ms_stats_kernel<<<grid, THREADS, 0, stream>>>(
        yc.data_ptr<float>(),
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        OC, slice_elems);

    launch_finalize(slice_sum_p, chan_sum_p, chan_sumsq_p,
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean_p, scale_p,
        NSLICE, OC, slice_elems, N,
        (float)eps, (float)momentum, use_batch, update_running,
        stream.stream());

    bn_ms_apply_kernel<<<grid, THREADS, 0, stream>>>(
        yc.data_ptr<float>(), out.data_ptr<float>(),
        scale_p, slice_mean_p,
        OC, slice_elems);

    return out;
}

// NHWC (channels_last_3d) fast path; requires C == 32.
torch::Tensor fused_bn_meansub_cl(
    torch::Tensor y,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine)
{
    TORCH_CHECK(y.is_cuda() && y.dtype() == torch::kFloat32,
                "fused_bn_meansub_cl expects CUDA float32 input");
    TORCH_CHECK(y.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "fused_bn_meansub_cl expects channels_last_3d input");
    const int N  = y.size(0);
    const int OC = y.size(1);
    TORCH_CHECK(OC == 32, "fused_bn_meansub_cl requires 32 channels");
    const long S = y.numel() / ((long)N * (long)OC);
    const int NSLICE = N * OC;

    // final output: NCHW contiguous
    auto out = torch::empty({N, OC, y.size(2), y.size(3), y.size(4)}, y.options());
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({OC}, y.options());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(y.device());
    auto dbuf = torch::zeros({(long)NSLICE + 2L * OC}, opts_d);
    auto fbuf = torch::empty({(long)NSLICE + OC}, y.options());
    double* slice_sum_p  = dbuf.data_ptr<double>();
    double* chan_sum_p   = slice_sum_p + NSLICE;
    double* chan_sumsq_p = chan_sum_p + OC;
    float*  slice_mean_p = fbuf.data_ptr<float>();
    float*  scale_p      = slice_mean_p + NSLICE;

    auto stream = at::cuda::getCurrentCUDAStream();

    const int tiles = (int)((S + 31) >> 5);
    int gx = tiles < 64 ? tiles : 64;
    dim3 sgrid(gx, N);
    stats_nhwc_kernel<<<sgrid, THREADS, 0, stream>>>(
        y.data_ptr<float>(),
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        OC, S);

    launch_finalize(slice_sum_p, chan_sum_p, chan_sumsq_p,
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean_p, scale_p,
        NSLICE, OC, (int)S, N,
        (float)eps, (float)momentum, use_batch, update_running,
        stream.stream());

    dim3 agrid(tiles, N);
    apply_transpose_nhwc_kernel<<<agrid, THREADS, 0, stream>>>(
        y.data_ptr<float>(), out.data_ptr<float>(),
        scale_p, slice_mean_p,
        OC, S);

    return out;
}
"""

cpp_source = """
torch::Tensor fused_bn_meansub(
    torch::Tensor y, torch::Tensor gamma,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine);
torch::Tensor fused_bn_meansub_cl(
    torch::Tensor y, torch::Tensor gamma,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine);
"""

fused_op = load_inline(
    name="fused_bn_meansub_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_bn_meansub", "fused_bn_meansub_cl"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cudnn, channels_last_3d, no bias kernel needed) followed
    by a fused custom-CUDA BatchNorm3d + spatial-mean-subtraction tail.

    Exact algebra used:
      BN(y)[n,c,s] - mean_s BN(y) == scale[c] * (g[n,c,s] - mean_s g),
    where y = g + bias[c], scale[c] = gamma[c] / sqrt(var_c + eps). The
    per-channel conv bias, BN beta and BN channel mean all cancel exactly,
    so the conv is run with bias=None and the NHWC conv output is consumed
    directly by the fused kernels (stats + transposing apply).
    """

    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):
        bn = self.batch_norm
        ct = self.conv_transpose

        if (x.is_cuda and x.dtype == torch.float32
                and bn.momentum is not None and ct.groups == 1):
            track = bn.track_running_stats
            use_batch = bool(bn.training) or (not track)
            update_running = bool(bn.training) and bool(track)

            if track:
                rm = bn.running_mean
                rv = bn.running_var
            else:
                rm = torch.zeros(bn.num_features, device=x.device, dtype=torch.float32)
                rv = torch.ones(bn.num_features, device=x.device, dtype=torch.float32)

            if update_running and bn.num_batches_tracked is not None:
                bn.num_batches_tracked += 1

            gamma = bn.weight if bn.affine else torch.empty(0, device=x.device)
            eps = float(bn.eps)
            mom = float(bn.momentum)

            # Bias cancels exactly in BN->mean-sub, so run the conv without it.
            x_cl = x.contiguous(memory_format=torch.channels_last_3d)
            w_cl = ct.weight.contiguous(memory_format=torch.channels_last_3d)
            y = F.conv_transpose3d(
                x_cl, w_cl, None, ct.stride, ct.padding,
                ct.output_padding, ct.groups, ct.dilation)

            if (y.is_contiguous(memory_format=torch.channels_last_3d)
                    and ct.out_channels == 32):
                return fused_op.fused_bn_meansub_cl(
                    y, gamma, rm, rv, eps, mom,
                    use_batch, update_running, bool(bn.affine))

            # layout fallback: NCHW fused kernels
            return fused_op.fused_bn_meansub(
                y.contiguous(), gamma, rm, rv, eps, mom,
                use_batch, update_running, bool(bn.affine))

        # pure-PyTorch fallback (keeps full bias/BN semantics)
        y = self.conv_transpose(x)
        z = self.batch_norm(y)
        return z - torch.mean(z, dim=(2, 3, 4), keepdim=True)
```