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

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

#define CO_TILE 16

// ---------------------------------------------------------------------------
// Fully fused kernel (eval mode):
//   ConvTranspose3d(k=3,s=2,p=1) + BatchNorm3d + AvgPool(2) + AvgPool(2)
//
// The two 2x2x2 average pools are exactly equivalent to averaging over
// non-overlapping 4x4x4 blocks (tail elements are dropped by both forms).
// Pool block (a,b,c) covers conv outputs [4a,4a+3]^3. A conv output at
// o = 2i-1+k (k in {0,1,2}) receives tap k from input i. Hence per dimension,
// input (2a+di) contributes to block a with tap subsets:
//   di=0 -> {1,2}, di=1 -> {0,1,2}, di=2 -> {0}.
// The 3D effective weight is the sum of the raw 3x3x3 kernel over the
// per-dimension tap subsets. BatchNorm (per-channel affine) commutes with
// the average pool, so it folds into an epilogue:
//   out = sum_{ci,d1,d2,d3} x[..] * Weff[ci,d1,d2,d3,co] + beta[co]
// where Weff already includes the BN scale / 64 factor and the conv bias
// folds into beta (every conv output carries it, and 64*bias/64 = bias).
// ---------------------------------------------------------------------------
__global__ void fused_ct3_bn_pp_kernel(
    const float* __restrict__ x,     // [N, Cin, D, H, W]
    const float* __restrict__ wt,    // [Cin, Cout, 3, 3, 3]
    const float* __restrict__ cbias, // [Cout] or nullptr
    const float* __restrict__ bnw,   // [Cout]
    const float* __restrict__ bnb,   // [Cout]
    const float* __restrict__ bnm,   // [Cout]
    const float* __restrict__ bnv,   // [Cout]
    const float eps,
    float* __restrict__ out,         // [N, Cout, PD, PH, PW]
    const int N, const int Cin, const int Cout, const int CoutPad,
    const int D, const int H, const int W,
    const int PD, const int PH, const int PW,
    const int total)
{
    extern __shared__ float smem[];
    float* sWeff  = smem;                               // Cin*27*CoutPad
    float* sAlpha = sWeff + (size_t)Cin * 27 * CoutPad; // CoutPad
    float* sBeta  = sAlpha + CoutPad;                   // CoutPad

    const int tid  = threadIdx.x;
    const int nthr = blockDim.x;

    for (int co = tid; co < CoutPad; co += nthr) {
        if (co < Cout) {
            float scale = bnw[co] * rsqrtf(bnv[co] + eps);
            float shift = bnb[co] - bnm[co] * scale;
            float cb = (cbias != nullptr) ? cbias[co] : 0.0f;
            sAlpha[co] = scale * (1.0f / 64.0f);
            sBeta[co]  = scale * cb + shift;
        } else {
            sAlpha[co] = 0.0f;
            sBeta[co]  = 0.0f;
        }
    }
    __syncthreads();

    // 1D tap subsets: di=0 -> {1,2}, di=1 -> {0,1,2}, di=2 -> {0}
    const int start[3] = {1, 0, 0};
    const int count[3] = {2, 3, 1};

    const int E = Cin * 27 * CoutPad;
    for (int e = tid; e < E; e += nthr) {
        int co = e % CoutPad;
        int r  = e / CoutPad;
        int d3 = r % 3; r /= 3;
        int d2 = r % 3; r /= 3;
        int d1 = r % 3;
        int ci = r / 3;
        float acc = 0.0f;
        if (co < Cout) {
            const float* wbase = wt + ((size_t)ci * Cout + co) * 27;
            for (int u = 0; u < count[d1]; ++u) {
                int kd = start[d1] + u;
                for (int v = 0; v < count[d2]; ++v) {
                    int kh = start[d2] + v;
                    #pragma unroll
                    for (int q = 0; q < 3; ++q) {
                        if (q < count[d3]) {
                            int kw = start[d3] + q;
                            acc += wbase[kd * 9 + kh * 3 + kw];
                        }
                    }
                }
            }
        }
        sWeff[e] = acc * sAlpha[co];
    }
    __syncthreads();

    const int g = blockIdx.x * nthr + tid;
    if (g >= total) return;

    int c = g % PW;
    int t = g / PW;
    int b = t % PH;
    t /= PH;
    int a = t % PD;
    int n = t / PD;

    const float* xn = x + (size_t)n * Cin * D * H * W;
    const int d0 = 2 * a, h0 = 2 * b, w0 = 2 * c;
    const size_t HW = (size_t)H * W;
    const size_t pvol = (size_t)PD * PH * PW;

    for (int co0 = 0; co0 < Cout; co0 += CO_TILE) {
        float acc[CO_TILE];
        #pragma unroll
        for (int j = 0; j < CO_TILE; ++j) acc[j] = 0.0f;

        for (int ci = 0; ci < Cin; ++ci) {
            const float* xc = xn + (size_t)ci * HW;
            #pragma unroll
            for (int d1 = 0; d1 < 3; ++d1) {
                #pragma unroll
                for (int d2 = 0; d2 < 3; ++d2) {
                    const float* xr = xc + ((size_t)(d0 + d1) * H + (h0 + d2)) * W + w0;
                    float x0 = __ldg(xr);
                    float x1 = __ldg(xr + 1);
                    float x2 = __ldg(xr + 2);
                    const float* wb = sWeff + ((size_t)((ci * 3 + d1) * 3 + d2) * 3) * CoutPad + co0;
                    #pragma unroll
                    for (int j = 0; j < CO_TILE; ++j) {
                        acc[j] += x0 * wb[j];
                        acc[j] += x1 * wb[CoutPad + j];
                        acc[j] += x2 * wb[2 * CoutPad + j];
                    }
                }
            }
        }

        int nco = Cout - co0; if (nco > CO_TILE) nco = CO_TILE;
        float* ob = out + (((size_t)n * Cout + co0) * PD + a) * PH * PW + (size_t)b * PW + c;
        #pragma unroll
        for (int j = 0; j < CO_TILE; ++j) {
            if (j < nco) ob[(size_t)j * pvol] = acc[j] + sBeta[co0 + j];
        }
    }
}

// ---------------------------------------------------------------------------
// Fused affine + double average pool (used in training mode with batch
// statistics): out = alpha[c] * sum_{4x4x4 block} y + beta[c]
// ---------------------------------------------------------------------------
__global__ void affine_pool4_kernel(
    const float* __restrict__ y,     // [N, C, D, H, W]
    const float* __restrict__ alpha, // [C]
    const float* __restrict__ beta,  // [C]
    float* __restrict__ out,         // [N, C, PD, PH, PW]
    const int C, const int D, const int H, const int W,
    const int PD, const int PH, const int PW,
    const int total)
{
    const int g = blockIdx.x * blockDim.x + threadIdx.x;
    if (g >= total) return;

    int c = g % PW;
    int t = g / PW;
    int b = t % PH;
    t /= PH;
    int a = t % PD;
    t /= PD;
    int ch = t % C;
    int n  = t / C;

    const float* base = y + ((((size_t)n * C + ch) * D + 4 * a) * H + 4 * b) * W + 4 * c;
    const size_t Ws = (size_t)W;
    const size_t Hs = (size_t)H * W;

    float s = 0.0f;
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const float* p = base + (size_t)i * Hs + (size_t)j * Ws;
            s += p[0]; s += p[1]; s += p[2]; s += p[3];
        }
    }

    out[g] = fmaf(s, alpha[ch], beta[ch]);
}

torch::Tensor fused_forward(
    torch::Tensor x, torch::Tensor wt, torch::Tensor cbias,
    torch::Tensor bnw, torch::Tensor bnb, torch::Tensor bnm, torch::Tensor bnv,
    double eps)
{
    const int N = x.size(0), Cin = x.size(1);
    const int D = x.size(2), H = x.size(3), W = x.size(4);
    const int Cout = wt.size(1);
    const int PD = (D - 1) / 2, PH = (H - 1) / 2, PW = (W - 1) / 2;
    if (PD <= 0 || PH <= 0 || PW <= 0 || Cout <= 0 || Cin <= 0) return torch::Tensor();

    const int CoutPad = ((Cout + CO_TILE - 1) / CO_TILE) * CO_TILE;
    const size_t smem = ((size_t)Cin * 27 + 2) * (size_t)CoutPad * sizeof(float);
    if (smem > 46 * 1024) return torch::Tensor();  // caller falls back

    auto out = torch::empty({N, Cout, PD, PH, PW}, x.options());
    const long total = (long)N * PD * PH * PW;
    const int block = 256;
    const int grid = (int)((total + block - 1) / block);
    const float* cb = (cbias.defined() && cbias.numel() > 0) ? cbias.data_ptr<float>() : nullptr;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_ct3_bn_pp_kernel<<<grid, block, smem, stream>>>(
        x.data_ptr<float>(), wt.data_ptr<float>(), cb,
        bnw.data_ptr<float>(), bnb.data_ptr<float>(), bnm.data_ptr<float>(), bnv.data_ptr<float>(),
        (float)eps, out.data_ptr<float>(),
        N, Cin, Cout, CoutPad, D, H, W, PD, PH, PW, (int)total);
    return out;
}

torch::Tensor affine_pool4(torch::Tensor y, torch::Tensor alpha, torch::Tensor beta)
{
    const int N = y.size(0), C = y.size(1);
    const int D = y.size(2), H = y.size(3), W = y.size(4);
    const int PD = ((D - 1) / 2) / 2, PH = ((H - 1) / 2) / 2, PW = ((W - 1) / 2) / 2;
    TORCH_CHECK(PD > 0 && PH > 0 && PW > 0, "bad dims");

    auto out = torch::empty({N, C, PD, PH, PW}, y.options());
    const long total = (long)N * C * PD * PH * PW;
    const int block = 256;
    const int grid = (int)((total + block - 1) / block);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    affine_pool4_kernel<<<grid, block, 0, stream>>>(
        y.data_ptr<float>(), alpha.data_ptr<float>(), beta.data_ptr<float>(),
        out.data_ptr<float>(), C, D, H, W, PD, PH, PW, (int)total);
    return out;
}
"""

_cpp_src = r"""
torch::Tensor fused_forward(
    torch::Tensor x, torch::Tensor wt, torch::Tensor cbias,
    torch::Tensor bnw, torch::Tensor bnb, torch::Tensor bnm, torch::Tensor bnv,
    double eps);
torch::Tensor affine_pool4(torch::Tensor y, torch::Tensor alpha, torch::Tensor beta);
"""

_ext = load_inline(
    name="fused_ct3_bn_pool_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward", "affine_pool4"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Key ideas:
      * The two consecutive AvgPool3d(2) layers are mathematically equivalent
        to a single average over non-overlapping 4x4x4 blocks (both drop the
        same tail elements), so they collapse into one reduction.
      * BatchNorm3d is a per-channel affine transform, which commutes with the
        average pool; in eval mode it folds entirely into an epilogue.
      * In eval mode the whole network (ConvTranspose3d k=3/s=2/p=1 + BN +
        pools) is executed by a single fused CUDA kernel that never
        materializes the large 63^3 intermediate: each pooled output is
        computed directly from the 3x3x3 input neighborhood through
        pre-computed "effective" 3x3x3 weights (sums of kernel tap subsets).
      * In training mode, cuDNN computes the transposed convolution, batch
        statistics come from one reduction, and a fused affine+pool4 kernel
        finishes (running stats are updated exactly like BatchNorm3d).
    """

    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)

    # -- fused eval path applicability ------------------------------------
    def _can_fuse_eval(self, x):
        ct, bn = self.conv_transpose, self.batch_norm
        if self.training:
            return False
        if not (x.is_cuda and x.dtype == torch.float32 and x.dim() == 5):
            return False
        if tuple(ct.kernel_size) != (3, 3, 3) or tuple(ct.stride) != (2, 2, 2):
            return False
        if tuple(ct.padding) != (1, 1, 1) or tuple(ct.output_padding) != (0, 0, 0):
            return False
        if tuple(ct.dilation) != (1, 1, 1) or ct.groups != 1:
            return False
        if not bn.affine or not bn.track_running_stats:
            return False
        if min(x.size(2), x.size(3), x.size(4)) < 3:
            return False
        return True

    # -- fused train path applicability -----------------------------------
    def _can_fuse_train(self, x):
        bn = self.batch_norm
        if not self.training:
            return False
        if not (x.is_cuda and x.dtype == torch.float32 and x.dim() == 5):
            return False
        if not bn.affine or not bn.track_running_stats or bn.momentum is None:
            return False
        return True

    def forward(self, x):
        # Fast path 1: eval mode, everything in one kernel.
        if self._can_fuse_eval(x):
            bias = self.conv_transpose.bias
            if bias is None:
                bias = torch.empty(0, device=x.device, dtype=torch.float32)
            out = _ext.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                bias.contiguous(),
                self.batch_norm.weight.contiguous(),
                self.batch_norm.bias.contiguous(),
                self.batch_norm.running_mean.contiguous(),
                self.batch_norm.running_var.contiguous(),
                float(self.batch_norm.eps),
            )
            if out is not None:
                return out

        # Fast path 2: training mode, cuDNN conv + fused BN(batch stats) + pools.
        if self._can_fuse_train(x):
            ct, bn = self.conv_transpose, self.batch_norm
            y = F.conv_transpose3d(
                x, ct.weight, ct.bias, ct.stride, ct.padding,
                ct.output_padding, ct.groups, ct.dilation,
            )
            reduce_dims = (0, 2, 3, 4)
            var, mean = torch.var_mean(y, dim=reduce_dims, unbiased=False, keepdim=False)
            count = y.numel() // y.size(1)
            if count <= 1:
                # Match BatchNorm3d's behavior for degenerate batches.
                return self.avg_pool2(self.avg_pool1(bn(self.conv_transpose(x))))
            with torch.no_grad():
                m = bn.momentum
                bn.running_mean.mul_(1.0 - m).add_(mean, alpha=m)
                bn.running_var.mul_(1.0 - m).add_(var * (float(count) / (count - 1)), alpha=m)
                bn.num_batches_tracked.add_(1)
            alpha = bn.weight * torch.rsqrt(var + bn.eps) * (1.0 / 64.0)
            beta = bn.bias - mean * (bn.weight * torch.rsqrt(var + bn.eps))
            return _ext.affine_pool4(
                y.contiguous(), alpha.contiguous(), beta.contiguous()
            )

        # Reference fallback for any unsupported configuration.
        y = self.conv_transpose(x)
        y = self.batch_norm(y)
        y = self.avg_pool1(y)
        y = self.avg_pool2(y)
        return y
```