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


# ============================================================================
# Custom CUDA kernels
#
# Optimizations vs. the previous version:
#  1. Training mode no longer runs a separate full-volume statistics kernel.
#     A single fused kernel computes the raw ConvTranspose3d output for the
#     whole 63^3 volume (organized as 8x4x4 voxel tiles), writes the 4x4x4
#     block sums (pooled pre-BN values) for the interior 60^3 region, and
#     accumulates per-channel sum / sum-of-squares for BatchNorm on the fly.
#  2. Threads are assigned to voxels sorted by parity class so that warps are
#     (nearly) uniform in the number of kernel taps they execute.  This
#     removes the divergent-loop serialization of the naive mapping and
#     roughly doubles the conv throughput.
#  3. BatchNorm scale/shift are finalized in a tiny dedicated kernel that
#     also updates the running statistics, and BN+pool scaling is applied
#     with a single addcmul, so training mode adds only ~50us of overhead.
#  4. Inference mode keeps everything in one kernel: ConvTranspose3d +
#     BatchNorm(running stats) + AvgPool(4,4,4), never materializing the
#     63^3 intermediate in DRAM.
# ============================================================================

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

#define CI 3
#define CO 16
#define IDIM 32
#define ODIM 15
#define NTHREADS 128
#define PITCH (CO + 1)
#define SXD 45   // 5*3*3 input patch per channel for an 8x4x4 output tile

// Parity classes ordered by tap count so warps are near-uniform:
// 0:eee(1 tap) 1:eeo(2) 2:eoe(2) 3:oee(2) 4:eoo(4) 5:oeo(4) 6:ooe(4) 7:ooo(8)
__constant__ int c_pd[8] = {0, 0, 0, 1, 0, 1, 1, 1};
__constant__ int c_ph[8] = {0, 0, 1, 0, 1, 0, 1, 1};
__constant__ int c_pw[8] = {0, 1, 0, 0, 1, 1, 0, 1};

// For local coordinate l (parity dependent):
//   global D = base + l ; input d = (D + 1 - kd) / 2
//   even l: 1 tap  -> a = l/2   , kd = 1
//   odd  l: 2 taps -> a = l/2   , kd = 2
//                     a = l/2+1 , kd = 0
__device__ __forceinline__ int setup_taps(int l, int* a, int* k) {
    if ((l & 1) == 0) { a[0] = l >> 1; k[0] = 1; return 1; }
    a[0] = l >> 1; k[0] = 2; a[1] = (l >> 1) + 1; k[1] = 0;
    return 2;
}

__device__ __forceinline__ void conv_acc(
    const float* s_x, const float* s_w,
    const int* ad, const int* ah, const int* aw,
    const int* kd, const int* kh, const int* kw,
    int nd, int nh, int nw, float* acc)
{
    for (int i = 0; i < nd; ++i) {
        for (int j = 0; j < nh; ++j) {
            for (int m = 0; m < nw; ++m) {
                const int xs = ad[i] * 9 + ah[j] * 3 + aw[m];
                const float x0 = s_x[xs];
                const float x1 = s_x[xs + SXD];
                const float x2 = s_x[xs + 2 * SXD];
                const int woff = kd[i] * 9 + kh[j] * 3 + kw[m];
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) {
                    const float* wp = s_w + oc * 27 + woff;
                    acc[oc] += x0 * wp[0] + x1 * wp[CO * 27] + x2 * wp[2 * CO * 27];
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Training kernel: raw ConvTranspose3d over the full 63^3 volume, fused
// 4x4x4 block pooling (raw block sums) + BatchNorm statistics accumulation.
// Tiles: 8x4x4 voxels, grid per image = 8*16*16 tiles (covers 64^3 >= 63^3).
// ---------------------------------------------------------------------------
__global__ void fused_train_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    double* __restrict__ gsum,
    double* __restrict__ gsq,
    int total_tiles)
{
    __shared__ float s_x[CI * SXD];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_bsum[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) s_bsum[tid] = 0.0f;

    const int cls = tid >> 4, r = tid & 15;
    const int pd = c_pd[cls], ph = c_ph[cls], pw = c_pw[cls];
    const int hd = r >> 2, hh = (r >> 1) & 1, hw = r & 1;
    const int vd = 2 * hd + pd, vh = 2 * hh + ph, vw = 2 * hw + pw;
    const int vid = vd * 16 + vh * 4 + vw;

    int ad[2], ah[2], aw[2], kd[2], kh[2], kw[2];
    const int nd = setup_taps(vd, ad, kd);
    const int nh = setup_taps(vh, ah, kh);
    const int nw = setup_taps(vw, aw, kw);

    float tsq[CO];
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) tsq[oc] = 0.0f;
    __syncthreads();

    const int TPI = 8 * 16 * 16;
    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int td  = rem >> 8;
        rem          &= 255;
        const int th  = rem >> 4;
        const int tw  = rem & 15;
        const bool interior = (td < 7) && (th < 15) && (tw < 15);

        const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        if (interior) {
            #pragma unroll
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                s_x[i] = xb[((ic * IDIM + 4 * td + a) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
            }
        } else {
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                const int gi = 4 * td + a, gj = 2 * th + b, gk = 2 * tw + c;
                float v = 0.0f;
                if (gi < IDIM && gj < IDIM && gk < IDIM)
                    v = xb[((ic * IDIM + gi) * IDIM + gj) * IDIM + gk];
                s_x[i] = v;
            }
        }
        __syncthreads();

        float acc[CO];
        const bool valid = interior ||
            ((8 * td + vd) < 63 && (4 * th + vh) < 63 && (4 * tw + vw) < 63);
        if (valid) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = bias[oc];
            conv_acc(s_x, s_w, ad, ah, aw, kd, kh, kw, nd, nh, nw, acc);
        } else {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;
        }

        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) {
            s_red[vid][oc] = acc[oc];
            tsq[oc] += acc[oc] * acc[oc];
        }
        __syncthreads();

        if (tid < 32) {
            const int oc = tid & 15, half = tid >> 4;
            float s = 0.0f;
            #pragma unroll 8
            for (int i = 0; i < 64; ++i) s += s_red[half * 64 + i][oc];
            atomicAdd(&s_bsum[oc], s);
            const int bd4 = 2 * td + half;
            if (bd4 < ODIM && th < ODIM && tw < ODIM) {
                const long oidx = (((long)n * CO + oc) * ODIM + bd4) * (ODIM * ODIM)
                                + (long)th * ODIM + tw;
                out[oidx] = s;
            }
        }
    }

    __syncthreads();
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) s_red[tid][oc] = tsq[oc];
    __syncthreads();
    if (tid < CO) {
        float q = 0.0f;
        #pragma unroll 8
        for (int i = 0; i < NTHREADS; ++i) q += s_red[i][tid];
        atomicAdd(gsq + tid, (double)q);
        atomicAdd(gsum + tid, (double)s_bsum[tid]);
    }
}

// ---------------------------------------------------------------------------
// Inference kernel: ConvTranspose3d + BatchNorm(scale/shift) + AvgPool(4)
// fully fused.  Only voxels in the 60^3 region that pooling reads are
// computed.  Tiles per image: 8*15*15.
// ---------------------------------------------------------------------------
__global__ void fused_eval_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    const float* __restrict__ shift,
    float* __restrict__ out,
    int total_tiles)
{
    __shared__ float s_x[CI * SXD];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_sc[CO];
    __shared__ float s_sh[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) { s_sc[tid] = scale[tid]; s_sh[tid] = shift[tid]; }

    const int cls = tid >> 4, r = tid & 15;
    const int pd = c_pd[cls], ph = c_ph[cls], pw = c_pw[cls];
    const int hd = r >> 2, hh = (r >> 1) & 1, hw = r & 1;
    const int vd = 2 * hd + pd, vh = 2 * hh + ph, vw = 2 * hw + pw;
    const int vid = vd * 16 + vh * 4 + vw;

    int ad[2], ah[2], aw[2], kd[2], kh[2], kw[2];
    const int nd = setup_taps(vd, ad, kd);
    const int nh = setup_taps(vh, ah, kh);
    const int nw = setup_taps(vw, aw, kw);
    __syncthreads();

    const int TPI = 8 * ODIM * ODIM;
    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int td  = rem / (ODIM * ODIM);
        rem          %= ODIM * ODIM;
        const int th  = rem / ODIM;
        const int tw  = rem % ODIM;
        const bool gd = (td == 7);

        const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        if (!gd) {
            #pragma unroll
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                s_x[i] = xb[((ic * IDIM + 4 * td + a) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
            }
        } else {
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                const int gi = 4 * td + a;
                float v = 0.0f;
                if (gi < IDIM)
                    v = xb[((ic * IDIM + gi) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
                s_x[i] = v;
            }
        }
        __syncthreads();

        float acc[CO];
        if (!gd || vd < 4) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = bias[oc];
            conv_acc(s_x, s_w, ad, ah, aw, kd, kh, kw, nd, nh, nw, acc);
        } else {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;
        }

        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) s_red[vid][oc] = acc[oc];
        __syncthreads();

        if (tid < 32) {
            const int oc = tid & 15, half = tid >> 4;
            const int bd4 = 2 * td + half;
            if (bd4 < ODIM) {
                float s = 0.0f;
                #pragma unroll 8
                for (int i = 0; i < 64; ++i) s += s_red[half * 64 + i][oc];
                const long oidx = (((long)n * CO + oc) * ODIM + bd4) * (ODIM * ODIM)
                                + (long)th * ODIM + tw;
                out[oidx] = s * s_sc[oc] + s_sh[oc];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// BatchNorm finalization: computes mean/var from the accumulated moments,
// updates running statistics, and produces the fused scale/shift that also
// folds the 1/64 average-pool factor.
// ---------------------------------------------------------------------------
__global__ void bn_finalize_kernel(
    const double* __restrict__ gsum,
    const double* __restrict__ gsq,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    long long* __restrict__ nbt,
    float* __restrict__ scale,
    float* __restrict__ shift,
    double m, double momentum, double eps, int track)
{
    const int c = threadIdx.x;
    if (c >= CO) return;
    const double mean = gsum[c] / m;
    double var = gsq[c] / m - mean * mean;
    if (var < 0.0) var = 0.0;
    const double invstd = rsqrt(var + eps);
    if (track) {
        const double var_unb = var * (m / (m - 1.0));
        running_mean[c] = (float)((double)running_mean[c] * (1.0 - momentum) + momentum * mean);
        running_var[c]  = (float)((double)running_var[c] * (1.0 - momentum) + momentum * var_unb);
        if (c == 0) nbt[0] += 1;
    }
    const double g = (double)gamma[c];
    scale[c] = (float)(g * invstd * (1.0 / 64.0));
    shift[c] = (float)((double)beta[c] - mean * g * invstd);
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                          torch::Tensor gsum, torch::Tensor gsq) {
    const int N = x.size(0);
    auto out = torch::empty({N, CO, ODIM, ODIM, ODIM}, x.options());
    const int total = N * 8 * 16 * 16;
    int blocks = total < 4096 ? total : 4096;
    fused_train_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(),
        gsum.data_ptr<double>(), gsq.data_ptr<double>(), total);
    return out;
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                         torch::Tensor scale, torch::Tensor shift) {
    const int N = x.size(0);
    auto out = torch::empty({N, CO, ODIM, ODIM, ODIM}, x.options());
    const int total = N * 8 * ODIM * ODIM;
    int blocks = total < 4096 ? total : 4096;
    fused_eval_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>(), total);
    return out;
}

void bn_finalize(torch::Tensor gsum, torch::Tensor gsq, torch::Tensor gamma,
                 torch::Tensor beta, torch::Tensor running_mean,
                 torch::Tensor running_var, torch::Tensor nbt,
                 torch::Tensor scale, torch::Tensor shift,
                 double m, double momentum, double eps, int64_t track) {
    bn_finalize_kernel<<<1, CO>>>(
        gsum.data_ptr<double>(), gsq.data_ptr<double>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        (long long*)nbt.data_ptr<int64_t>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        m, momentum, eps, (int)track);
}
"""

cpp_source = (
    "torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor gsum, torch::Tensor gsq);\n"
    "torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor scale, torch::Tensor shift);\n"
    "void bn_finalize(torch::Tensor gsum, torch::Tensor gsq, torch::Tensor gamma, "
    "torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, "
    "torch::Tensor nbt, torch::Tensor scale, torch::Tensor shift, "
    "double m, double momentum, double eps, int64_t track);"
)

_ext = load_inline(
    name="fused_convT_bn_pool_v2",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_train", "fused_eval", "bn_finalize"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: ConvTranspose3d + BatchNorm3d + two AvgPool3d
    layers fused into custom CUDA kernels.  A single kernel handles the whole
    forward pass in inference mode; training mode uses one fused kernel for
    conv+pool+statistics, a tiny finalization kernel for the BatchNorm
    moments, and one elementwise scale.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
        self.avg_pool2 = nn.AvgPool3d(kernel_size=2)
        self.ext = _ext

    def _fallback(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.avg_pool1(x)
        x = self.avg_pool2(x)
        return x

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

        supported = (
            x.is_cuda and x.dtype == torch.float32 and
            ct.in_channels == 3 and ct.out_channels == 16 and
            ct.kernel_size == (3, 3, 3) and ct.stride == (2, 2, 2) and
            ct.padding == (1, 1, 1) and ct.output_padding == (0, 0, 0) and
            ct.dilation == (1, 1, 1) and ct.groups == 1 and ct.bias is not None and
            tuple(x.shape[-3:]) == (32, 32, 32) and
            bn.affine and bn.momentum is not None and
            (self.training or bn.track_running_stats)
        )
        if not supported:
            return self._fallback(x)

        x = x.contiguous()
        w = ct.weight.contiguous()
        b = ct.bias.contiguous()
        gamma = bn.weight.contiguous()
        beta = bn.bias.contiguous()
        dev = x.device

        if self.training:
            n = x.shape[0]
            m_count = float(n) * 250047.0  # n * 63^3, exactly representable
            gsum = torch.zeros(16, dtype=torch.float64, device=dev)
            gsq = torch.zeros(16, dtype=torch.float64, device=dev)
            raw = self.ext.fused_train(x, w, b, gsum, gsq)

            scale = torch.empty(16, dtype=torch.float32, device=dev)
            shift = torch.empty(16, dtype=torch.float32, device=dev)
            track = bn.track_running_stats
            rm = bn.running_mean if track else gamma
            rv = bn.running_var if track else gamma
            nbt = bn.num_batches_tracked if track else bn.num_batches_tracked
            self.ext.bn_finalize(gsum, gsq, gamma, beta, rm, rv, nbt,
                                 scale, shift, m_count, float(bn.momentum),
                                 float(bn.eps), 1 if track else 0)

            return torch.addcmul(shift.view(1, 16, 1, 1, 1), raw,
                                 scale.view(1, 16, 1, 1, 1))
        else:
            scale = gamma * torch.rsqrt(bn.running_var + bn.eps) * (1.0 / 64.0)
            shift = beta - bn.running_mean * (gamma * torch.rsqrt(bn.running_var + bn.eps))
            return self.ext.fused_eval(x, w, b, scale.contiguous(), shift.contiguous())


batch_size = 64
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (out_channels, 1, 1, 1)


def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, bias_shape]
```