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


# ============================================================================
# Fused ConvTranspose3d + BatchNorm3d + AvgPool3d(2) + AvgPool3d(2)
#
# Key algorithmic insight (eval path):
#   The two AvgPool3d(2) layers together compute 4x4x4 block averages of the
#   63^3 transposed-conv output, using only voxels [0,60)^3.  In 1D, pooling
#   and the transposed convolution commute, giving (for pooled output q):
#       out4[q] = w0*(x[2q+1]+x[2q+2]) + (w1+w2)*(x[2q]+x[2q+1])
#   i.e. each pooled output voxel is a 3x3x3 tap stencil over 2x2x2 "box
#   sums" of the input, instead of 64 individual conv voxels.  This removes
#   ~8x of the arithmetic, requires zero boundary handling (input indices
#   stay in [0,30]), and the BatchNorm affine (being linear) is folded into
#   a single scale/shift applied to the pooled result.
#
#   Each thread produces one (16-channel) pooled output voxel.  A block
#   covers an 8x8x4 output tile (256 threads); the needed 17x17x9x3 input
#   slab is loaded into shared memory cooperatively with fully coalesced
#   32-float rows.
#
# Training path:
#   BatchNorm needs per-voxel statistics over the full 63^3 volume, so a
#   direct fused conv+pool+statistics kernel is used (no separate stats
#   pass), followed by a tiny finalization kernel that also applies the
#   scale/shift.
# ============================================================================

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

#define CI 3
#define CO 16
#define IDIM 32
#define ODIM 15
#define NT 256

// Eval geometry: 8x8x4 output tile -> 17x17x9 input slab per image.
#define SLD 17
#define SLH 17
#define SLW 9
#define SLICE (SLD * SLH * SLW)   // 2601
#define PADSL 16                  // padded slice stride -> conflict-free
#define SLPAD (SLICE + PADSL)     // 2617

// Training geometry: 16x4x4 conv-output tile -> 9x3x3 input patch per ic.
#define SXD 81
#define PITCH (CO + 1)

// ---------------------------------------------------------------------------
// Eval kernel: fully fused, one thread per pooled output voxel.
// ---------------------------------------------------------------------------
__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)
{
    __shared__ float s_x[CI * SLPAD];
    __shared__ float s_w[CI * 27 * CO];

    const int tid = threadIdx.x;
    // grid: (tiles_per_image = 2*2*4 = 16, batch)
    const int qd0 = (blockIdx.x & 1) << 3;
    const int qh0 = (blockIdx.x & 2) << 2;
    const int qw0 = (blockIdx.x >> 2) << 2;
    const int n   = blockIdx.y;

    const float* __restrict__ xb = x + (long)n * (CI * IDIM * IDIM * IDIM);
    const float* __restrict__ xr = xb + ((long)2 * qd0 * IDIM + 2 * qh0) * IDIM + 2 * qw0;

    // Cooperative load: 3ic x 17d x 17h rows of 9 contiguous floats.
    for (int i = tid; i < CI * SLD * SLH; i += NT) {
        const int c  = i / (SLD * SLH);
        const int r  = i - c * (SLD * SLH);
        const int a  = r / SLH;
        const int b  = r - a * SLH;
        const float* src = xr + ((long)c * IDIM + a) * IDIM + b;
        float* dst = s_x + c * SLPAD + r * SLW;
        #pragma unroll
        for (int m = 0; m < SLW; ++m) dst[m] = src[m];
    }
    for (int i = tid; i < CI * 27 * CO; i += NT) s_w[i] = w[i];
    __syncthreads();

    const int qd = qd0 + (tid >> 5);
    const int qh = qh0 + ((tid >> 2) & 7);
    const int qw = qw0 + (tid & 3);
    const float* __restrict__ xs =
        s_x + ((tid >> 5) * SLH + ((tid >> 2) & 7)) * SLW + (tid & 3) * 2;

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

    const float* __restrict__ wp = s_w;
    #pragma unroll
    for (int ic = 0; ic < CI; ++ic, xs += SLPAD) {
        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            const int os = (kd == 0) ? SLW : 0;
            float A[3][3];
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                #pragma unroll
                for (int m = 0; m < 3; ++m)
                    A[j][m] = xs[os + j * SLW + m] + xs[os + j * SLW + m + SLW];
            }
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int oh = (kh == 0) ? 1 : 0;
                float B0 = A[oh][0] + A[oh + 1][0];
                float B1 = A[oh][1] + A[oh + 1][1];
                float B2 = A[oh][2] + A[oh + 1][2];
                const float X0 = B1 + B2;   // kw = 0
                const float X1 = B0 + B1;   // kw = 1, 2
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) acc[oc] += X0 * wp[oc];
                wp += CO;
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) acc[oc] += X1 * wp[oc];
                wp += CO;
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) acc[oc] += X1 * wp[oc];
                wp += CO;
            }
        }
    }

    const long oidx = (((long)n * CO) * ODIM + qd) * (ODIM * ODIM)
                    + (long)qh * ODIM + qw;
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc)
        out[oidx + (long)oc * (ODIM * ODIM * ODIM)] =
            (acc[oc] + bias[oc]) * scale[oc] + shift[oc];
}

// ---------------------------------------------------------------------------
// Training kernel: raw ConvTranspose3d over the full 63^3 volume with fused
// 4x4x4 block pooling (raw block sums) and BatchNorm moment accumulation.
// Tiles: 16x4x4 voxels (256 threads, one parity class per warp).
// ---------------------------------------------------------------------------
__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[NT][PITCH];
    __shared__ float s_bsum[CO];

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

    const int cls = tid >> 5, lane = tid & 31;
    const int pd = cls & 1, ph = (cls >> 1) & 1, pw = cls >> 2;
    const int hd = lane >> 2, hh = (lane >> 1) & 1, hw = lane & 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];
    if (pd) { ad[0] = hd; kd[0] = 2; ad[1] = hd + 1; kd[1] = 0; }
    else    { ad[0] = hd; kd[0] = 1; }
    if (ph) { ah[0] = hh; kh[0] = 2; ah[1] = hh + 1; kh[1] = 0; }
    else    { ah[0] = hh; kh[0] = 1; }
    if (pw) { aw[0] = hw; kw[0] = 2; aw[1] = hw + 1; kw[1] = 0; }
    else    { aw[0] = hw; kw[0] = 1; }
    const int nd = pd + 1, nh = ph + 1, nw = pw + 1;

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

    const int TPI = 4 * 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 float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        for (int i = tid; i < CI * SXD; i += NT) {
            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 = 8 * 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 =
            (16 * 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];
            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* wpp = s_w + oc * 27 + woff;
                            acc[oc] += x0 * wpp[0] + x1 * wpp[CO * 27]
                                     + x2 * wpp[2 * CO * 27];
                        }
                    }
                }
            }
        } 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 < 64) {
            const int oc = tid & 15, qrt = tid >> 4;
            float s = 0.0f;
            #pragma unroll 8
            for (int i = 0; i < 64; ++i) s += s_red[qrt * 64 + i][oc];
            atomicAdd(&s_bsum[oc], s);
            const int bd4 = 4 * td + qrt;
            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 < NT; ++i) q += s_red[i][tid];
        atomicAdd(gsq + tid, (double)q);
        atomicAdd(gsum + tid, (double)s_bsum[tid]);
    }
}

// ---------------------------------------------------------------------------
// BatchNorm finalization + fused scale/shift application (training path).
// Block 0 finalizes the moments, updates running stats and writes
// scale/shift; all blocks then apply out = raw * scale + shift.
// ---------------------------------------------------------------------------
__global__ void bn_finalize_apply_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,
    float* __restrict__ raw,
    double m, double momentum, double eps, int track, long total)
{
    __shared__ float s_sc[CO];
    __shared__ float s_sh[CO];
    if (threadIdx.x < CO) {
        const int c = threadIdx.x;
        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 (blockIdx.x == 0) {
            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);
        }
        const double g = (double)gamma[c];
        s_sc[c] = (float)(g * invstd * (1.0 / 64.0));
        s_sh[c] = (float)((double)beta[c] - mean * g * invstd);
    }
    __syncthreads();

    const long stride = (long)gridDim.x * blockDim.x;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += stride) {
        const int c = (int)(idx / 3375) & 15;
        raw[idx] = raw[idx] * s_sc[c] + s_sh[c];
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
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());
    dim3 grid(16, N);
    fused_eval_kernel<<<grid, NT>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>());
    return out;
}

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 * 4 * 16 * 16;
    int blocks = total < 8192 ? total : 8192;
    fused_train_kernel<<<blocks, NT>>>(
        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;
}

void bn_finalize_apply(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,
                       torch::Tensor raw,
                       double m, double momentum, double eps, int64_t track) {
    const long total = raw.numel();
    int blocks = (int)((total + 255) / 256);
    if (blocks > 1024) blocks = 1024;
    bn_finalize_apply_kernel<<<blocks, 256>>>(
        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>(),
        raw.data_ptr<float>(),
        m, momentum, eps, (int)track, total);
}
"""

cpp_source = (
    "torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor scale, torch::Tensor shift);\n"
    "torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor gsum, torch::Tensor gsq);\n"
    "void bn_finalize_apply(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, torch::Tensor raw, "
    "double m, double momentum, double eps, int64_t track);"
)

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


class ModelNew(nn.Module):
    """
    Optimized version of Model: ConvTranspose3d + BatchNorm3d + two AvgPool3d
    layers fused into custom CUDA kernels.  In inference mode the whole
    network collapses into a single kernel that exploits the algebraic
    commutation of the transposed convolution with the two pooling layers,
    skipping the 63^3 intermediate entirely.  Training mode uses a fused
    conv+pool+statistics kernel and a fused BatchNorm-finalize/apply kernel.
    """

    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
            self.ext.bn_finalize_apply(gsum, gsq, gamma, beta, rm, rv, nbt,
                                       scale, shift, raw, m_count,
                                       float(bn.momentum), float(bn.eps),
                                       1 if track else 0)
            return raw
        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]
```