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

# ---------------------------------------------------------------------------
# Algebra behind the fusion (kernel=3, stride=2, padding=1 transposed conv,
# then two stride-2 2x2x2 average pools):
#   * Two odd-sized pools (63->31->15) only touch convT outputs p in [0,59],
#     so final[q] = (1/64) * sum_{d in [0,3]^3} y[4q+d] collapses into a
#     *valid* stride-2 kernel-3 convolution on x[0..30] with an "effective"
#     kernel E stored transposed as float4 in constant memory.
#   * BatchNorm is a per-channel affine map and commutes with averaging, so
#     it becomes a scale/shift applied once per output value.
#   * Train mode: batch statistics of the full 63^3 convT output y are
#     computed on the fly (y never materialized).  Each thread produces the
#     y values for one (py, px) over all pz and all 16 channels with parity-
#     specialized, branch-free tap loops reading a shared input slab and
#     constant-memory weights, accumulating sum/sumsq in registers.
# ---------------------------------------------------------------------------

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

#define SLICE_ELEMS 250047   // 63*63*63

// Effective pooled kernel, layout E2[(ic*27 + j)*16 + oc], as float4.
__constant__ float4 Ec4[3 * 27 * 4];     // 5184 B
// Raw convT weights W[ic][oc][27] regrouped as [(ic*27+k)*4 + oc/4] float4.
__constant__ float4 Wc4[3 * 27 * 4];     // 5184 B

// grid = (IC*27) blocks, block = 16 threads (one per oc)
__global__ void build_eff_kernel(const float* __restrict__ W,
                                 float* __restrict__ E2,
                                 float* __restrict__ Wt,
                                 int IC, int OC) {
    const int ic = blockIdx.x / 27;
    const int j  = blockIdx.x % 27;
    const int oc = threadIdx.x;
    if (oc >= OC) return;
    const int jz = j / 9, jy = (j / 3) % 3, jx = j % 3;

    int kzl[3], kyl[3], kxl[3], nz, ny, nx;
    if (jz == 0)      { kzl[0]=1; kzl[1]=2; nz=2; }
    else if (jz == 1) { kzl[0]=0; kzl[1]=1; kzl[2]=2; nz=3; }
    else              { kzl[0]=0; nz=1; }
    if (jy == 0)      { kyl[0]=1; kyl[1]=2; ny=2; }
    else if (jy == 1) { kyl[0]=0; kyl[1]=1; kyl[2]=2; ny=3; }
    else              { kyl[0]=0; ny=1; }
    if (jx == 0)      { kxl[0]=1; kxl[1]=2; nx=2; }
    else if (jx == 1) { kxl[0]=0; kxl[1]=1; kxl[2]=2; nx=3; }
    else              { kxl[0]=0; nx=1; }

    const float* wb = W + ((long)ic * OC + oc) * 27;
    const float w = wb[j];
    float acc = 0.f;
    for (int a = 0; a < nz; ++a)
        for (int b = 0; b < ny; ++b)
            for (int c = 0; c < nx; ++c)
                acc += wb[kzl[a] * 9 + kyl[b] * 3 + kxl[c]];
    E2[(ic * 27 + j) * OC + oc] = acc * (1.0f / 64.0f);
    Wt[(ic * 27 + j) * OC + oc] = w;
}

// Fully fused eval path: effective stride-2 valid conv (k=3) + BN + pools.
// grid = (15, N), block = 256; thread -> (qy, qx).
__global__ void __launch_bounds__(256) fused_conv_bn_pool_kernel(
        const float* __restrict__ x,
        const float* __restrict__ bn_w,
        const float* __restrict__ bn_b,
        const float* __restrict__ bn_mean,
        const float* __restrict__ bn_var,
        const float* __restrict__ conv_bias,
        float* __restrict__ out,
        float eps) {
    __shared__ float xs[3 * 3 * 31 * 31];

    const int n   = blockIdx.y;
    const int qz  = blockIdx.x;
    const int tid = threadIdx.x;
    const int qy  = tid >> 4;
    const int qx  = tid & 15;

    const long xb = (long)n * 98304 + (long)(2 * qz) * 1024;
    #pragma unroll
    for (int c = 0; c < 9; ++c) {
        const float* src = x + xb + (long)(c / 3) * 32768 + (long)(c % 3) * 1024;
        float* dst = xs + c * 961;
        #pragma unroll
        for (int i = tid; i < 961; i += 256) {
            const int gy = i / 31;
            const int xx = i - gy * 31;
            dst[i] = src[gy * 32 + xx];
        }
    }
    __syncthreads();

    if (qy >= 15 || qx >= 15) return;

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

    #pragma unroll
    for (int ic = 0; ic < 3; ++ic) {
        #pragma unroll
        for (int j = 0; j < 27; ++j) {
            const int jz = j / 9, jy = (j / 3) % 3, jx = j % 3;
            const float v = xs[(((ic * 3 + jz) * 31) + (2 * qy + jy)) * 31 + 2 * qx + jx];
            const float4* e4 = Ec4 + (ic * 27 + j) * 4;
            #pragma unroll
            for (int g = 0; g < 4; ++g) {
                const float4 w4 = e4[g];
                acc[4*g+0] = fmaf(w4.x, v, acc[4*g+0]);
                acc[4*g+1] = fmaf(w4.y, v, acc[4*g+1]);
                acc[4*g+2] = fmaf(w4.z, v, acc[4*g+2]);
                acc[4*g+3] = fmaf(w4.w, v, acc[4*g+3]);
            }
        }
    }

    const long ob = ((long)n * 240 + qz) * 225 + qy * 15 + qx;
    #pragma unroll
    for (int oc = 0; oc < 16; ++oc) {
        const float scale = bn_w[oc] * rsqrtf(bn_var[oc] + eps);
        const float shift = bn_b[oc] + scale * (conv_bias[oc] - bn_mean[oc]);
        out[ob + (long)oc * 3375] = fmaf(acc[oc], scale, shift);
    }
}

// Train-mode batch statistics of the convT output y (never materialized).
// Slab layout: [2 iz][3 ic][9 iy][9 ix] = 486 floats.
// grid = (64, 4, 4): blockIdx.x = n, blockIdx.y = py-tile, blockIdx.z = px-tile.
// block = 256 threads -> (ty, tx); py = 16*by + ty, px = 16*bz + tx.
__global__ void __launch_bounds__(256) convt_stats_kernel(
        const float* __restrict__ x,
        const float* __restrict__ bias,
        double* __restrict__ gsum,
        double* __restrict__ gsumsq) {
    __shared__ float slab[486];
    __shared__ double red[16][8];

    const int n  = blockIdx.x;
    const int by = blockIdx.y;
    const int bz = blockIdx.z;
    const int ty = threadIdx.y;
    const int tx = threadIdx.x;
    const int tid = ty * 16 + tx;
    const int py = by * 16 + ty;
    const int px = bz * 16 + tx;
    const bool act = (py < 63) && (px < 63);

    const int iy = act ? (py >> 1) : 0;
    const int ix = act ? (px >> 1) : 0;
    const int iyr = iy - by * 8;
    const int ixr = ix - bz * 8;
    const bool pye = ((py & 1) == 0);
    const bool pxe = ((px & 1) == 0);
    // per-dim tap lists (shared-flattened index = ic*81 + (iyr+di)*9 + ixr+dj)
    int tapI[4], tapW[4], ntap;
    if (pye && pxe) {
        ntap = 1;
        tapI[0] = iyr * 9 + ixr;                 tapW[0] = 13;
    } else if (pye) {
        ntap = 2;
        tapI[0] = iyr * 9 + ixr;                 tapW[0] = 14;  // kx=2
        tapI[1] = iyr * 9 + ixr + 1;             tapW[1] = 12;  // kx=0
    } else if (pxe) {
        ntap = 2;
        tapI[0] = iyr * 9 + ixr;                 tapW[0] = 16;  // ky=2
        tapI[1] = (iyr + 1) * 9 + ixr;           tapW[1] = 10;  // ky=0
    } else {
        ntap = 4;
        tapI[0] = iyr * 9 + ixr;                 tapW[0] = 17;  // ky=2,kx=2
        tapI[1] = iyr * 9 + ixr + 1;             tapW[1] = 15;  // ky=2,kx=0
        tapI[2] = (iyr + 1) * 9 + ixr;           tapW[2] = 11;  // ky=0,kx=2
        tapI[3] = (iyr + 1) * 9 + ixr + 1;       tapW[3] = 9;   // ky=0,kx=0
    }

    const long xn = (long)n * 98304;   // 3*32*32*32

    float sumv[16], sqv[16], yv[16];
    #pragma unroll
    for (int oc = 0; oc < 16; ++oc) { sumv[oc] = 0.f; sqv[oc] = 0.f; }

    // load iz = 0 planes into slot 0
    {
        const float* src = x + xn;
        for (int i = tid; i < 243; i += 256) {
            const int ic = i / 81, r = i - (i / 81) * 81;
            const int ryy = r / 9, rxx = r - (r / 9) * 9;
            slab[i] = src[ic * 32768 + (by * 8 + ryy) * 32 + bz * 8 + rxx];
        }
    }
    __syncthreads();

    for (int pz = 0; pz < 63; ++pz) {
        const int m = pz >> 1;
        if ((pz & 1) && (m + 1 < 32)) {
            // load iz = m+1 into slot (m+1)&1
            float* dst = slab + (((unsigned)(m + 1) & 1u) ? 243 : 0);
            const float* src = x + xn + (long)(m + 1) * 1024;
            for (int i = tid; i < 243; i += 256) {
                const int ic = i / 81, r = i - (i / 81) * 81;
                const int ryy = r / 9, rxx = r - (r / 9) * 9;
                dst[i] = src[ic * 32768 + (by * 8 + ryy) * 32 + bz * 8 + rxx];
            }
        }
        __syncthreads();

        if (act) {
            const float biasv[16];
            #pragma unroll
            for (int oc = 0; oc < 16; ++oc) { yv[oc] = 0.f; }
            (void)biasv;

            const float* base = slab + ((unsigned)m & 1u ? 243 : 0);
            if (pz & 1) {
                // odd pz: iz=m (kz=2) and iz=m+1 (kz=0)
                #pragma unroll
                for (int ic = 0; ic < 3; ++ic) {
                    #pragma unroll
                    for (int t = 0; t < 4; ++t) {
                        if (t < ntap) {
                            const float v0 = base[ic * 81 + tapI[t]];
                            const float v1 = base[243 + ic * 81 + tapI[t]];
                            const float4* w0 = Wc4 + (ic * 27 + tapW[t] + 9) * 4;  // kz=2
                            const float4* w1 = Wc4 + (ic * 27 + tapW[t]) * 4;      // kz=0
                            #pragma unroll
                            for (int g = 0; g < 4; ++g) {
                                const float4 a = w0[g];
                                const float4 b = w1[g];
                                yv[4*g+0] = fmaf(a.x, v0, fmaf(b.x, v1, yv[4*g+0]));
                                yv[4*g+1] = fmaf(a.y, v0, fmaf(b.y, v1, yv[4*g+1]));
                                yv[4*g+2] = fmaf(a.z, v0, fmaf(b.z, v1, yv[4*g+2]));
                                yv[4*g+3] = fmaf(a.w, v0, fmaf(b.w, v1, yv[4*g+3]));
                            }
                        }
                    }
                }
            } else {
                // even pz: iz=m, kz=1
                #pragma unroll
                for (int ic = 0; ic < 3; ++ic) {
                    #pragma unroll
                    for (int t = 0; t < 4; ++t) {
                        if (t < ntap) {
                            const float v = base[ic * 81 + tapI[t]];
                            const float4* w0 = Wc4 + (ic * 27 + tapW[t] + 18) * 4;  // kz=1
                            #pragma unroll
                            for (int g = 0; g < 4; ++g) {
                                const float4 a = w0[g];
                                yv[4*g+0] = fmaf(a.x, v, yv[4*g+0]);
                                yv[4*g+1] = fmaf(a.y, v, yv[4*g+1]);
                                yv[4*g+2] = fmaf(a.z, v, yv[4*g+2]);
                                yv[4*g+3] = fmaf(a.w, v, yv[4*g+3]);
                            }
                        }
                    }
                }
            }
            #pragma unroll
            for (int oc = 0; oc < 16; ++oc) {
                const float yy = yv[oc] + bias[oc];
                sumv[oc] += yy;
                sqv[oc] = fmaf(yy, yy, sqv[oc]);
            }
        }
        __syncthreads();
    }

    // reduction: warp shfl (double), then across 8 warps in shared, then atomics
    #pragma unroll
    for (int oc = 0; oc < 16; ++oc) {
        double ds = (double)sumv[oc];
        double dq = (double)sqv[oc];
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            ds += __shfl_down_sync(0xffffffffu, ds, off);
            dq += __shfl_down_sync(0xffffffffu, dq, off);
        }
        if ((tid & 31) == 0) {
            red[oc][tid >> 5]     = ds;
            red[oc][8 + (tid >> 5)] = dq;
        }
    }
    __syncthreads();
    if (tid < 16) {
        double s = 0.0, q = 0.0;
        #pragma unroll
        for (int w = 0; w < 8; ++w) { s += red[tid][w]; q += red[tid][8 + w]; }
        atomicAdd(&gsum[tid], s);
        atomicAdd(&gsumsq[tid], q);
    }
}

// compute batch mean/var and update running stats like aten batch_norm
__global__ void bn_finalize_kernel(const double* __restrict__ sum,
                                   const double* __restrict__ sumsq,
                                   float* __restrict__ running_mean,
                                   float* __restrict__ running_var,
                                   float* __restrict__ out_mean,
                                   float* __restrict__ out_var,
                                   long count, double momentum, int C) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    double m  = sum[c] / (double)count;
    double sq = sumsq[c] / (double)count;
    double v  = sq - m * m;
    if (v < 0.0) v = 0.0;
    out_mean[c] = (float)m;
    out_var[c]  = (float)v;
    double unb = (count > 1) ? v * ((double)count / (double)(count - 1)) : 0.0;
    running_mean[c] = (float)((1.0 - momentum) * (double)running_mean[c] + momentum * m);
    running_var[c]  = (float)((1.0 - momentum) * (double)running_var[c]  + momentum * unb);
}

static inline void upload_coeffs(const torch::Tensor& W, cudaStream_t stream) {
    auto buf = torch::empty({2 * 3 * 27 * 16}, W.options());
    float* E2 = buf.data_ptr<float>();
    float* Wt = E2 + 3 * 27 * 16;
    build_eff_kernel<<<dim3(81), 16, 0, stream>>>(
        W.data_ptr<float>(), E2, Wt, 3, 16);
    cudaMemcpyToSymbolAsync(Ec4, E2, 3 * 27 * 16 * sizeof(float), 0,
                            cudaMemcpyDeviceToDevice, stream);
    cudaMemcpyToSymbolAsync(Wc4, Wt, 3 * 27 * 16 * sizeof(float), 0,
                            cudaMemcpyDeviceToDevice, stream);
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor W,
                         torch::Tensor bn_w, torch::Tensor bn_b,
                         torch::Tensor bn_mean, torch::Tensor bn_var,
                         torch::Tensor conv_bias,
                         int64_t N, double eps) {
    auto stream = at::cuda::getCurrentCUDAStream();
    upload_coeffs(W, stream.stream());
    auto out = torch::empty({N, 16, 15, 15, 15}, x.options());
    dim3 grid(15, (unsigned)N);
    fused_conv_bn_pool_kernel<<<grid, 256, 0, stream>>>(
        x.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        bn_mean.data_ptr<float>(), bn_var.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(), (float)eps);
    return out;
}

torch::Tensor fused_train(torch::Tensor x, torch::Tensor W,
                          torch::Tensor bn_w, torch::Tensor bn_b,
                          torch::Tensor running_mean, torch::Tensor running_var,
                          torch::Tensor conv_bias,
                          int64_t N, double eps, double momentum) {
    auto stream = at::cuda::getCurrentCUDAStream();
    upload_coeffs(W, stream.stream());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto sum    = torch::zeros({16}, opts_d);
    auto sumsq  = torch::zeros({16}, opts_d);

    dim3 sgrid((unsigned)N, 4, 4);
    dim3 sblock(16, 16);
    convt_stats_kernel<<<sgrid, sblock, 0, stream>>>(
        x.data_ptr<float>(), conv_bias.data_ptr<float>(),
        sum.data_ptr<double>(), sumsq.data_ptr<double>());

    auto bmean = torch::empty({16}, x.options());
    auto bvar  = torch::empty({16}, x.options());
    long count = (long)N * SLICE_ELEMS;
    bn_finalize_kernel<<<1, 64, 0, stream>>>(
        sum.data_ptr<double>(), sumsq.data_ptr<double>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        bmean.data_ptr<float>(), bvar.data_ptr<float>(),
        count, momentum, 16);

    auto out = torch::empty({N, 16, 15, 15, 15}, x.options());
    dim3 grid(15, (unsigned)N);
    fused_conv_bn_pool_kernel<<<grid, 256, 0, stream>>>(
        x.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        bmean.data_ptr<float>(), bvar.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(), (float)eps);
    return out;
}
"""

cpp_source = r"""
torch::Tensor fused_eval(torch::Tensor x, torch::Tensor W,
                         torch::Tensor bn_w, torch::Tensor bn_b,
                         torch::Tensor bn_mean, torch::Tensor bn_var,
                         torch::Tensor conv_bias,
                         int64_t N, double eps);
torch::Tensor fused_train(torch::Tensor x, torch::Tensor W,
                          torch::Tensor bn_w, torch::Tensor bn_b,
                          torch::Tensor running_mean, torch::Tensor running_var,
                          torch::Tensor conv_bias,
                          int64_t N, double eps, double momentum);
"""

_fused = load_inline(
    name="fused_convT_bn_avgpool_v4",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_eval", "fused_train"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized Model: conv-transpose (k=3, s=2, p=1) + batch-norm + two 2x2x2
    average pools fused into custom CUDA kernels.  The eval path collapses the
    whole chain into one valid stride-2 kernel-3 convolution with an
    effective kernel held in constant memory.  The train path additionally
    computes the batch statistics of the transposed-conv output on the fly
    with a parity-specialized shared-memory kernel (the 63^3-per-channel
    intermediate is never materialized), removing the cuDNN dgrad, memset,
    pooling and statistics passes entirely.
    """

    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._in_channels = in_channels
        self._out_channels = out_channels

    def _fallback(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = F.avg_pool3d(x, kernel_size=2)
        x = F.avg_pool3d(x, kernel_size=2)
        return x

    def forward(self, x):
        ic, oc = self._in_channels, self._out_channels
        bn = self.batch_norm
        ct = self.conv_transpose

        fast_ok = (
            x.is_cuda and x.dtype == torch.float32 and x.dim() == 5
            and x.size(1) == ic and x.size(2) == 32 and x.size(3) == 32 and x.size(4) == 32
            and ic == 3 and oc == 16
            and ct.bias is not None and bn.affine and bn.track_running_stats
        )
        if not fast_ok:
            return self._fallback(x)

        x = x.contiguous()
        w = ct.weight.contiguous()
        conv_bias = ct.bias.contiguous()
        N = x.shape[0]

        if not self.training:
            return _fused.fused_eval(
                x, w,
                bn.weight.contiguous(), bn.bias.contiguous(),
                bn.running_mean.contiguous(), bn.running_var.contiguous(),
                conv_bias, N, float(bn.eps),
            )
        else:
            return _fused.fused_train(
                x, w,
                bn.weight.contiguous(), bn.bias.contiguous(),
                bn.running_mean, bn.running_var,
                conv_bias, N, float(bn.eps), float(bn.momentum),
            )


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