```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.
#
# The architecture is: ConvTranspose3d(k=3,s=2,p=1) -> BatchNorm3d ->
# AvgPool3d(2) -> AvgPool3d(2).
#
# Key observations:
#   * Two consecutive non-overlapping 2x2x2 average pools are exactly one
#     4x4x4 average pool (the trailing elements dropped by floor() match).
#   * BatchNorm is affine per channel, so BN followed by average pooling is
#     just: out = scale * (block average of conv output) + shift.
#   * Therefore the whole network collapses to: for every 4x4x4 block of the
#     transposed-conv output, compute its sum directly from the input.
#     Each pooled voxel is a 3x3x3xIC stencil over the *input* with
#     precomputed coefficients obtained by contracting the conv weights with
#     a small multiplicity matrix M.  The huge 63^3 intermediate is never
#     materialized.
#   * In training mode BatchNorm needs the batch mean/var of the (never
#     materialized) conv output, so a second kernel computes per-channel
#     sum and sum-of-squares of every transposed-conv output voxel on the fly
#     with double precision accumulators.
# ---------------------------------------------------------------------------

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

// --------------------------------------------------------------------------
// Kernel A: per-channel sum and sum-of-squares of the full transposed-conv
// output y (never materialized).  One block per (n, oc) pair.
// y[o] = sum_{ic,k} x[ic, i] * w[ic, oc, k] with o = 2*i - 1 + k
// => for a given o, k must satisfy k == (o+1) mod 2, i = (o+1-k)/2.
// --------------------------------------------------------------------------
__global__ void tconv_stats_kernel(
        const float* __restrict__ x,      // (N, IC, ID, ID, ID)
        const float* __restrict__ w,      // (IC, OC, 3, 3, 3)
        double* __restrict__ g_sum,       // (OC)
        double* __restrict__ g_sumsq,     // (OC)
        int N, int IC, int OC, int ID, int OD) {

    const int oc = blockIdx.x % OC;
    const int n  = blockIdx.x / OC;

    extern __shared__ float ws[];         // IC * 27 weights for this oc
    for (int t = threadIdx.x; t < IC * 27; t += blockDim.x) {
        int ic = t / 27;
        int k  = t % 27;
        ws[t] = w[((long)(ic * OC + oc)) * 27 + k];
    }
    __syncthreads();

    const long xbase = (long)n * IC * ID * ID * ID;
    const int total  = OD * OD * OD;

    double s  = 0.0;
    double s2 = 0.0;

    for (int v = threadIdx.x; v < total; v += blockDim.x) {
        const int d = v / (OD * OD);
        const int r = v - d * OD * OD;
        const int h = r / OD;
        const int ww = r - h * OD;

        const int kpd = (d + 1) & 1;
        const int kph = (h + 1) & 1;
        const int kpw = (ww + 1) & 1;
        const int nd = kpd ? 1 : 2;
        const int nh = kph ? 1 : 2;
        const int nw = kpw ? 1 : 2;

        float y = 0.0f;
        for (int td = 0; td < nd; ++td) {
            const int kd = kpd + 2 * td;
            const int id = (d + 1 - kd) >> 1;
            for (int th = 0; th < nh; ++th) {
                const int kh = kph + 2 * th;
                const int ih = (h + 1 - kh) >> 1;
                for (int tw = 0; tw < nw; ++tw) {
                    const int kw = kpw + 2 * tw;
                    const int iw = (ww + 1 - kw) >> 1;
                    const int woff = kd * 9 + kh * 3 + kw;
                    const long xoff = ((long)id * ID + ih) * ID + iw;
                    #pragma unroll
                    for (int ic = 0; ic < 3; ++ic) {
                        y += x[xbase + ((long)ic * ID * ID * ID) + xoff]
                           * ws[ic * 27 + woff];
                    }
                }
            }
        }
        s  += (double)y;
        s2 += (double)y * (double)y;
    }

    // Block reduction (double precision).
    __shared__ double rs[32];
    __shared__ double rs2[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s  += __shfl_down_sync(0xffffffffu, s,  off);
        s2 += __shfl_down_sync(0xffffffffu, s2, off);
    }
    if (lane == 0) { rs[wid] = s; rs2[wid] = s2; }
    __syncthreads();
    if (wid == 0) {
        const int nwarp = (blockDim.x + 31) >> 5;
        if (lane < nwarp) { s = rs[lane]; s2 = rs2[lane]; }
        else { s = 0.0; s2 = 0.0; }
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s  += __shfl_down_sync(0xffffffffu, s,  off);
            s2 += __shfl_down_sync(0xffffffffu, s2, off);
        }
        if (lane == 0) {
            atomicAdd(&g_sum[oc],   s);
            atomicAdd(&g_sumsq[oc], s2);
        }
    }
}

// --------------------------------------------------------------------------
// Kernel B: fused ConvTranspose3d-block-sum + BatchNorm + both AvgPools.
// Each thread handles one pooled spatial voxel for one batch entry and loops
// over all output channels.  C holds the precomputed stencil coefficients
// with layout (OC, IC, 3, 3, 3) (flattened last 3 dims).
// --------------------------------------------------------------------------
template <int IC, int OC>
__global__ void fused_tconv_bn_pool_kernel(
        const float* __restrict__ x,      // (N, IC, ID, ID, ID)
        const float* __restrict__ C,      // (OC, IC, 27)
        const float* __restrict__ scale,  // (OC)
        const float* __restrict__ shift,  // (OC)
        float* __restrict__ out,          // (N, OC, BD, BD, BD)
        int N, int ID, int BD) {

    __shared__ float Cs[OC * IC * 27];
    __shared__ float scl[OC];
    __shared__ float shf[OC];

    for (int t = threadIdx.x; t < OC * IC * 27; t += blockDim.x) Cs[t] = C[t];
    if (threadIdx.x < OC) {
        scl[threadIdx.x] = scale[threadIdx.x];
        shf[threadIdx.x] = shift[threadIdx.x];
    }
    __syncthreads();

    const int spatial = BD * BD * BD;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * spatial) return;

    const int n  = idx / spatial;
    int r        = idx - n * spatial;
    const int bd = r / (BD * BD);
    r -= bd * BD * BD;
    const int bh = r / BD;
    const int bw = r - bh * BD;

    const int id0 = 2 * bd;
    const int ih0 = 2 * bh;
    const int iw0 = 2 * bw;

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

    const long xn = (long)n * IC * ID * ID * ID;

    #pragma unroll
    for (int ic = 0; ic < IC; ++ic) {
        #pragma unroll
        for (int a = 0; a < 3; ++a) {
            #pragma unroll
            for (int b = 0; b < 3; ++b) {
                #pragma unroll
                for (int c = 0; c < 3; ++c) {
                    const float xv = x[xn + (((long)ic * ID + (id0 + a)) * ID + (ih0 + b)) * ID + (iw0 + c)];
                    const int coff = ic * 27 + a * 9 + b * 3 + c;
                    #pragma unroll
                    for (int oc = 0; oc < OC; ++oc) {
                        acc[oc] += xv * Cs[oc * (IC * 27) + coff];
                    }
                }
            }
        }
    }

    const long obase = (long)n * OC * spatial + bd * BD * BD + bh * BD + bw;
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) {
        out[obase + (long)oc * spatial] = acc[oc] * scl[oc] + shf[oc];
    }
}

// --------------------------------------------------------------------------
// Host wrappers
// --------------------------------------------------------------------------
std::vector<torch::Tensor> tconv_stats(torch::Tensor x, torch::Tensor w,
                                       int64_t N, int64_t IC, int64_t OC,
                                       int64_t ID, int64_t OD) {
    auto opts = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto sum   = torch::zeros({OC}, opts);
    auto sumsq = torch::zeros({OC}, opts);

    const int blocks = (int)(N * OC);
    const int threads = 256;
    const int shmem = (int)(IC * 27 * sizeof(float));

    tconv_stats_kernel<<<blocks, threads, shmem>>>(
        x.data_ptr<float>(), w.data_ptr<float>(),
        sum.data_ptr<double>(), sumsq.data_ptr<double>(),
        (int)N, (int)IC, (int)OC, (int)ID, (int)OD);

    return {sum, sumsq};
}

torch::Tensor fused_tconv_bn_pool(torch::Tensor x, torch::Tensor C,
                                  torch::Tensor scale, torch::Tensor shift,
                                  int64_t N, int64_t OC, int64_t ID, int64_t BD) {
    auto out = torch::empty({N, OC, BD, BD, BD}, x.options());

    const int total = (int)(N * BD * BD * BD);
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;

    fused_tconv_bn_pool_kernel<3, 16><<<blocks, threads>>>(
        x.data_ptr<float>(), C.data_ptr<float>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>(), (int)N, (int)ID, (int)BD);

    return out;
}
"""

cpp_source = (
    "std::vector<torch::Tensor> tconv_stats(torch::Tensor x, torch::Tensor w, "
    "int64_t N, int64_t IC, int64_t OC, int64_t ID, int64_t OD);\n"
    "torch::Tensor fused_tconv_bn_pool(torch::Tensor x, torch::Tensor C, "
    "torch::Tensor scale, torch::Tensor shift, int64_t N, int64_t OC, "
    "int64_t ID, int64_t BD);"
)

_fused_ops = load_inline(
    name="fused_tconv_bn_pool",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["tconv_stats", "fused_tconv_bn_pool"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d + BatchNorm3d + AvgPool3d(2) + AvgPool3d(2), fully fused.

    The two 2x2x2 average pools collapse into one 4x4x4 average pool, and the
    affine BatchNorm commutes with the averaging, so the entire pipeline is
    evaluated as a single 3x3x3xIC stencil over the *input* per pooled voxel.
    The 63^3 transposed-conv output is never materialized.
    """

    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._in_channels = in_channels
        self._out_channels = out_channels
        self._kernel_size = kernel_size
        self._stride = stride
        self._padding = padding

        # Multiplicity matrix: M[di, k] = number of output positions inside a
        # pooled 4-block that input offset di reaches through kernel tap k.
        self.register_buffer(
            "_M",
            torch.tensor(
                [[0.0, 1.0, 1.0],
                 [1.0, 1.0, 1.0],
                 [1.0, 0.0, 0.0]]
            ),
            persistent=False,
        )

    def _fused_path_ok(self, x):
        return (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
            and self._kernel_size == 3
            and self._stride == 2
            and self._padding == 1
            and self._in_channels == 3
            and self._out_channels == 16
            and x.shape[1] == 3
        )

    def forward(self, x):
        if not self._fused_path_ok(x):
            y = self.conv_transpose(x)
            y = self.batch_norm(y)
            y = self.avg_pool1(y)
            y = self.avg_pool2(y)
            return y

        x = x.contiguous()
        N, IC, ID, IH, IW = x.shape
        w = self.conv_transpose.weight          # (IC, OC, 3, 3, 3)
        cbias = self.conv_transpose.bias        # (OC,)
        OC = w.shape[1]

        # Output size of the transposed conv and of the two pooled stages.
        OD = (ID - 1) * self._stride - 2 * self._padding + self._kernel_size
        BD = (OD // 2) // 2

        # Fold the two avg-pools into per-input-offset stencil coefficients:
        # C[oc, ic, a, b, c] = sum_{p,q,r} w[ic, oc, p, q, r] M[a,p] M[b,q] M[c,r]
        M = self._M.to(device=x.device, dtype=x.dtype)
        C = torch.einsum("iopqr,ap,bq,cr->oiabc", w, M, M, M).contiguous()

        bn = self.batch_norm
        gamma = bn.weight
        beta = bn.bias
        eps = bn.eps

        if self.training:
            stats = _fused_ops.tconv_stats(x, w.contiguous(), N, IC, OC, ID, OD)
            cnt = float(N) * float(OD) * float(OD) * float(OD)
            mean_d = stats[0] / cnt
            var_d = stats[1] / cnt - mean_d * mean_d
            var_d = torch.clamp(var_d, min=0.0)
            mean = mean_d.float()
            var = var_d.float()
            # Replicate BatchNorm running-stat bookkeeping.
            with torch.no_grad():
                mom = bn.momentum
                bn.running_mean.mul_(1.0 - mom).add_(mean, alpha=mom)
                unbiased = var * (cnt / (cnt - 1.0))
                bn.running_var.mul_(1.0 - mom).add_(unbiased, alpha=mom)
                bn.num_batches_tracked.add_(1)
        else:
            mean = bn.running_mean
            var = bn.running_var

        inv = torch.rsqrt(var + eps)
        # out = gamma/sd * (S/64 + cbias - mean) + beta
        #     = S * (gamma / (64*sd)) + (beta + gamma*(cbias-mean)/sd)
        scale = (gamma * inv / 64.0).contiguous()
        shift = (beta + gamma * (cbias - mean) * inv).contiguous()

        return _fused_ops.fused_tconv_bn_pool(x, C, scale, shift, N, OC, ID, BD)
```