```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) stats_nhwc_kernel: the double-precision shared-memory atomicAdd CAS
#     loop is replaced by an in-warp shuffle reduction + float shared
#     staging, with only 3 global double atomics per channel per block.
#  2) apply_transpose_nhwc_kernel: rewritten as a 64(s) x 32(c) shared-mem
#     tile where BOTH the global loads (NHWC float4) and the global stores
#     (NCHW, 32 consecutive spatial positions per warp) are fully coalesced
#     128B/warp transactions. The old version issued strided 4B stores
#     (~25% store efficiency) which made it 4.5x slower than the traffic
#     bound. Scalar stores are used because rows of the NCHW output are not
#     16B-aligned when S is not a multiple of 4 (S = 31^3 here).
#  3) Workspace zero-init uses cudaMemsetAsync instead of a fill kernel.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256

// ---------------------------------------------------------------------------
// NCHW fallback kernels (layout-generic path)
// ---------------------------------------------------------------------------
__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] stored as
// [N, S=D*H*W, C] with C == 32 (guarantees 16B alignment of every float4).
// ---------------------------------------------------------------------------

// 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.
// Reduction: in-warp shfl over the 4 spatial lanes sharing a channel group,
// float shared staging across 8 warps, 3 global double atomics per channel.
__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);
        }
    }

    // Reduce over the 4 spatial lanes that share a channel group (lanes
    // l, l+8, l+16, l+24 have identical channel group).
    #pragma unroll
    for (int off = 16; off >= 8; off >>= 1) {
        s0 += __shfl_down_sync(0xffffffffu, s0, off);
        s1 += __shfl_down_sync(0xffffffffu, s1, off);
        s2 += __shfl_down_sync(0xffffffffu, s2, off);
        s3 += __shfl_down_sync(0xffffffffu, s3, off);
        q0 += __shfl_down_sync(0xffffffffu, q0, off);
        q1 += __shfl_down_sync(0xffffffffu, q1, off);
        q2 += __shfl_down_sync(0xffffffffu, q2, off);
        q3 += __shfl_down_sync(0xffffffffu, q3, off);
    }

    __shared__ float sh_sum[8][32];
    __shared__ float sh_sq[8][32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane < 8) {
        sh_sum[warp][c0 + 0] = s0; sh_sum[warp][c0 + 1] = s1;
        sh_sum[warp][c0 + 2] = s2; sh_sum[warp][c0 + 3] = s3;
        sh_sq[warp][c0 + 0]  = q0; sh_sq[warp][c0 + 1]  = q1;
        sh_sq[warp][c0 + 2]  = q2; sh_sq[warp][c0 + 3]  = q3;
    }
    __syncthreads();

    if (threadIdx.x < 32) {
        const int c = threadIdx.x;
        float fs = 0.f, fq = 0.f;
        #pragma unroll
        for (int w = 0; w < 8; ++w) {
            fs += sh_sum[w][c];
            fq += sh_sq[w][c];
        }
        const double ds = (double)fs;
        atomicAdd(&slice_sum[n * C + c], ds);
        atomicAdd(&chan_sum[c],          ds);
        atomicAdd(&chan_sumsq[c],        (double)fq);
    }
}

// Pass 2 (NHWC -> NCHW transpose with fused affine):
//   out[n, c, s] = scale[c] * (y[n, s, c] - slice_mean[n, c])
// 64(s) x 32(c) shared-memory tile. Loads: coalesced float4 over channels.
// Stores: each warp writes 32 consecutive spatial positions of one channel
// (128B/warp, fully coalesced, no alignment requirement). All shared-mem
// accesses are bank-conflict free.
__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][65];
    const int n = blockIdx.y;
    const long s0 = (long)blockIdx.x << 6;

    const int cg = threadIdx.x & 7;        // float4 channel group 0..7
    const int s_sub = threadIdx.x >> 3;    // 0..31
    const float4 sc4 = *reinterpret_cast<const float4*>(scale + (cg << 2));
    const float4 mn4 = *reinterpret_cast<const float4*>(slice_mean + n * C + (cg << 2));
    const long nbase = (long)n * S;

    #pragma unroll
    for (int half = 0; half < 2; ++half) {
        const int ss = s_sub + (half << 5);   // 0..63
        const long s = s0 + ss;
        if (s < S) {
            const float4 v = __ldcs(reinterpret_cast<const float4*>(
                y + (nbase + s) * C + (cg << 2)));
            tile[(cg << 2) + 0][ss] = sc4.x * (v.x - mn4.x);
            tile[(cg << 2) + 1][ss] = sc4.y * (v.y - mn4.y);
            tile[(cg << 2) + 2][ss] = sc4.z * (v.z - mn4.z);
            tile[(cg << 2) + 3][ss] = sc4.w * (v.w - mn4.w);
        }
    }
    __syncthreads();

    const long rem = S - s0;
    const int s = threadIdx.x & 63;
    if (s < rem) {
        int c = threadIdx.x >> 6;   // 0..3, strides of 4 over 32 channels
        float* obase = out + ((long)n * C) * S + s0 + s;
        #pragma unroll
        for (int i = 0; i < 8; ++i, c += 4) {
            __stcs(obase + (long)c * S, tile[c][s]);
        }
    }
}

// ---------------------------------------------------------------------------
// 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::empty({(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();
    cudaMemsetAsync(dbuf.data_ptr(), 0, dbuf.numel() * sizeof(double), stream.stream());

    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::empty({(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();
    cudaMemsetAsync(dbuf.data_ptr(), 0, dbuf.numel() * sizeof(double), stream.stream());

    const int tiles32 = (int)((S + 31) >> 5);
    int gx = tiles32 < 32 ? tiles32 : 32;
    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((unsigned)((S + 63) >> 6), 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_v4",
    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)
```