```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
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernels:
#  1) fused_convT_bn_pool_kernel: a fully fused ConvTranspose3d(k=3,s=2,p=1)
#     + BatchNorm3d(affine) + AvgPool3d(2) + AvgPool3d(2) kernel.
#     Key observations used:
#       - Two 2x2x2 average pools = one 4x4x4 average (all windows are full).
#       - Per-channel BN affine commutes with the spatial average, so the
#         final output is a_c * (sum of 4x4x4 conv outputs)/64 + b_c.
#       - With k=3, s=2, p=1 and output size 2*D-1 there is no border
#         clipping, and each 4x4x4 output region only touches a 3x3x3 input
#         neighborhood.  The 1536 tap/position products factor into
#         8 reusable sub-block sums per input channel and 27 weights.
#  2) conv_trans_stats_kernel: computes per-channel sum / sum-of-squares of
#     the transposed-convolution output so that batch statistics (training
#     mode) can also use the fused fast path.

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

// ---------------------------------------------------------------------------
// Fused ConvTranspose3d(k=3,s=2,p=1) + BN(affine) + pool(2) + pool(2)
// out[n,c,qz,qy,qx] = abn[c] * (sum of conv outputs over the 4x4x4 region
//                                starting at (4*qz,4*qy,4*qx)) / 64 + bbn[c]
// where bbn already includes the conv bias contribution.
// ---------------------------------------------------------------------------
__global__ void fused_convT_bn_pool_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ abn,
    const float* __restrict__ bbn,
    float* __restrict__ out,
    int N, int Cin, int D, int H, int Wd,
    int Cout, int OD, int OH, int OW)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * Cout * OD * OH * OW;
    if (idx >= total) return;

    int ox = (int)(idx % OW);
    long long t = idx / OW;
    int oy = (int)(t % OH); t /= OH;
    int oz = (int)(t % OD); t /= OD;
    int c  = (int)(t % Cout);
    int n  = (int)(t / Cout);

    const int i0 = 2 * oz, j0 = 2 * oy, k0 = 2 * ox;
    const float* xn = x + (long long)n * Cin * D * H * Wd;
    const int HW = H * Wd;

    float acc = 0.0f;
    for (int ci = 0; ci < Cin; ++ci) {
        const float* xp = xn + (long long)ci * HW;
        float v[3][3][3];
        #pragma unroll
        for (int a = 0; a < 3; ++a) {
            #pragma unroll
            for (int b = 0; b < 3; ++b) {
                #pragma unroll
                for (int e = 0; e < 3; ++e) {
                    v[a][b][e] = xp[(long long)(i0 + a) * HW + (j0 + b) * Wd + (k0 + e)];
                }
            }
        }
        // sub-block sums along w
        float s1[3][3][2];
        #pragma unroll
        for (int a = 0; a < 3; ++a) {
            #pragma unroll
            for (int b = 0; b < 3; ++b) {
                s1[a][b][0] = v[a][b][0] + v[a][b][1];
                s1[a][b][1] = v[a][b][1] + v[a][b][2];
            }
        }
        // along h
        float s2[3][2][2];
        #pragma unroll
        for (int a = 0; a < 3; ++a) {
            #pragma unroll
            for (int se = 0; se < 2; ++se) {
                s2[a][0][se] = s1[a][0][se] + s1[a][1][se];
                s2[a][1][se] = s1[a][1][se] + s1[a][2][se];
            }
        }
        // along d : s3[sa][sb][se], index 0 -> {0,1}, 1 -> {1,2}
        float s3[2][2][2];
        #pragma unroll
        for (int sb = 0; sb < 2; ++sb) {
            #pragma unroll
            for (int se = 0; se < 2; ++se) {
                s3[0][sb][se] = s2[0][sb][se] + s2[1][sb][se];
                s3[1][sb][se] = s2[1][sb][se] + s2[2][sb][se];
            }
        }
        // contract with weights; kernel tap 0 selects the {1,2} sum,
        // taps 1 and 2 select the {0,1} sum.
        const float* wp = w + ((long long)ci * Cout + c) * 27;
        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    acc += wp[kd * 9 + kh * 3 + kw] *
                           s3[kd == 0 ? 1 : 0][kh == 0 ? 1 : 0][kw == 0 ? 1 : 0];
                }
            }
        }
    }
    out[idx] = abn[c] * (acc * (1.0f / 64.0f)) + bbn[c];
}

// ---------------------------------------------------------------------------
// Per-channel sum / sum-of-squares of the ConvTranspose3d output
// (used for BatchNorm batch statistics in training mode).
// grid: (ceil(spatial/256), Cout, N)
// ---------------------------------------------------------------------------
__global__ void conv_trans_stats_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    double* __restrict__ sum,
    double* __restrict__ sumsq,
    int N, int Cin, int D, int H, int Wd,
    int Cout, int Dc, int Hc, int Wc)
{
    int c = blockIdx.y;
    int n = blockIdx.z;
    int spatial = Dc * Hc * Wc;
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    float val = 0.0f;
    bool active = tid < spatial;
    if (active) {
        int ow = tid % Wc;
        int t = tid / Wc;
        int oh = t % Hc;
        int od = t / Hc;

        int id[2], kd_[2], ih[2], kh_[2], iw[2], kw_[2];
        int nd, nh, nw;
        if ((od & 1) == 0) { nd = 1; id[0] = od >> 1;        kd_[0] = 1; }
        else { nd = 2; id[0] = (od + 1) >> 1; kd_[0] = 0; id[1] = (od - 1) >> 1; kd_[1] = 2; }
        if ((oh & 1) == 0) { nh = 1; ih[0] = oh >> 1;        kh_[0] = 1; }
        else { nh = 2; ih[0] = (oh + 1) >> 1; kh_[0] = 0; ih[1] = (oh - 1) >> 1; kh_[1] = 2; }
        if ((ow & 1) == 0) { nw = 1; iw[0] = ow >> 1;        kw_[0] = 1; }
        else { nw = 2; iw[0] = (ow + 1) >> 1; kw_[0] = 0; iw[1] = (ow - 1) >> 1; kw_[1] = 2; }

        const int HW = H * Wd;
        const float* xn = x + (long long)n * Cin * D * HW;
        for (int ci = 0; ci < Cin; ++ci) {
            const float* xp = xn + (long long)ci * HW;
            const float* wp = w + ((long long)ci * Cout + c) * 27;
            for (int a = 0; a < nd; ++a) {
                for (int b = 0; b < nh; ++b) {
                    for (int e = 0; e < nw; ++e) {
                        val += wp[kd_[a] * 9 + kh_[b] * 3 + kw_[e]] *
                               xp[(long long)id[a] * HW + ih[b] * Wd + iw[e]];
                    }
                }
            }
        }
        val += bias[c];
    }

    __shared__ float s_sum[256];
    __shared__ float s_sq[256];
    int lane = threadIdx.x;
    s_sum[lane] = active ? val : 0.0f;
    s_sq[lane]  = active ? val * val : 0.0f;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (lane < s) {
            s_sum[lane] += s_sum[lane + s];
            s_sq[lane]  += s_sq[lane + s];
        }
        __syncthreads();
    }
    if (lane == 0) {
        atomicAdd(&sum[c],   (double)s_sum[0]);
        atomicAdd(&sumsq[c], (double)s_sq[0]);
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
torch::Tensor fused_convT_bn_pool(
    torch::Tensor x, torch::Tensor w, torch::Tensor abn, torch::Tensor bbn,
    int64_t OD, int64_t OH, int64_t OW)
{
    int N = x.size(0), Cin = x.size(1), D = x.size(2), H = x.size(3), Wd = x.size(4);
    int Cout = w.size(1);
    auto out = torch::empty({N, Cout, OD, OH, OW}, x.options());
    long long total = (long long)N * Cout * OD * OH * OW;
    if (total == 0) return out;
    int block = 256;
    long long grid = (total + block - 1) / block;
    fused_convT_bn_pool_kernel<<<(unsigned int)grid, block>>>(
        x.data_ptr<float>(), w.data_ptr<float>(),
        abn.data_ptr<float>(), bbn.data_ptr<float>(),
        out.data_ptr<float>(),
        N, Cin, D, H, Wd, Cout, (int)OD, (int)OH, (int)OW);
    return out;
}

void conv_trans_stats(
    torch::Tensor x, torch::Tensor w, torch::Tensor bias,
    torch::Tensor sum, torch::Tensor sumsq,
    int64_t Dc, int64_t Hc, int64_t Wc)
{
    int N = x.size(0), Cin = x.size(1), D = x.size(2), H = x.size(3), Wd = x.size(4);
    int Cout = w.size(1);
    int spatial = (int)(Dc * Hc * Wc);
    int block = 256;
    dim3 grid((spatial + block - 1) / block, Cout, N);
    conv_trans_stats_kernel<<<grid, block>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        sum.data_ptr<double>(), sumsq.data_ptr<double>(),
        N, Cin, D, H, Wd, Cout, (int)Dc, (int)Hc, (int)Wc);
}
"""

cpp_source = """
torch::Tensor fused_convT_bn_pool(
    torch::Tensor x, torch::Tensor w, torch::Tensor abn, torch::Tensor bbn,
    int64_t OD, int64_t OH, int64_t OW);
void conv_trans_stats(
    torch::Tensor x, torch::Tensor w, torch::Tensor bias,
    torch::Tensor sum, torch::Tensor sumsq,
    int64_t Dc, int64_t Hc, int64_t Wc);
"""

_fused_mod = load_inline(
    name="fused_convT_bn_pool",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_convT_bn_pool", "conv_trans_stats"],
    verbose=False,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    """
    Optimized model: ConvTranspose3d + BatchNorm3d + 2x AvgPool3d(2) fused
    into custom CUDA kernels.  The two non-overlapping 2x2x2 average pools
    are combined into a single 4x4x4 average, and the per-channel BatchNorm
    affine transform commutes with that average, so the whole pipeline
    collapses into one fused kernel (plus a statistics kernel for
    BatchNorm's batch statistics when in training mode).
    """

    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._fused = _fused_mod
        # The fused kernels implement the k=3, stride=2, padding=1 case.
        self._fast_ok = (kernel_size == 3 and stride == 2 and padding == 1)

    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):
        if (not self._fast_ok) or (not x.is_cuda) or x.dtype != torch.float32:
            return self._fallback(x)

        bn = self.batch_norm
        conv = self.conv_transpose
        xc = x.contiguous()
        w = conv.weight.contiguous()
        Cout = w.size(1)

        D, H, Wd = xc.size(2), xc.size(3), xc.size(4)
        Dc, Hc, Wc = 2 * D - 1, 2 * H - 1, 2 * Wd - 1   # conv transpose output dims
        OD, OH, OW = Dc // 4, Hc // 4, Wc // 4          # after pool(2) + pool(2)
        if OD < 1 or OH < 1 or OW < 1:
            return self._fallback(x)

        gamma = bn.weight
        beta = bn.bias
        conv_bias = conv.bias

        if self.training:
            # Batch statistics of the conv-transpose output via custom kernel.
            cnt = float(xc.size(0)) * float(Dc) * float(Hc) * float(Wc)
            sums = torch.zeros(Cout, dtype=torch.float64, device=xc.device)
            sumsq = torch.zeros(Cout, dtype=torch.float64, device=xc.device)
            self._fused.conv_trans_stats(xc, w, conv_bias.contiguous(), sums, sumsq, Dc, Hc, Wc)
            mean = (sums / cnt).to(torch.float32)
            var = (sumsq / cnt - (sums / cnt) * (sums / cnt)).to(torch.float32)
            var = torch.clamp(var, min=0.0)
            # Faithful running-stat update (unbiased variance, like PyTorch).
            with torch.no_grad():
                m = bn.momentum
                bn.running_mean.mul_(1.0 - m).add_(mean, alpha=m)
                var_u = var * (cnt / (cnt - 1.0))
                bn.running_var.mul_(1.0 - m).add_(var_u, alpha=m)
                if bn.track_running_stats:
                    bn.num_batches_tracked.add_(1)
        else:
            mean = bn.running_mean
            var = bn.running_var

        abn = gamma * torch.rsqrt(var + bn.eps)
        # BN affine + conv bias folded together (both commute with the average).
        bbn = beta - mean * abn + conv_bias * abn

        return self._fused.fused_convT_bn_pool(xc, w, abn.contiguous(), bbn.contiguous(),
                                               OD, OH, OW)
```