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

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

#define TILE 8
#define OC_TILE 8
#define IC_CHUNK 4
#define K 5
#define NTAP 125

// ---------------------------------------------------------------------------
// Custom ConvTranspose3d (gather formulation) for stride=2, kernel=5.
// out[b, oc, D, H, W] = sum_{ic,kd,kh,kw} X[b,ic,(D+pad-kd)/2,...] * W[ic,oc,kd,kh,kw]
// Only the parity-matched kernel taps are evaluated (saves ~7x FLOPs vs dense).
// Weight must be pre-packed to layout [IC][125][OC] (tap-major, oc contiguous).
// ---------------------------------------------------------------------------
__global__ void convt3d_s2k5_kernel(
    const float* __restrict__ X,
    const float* __restrict__ Wp,
    const float* __restrict__ bias,
    float* __restrict__ Y,
    int IC, int Di, int Hi, int Wi,
    int OC, int Do, int Ho, int Wo, int pad)
{
    __shared__ __align__(16) float sx[IC_CHUNK * 216];
    __shared__ __align__(16) float sw[NTAP * IC_CHUNK * OC_TILE];

    const int tid = threadIdx.x;                 // 0..511
    const int tx = tid & 7;
    const int ty = (tid >> 3) & 7;
    const int tz = tid >> 6;

    const int ntw = (Wo + TILE - 1) / TILE;
    const int nth = (Ho + TILE - 1) / TILE;
    const int tw = blockIdx.x % ntw;
    const int th = (blockIdx.x / ntw) % nth;
    const int td = blockIdx.x / (ntw * nth);

    const int W0 = tw * TILE, H0 = th * TILE, D0 = td * TILE;
    const int oc0 = blockIdx.y * OC_TILE;
    const int b   = blockIdx.z;

    const int W = W0 + tx, H = H0 + ty, D = D0 + tz;
    const bool valid_out = (D < Do) && (H < Ho) && (W < Wo);

    // input region required by this output tile
    const int d_lo = max(0, (D0 + pad - (K - 1)) / 2);
    const int d_hi = min(Di - 1, (D0 + TILE - 1 + pad) / 2);
    const int h_lo = max(0, (H0 + pad - (K - 1)) / 2);
    const int h_hi = min(Hi - 1, (H0 + TILE - 1 + pad) / 2);
    const int w_lo = max(0, (W0 + pad - (K - 1)) / 2);
    const int w_hi = min(Wi - 1, (W0 + TILE - 1 + pad) / 2);
    const int region_d = d_hi - d_lo + 1;
    const int region_h = h_hi - h_lo + 1;
    const int region_w = w_hi - w_lo + 1;

    // per-thread local tap indices (-1 = invalid)
    int dloc[K], hloc[K], wloc[K];
    #pragma unroll
    for (int kd = 0; kd < K; ++kd) {
        int t = D + pad - kd;
        int d = t >> 1;
        dloc[kd] = (valid_out && t >= 0 && !(t & 1) && d >= d_lo && d <= d_hi) ? (d - d_lo) : -1;
        t = H + pad - kd;
        d = t >> 1;
        hloc[kd] = (valid_out && t >= 0 && !(t & 1) && d >= h_lo && d <= h_hi) ? (d - h_lo) : -1;
        t = W + pad - kd;
        d = t >> 1;
        wloc[kd] = (valid_out && t >= 0 && !(t & 1) && d >= w_lo && d <= w_hi) ? (d - w_lo) : -1;
    }

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

    for (int ic0 = 0; ic0 < IC; ic0 += IC_CHUNK) {
        // stage input tile into shared memory
        #pragma unroll
        for (int idx = tid; idx < IC_CHUNK * 216; idx += 512) {
            int ic = idx / 216;
            int r  = idx % 216;
            int dz = r / 36;
            int dy = (r / 6) % 6;
            int dx = r % 6;
            float v = 0.f;
            if (dz < region_d && dy < region_h && dx < region_w) {
                const long xb = ((long)(b * IC + ic0 + ic) * Di + (d_lo + dz)) * Hi * Wi;
                v = X[xb + (h_lo + dy) * Wi + (w_lo + dx)];
            }
            sx[idx] = v;
        }
        // stage packed weights into shared memory
        #pragma unroll
        for (int idx = tid; idx < NTAP * IC_CHUNK * OC_TILE; idx += 512) {
            int oc  = idx & (OC_TILE - 1);
            int ic  = (idx / OC_TILE) & (IC_CHUNK - 1);
            int tap = idx / (OC_TILE * IC_CHUNK);
            sw[idx] = Wp[((long)(ic0 + ic) * NTAP + tap) * OC + (oc0 + oc)];
        }
        __syncthreads();

        if (valid_out) {
            #pragma unroll
            for (int kd = 0; kd < K; ++kd) {
                const int dl = dloc[kd];
                if (dl < 0) continue;
                #pragma unroll
                for (int kh = 0; kh < K; ++kh) {
                    const int hl = hloc[kh];
                    if (hl < 0) continue;
                    #pragma unroll
                    for (int kw = 0; kw < K; ++kw) {
                        const int wl = wloc[kw];
                        if (wl < 0) continue;
                        const int base  = (dl * 6 + hl) * 6 + wl;
                        const int tapid = kd * 25 + kh * 5 + kw;
                        #pragma unroll
                        for (int ic = 0; ic < IC_CHUNK; ++ic) {
                            const float xv = sx[ic * 216 + base];
                            const float4* wp4 = reinterpret_cast<const float4*>(&sw[(tapid * IC_CHUNK + ic) * OC_TILE]);
                            const float4 w0 = wp4[0];
                            const float4 w1 = wp4[1];
                            acc[0] = fmaf(xv, w0.x, acc[0]);
                            acc[1] = fmaf(xv, w0.y, acc[1]);
                            acc[2] = fmaf(xv, w0.z, acc[2]);
                            acc[3] = fmaf(xv, w0.w, acc[3]);
                            acc[4] = fmaf(xv, w1.x, acc[4]);
                            acc[5] = fmaf(xv, w1.y, acc[5]);
                            acc[6] = fmaf(xv, w1.z, acc[6]);
                            acc[7] = fmaf(xv, w1.w, acc[7]);
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    if (valid_out) {
        const long yb = ((long)b * OC + oc0) * Do * Ho * Wo + (long)D * Ho * Wo + (long)H * Wo + W;
        const long ostride = (long)Do * Ho * Wo;
        #pragma unroll
        for (int o = 0; o < OC_TILE; ++o) {
            Y[yb + (long)o * ostride] = acc[o] + bias[oc0 + o];
        }
    }
}

// ---------------------------------------------------------------------------
// Fused MaxPool(2)+MaxPool(3)+sum over channels.
// Composition of non-overlapping 2x2x2 then 3x3x3 max-pools equals a single
// non-overlapping 6x6x6 max-pool. Kernel computes, per output element, the
// per-channel max over the 6x6x6 window and sums over all 64 channels.
// One block (256 threads) per output element.
// ---------------------------------------------------------------------------
__global__ void pool6_sum_kernel(
    const float* __restrict__ Y,
    float* __restrict__ O,
    int C, int D, int H, int W,
    int OD, int OH, int OW)
{
    const int oid = blockIdx.x;
    const int ow = oid % OW;
    const int oh = (oid / OW) % OH;
    const int od = oid / (OW * OH);
    const int b  = blockIdx.y;

    const int tid  = threadIdx.x;   // 0..255
    const int c    = tid & 63;      // channel
    const int part = tid >> 6;      // 0..3

    const long HW = (long)H * W;
    const float* yb = Y + ((long)b * C + c) * D * HW + (long)6 * od * HW + (long)6 * oh * W + 6 * ow;

    float m = -FLT_MAX;
    for (int v = part; v < 216; v += 4) {
        const int dz = v / 36;
        const int dy = (v / 6) % 6;
        const int dx = v % 6;
        m = fmaxf(m, yb[(long)dz * HW + (long)dy * W + dx]);
    }

    __shared__ float smax[4][64];
    __shared__ float warpsum[2];
    smax[part][c] = m;
    __syncthreads();

    if (tid < 64) {
        float mm = fmaxf(fmaxf(smax[0][tid], smax[1][tid]),
                         fmaxf(smax[2][tid], smax[3][tid]));
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            mm += __shfl_down_sync(0xffffffffu, mm, off);
        if ((tid & 31) == 0) warpsum[tid >> 5] = mm;
    }
    __syncthreads();
    if (tid == 0) {
        O[((long)b * OD + od) * OH * OW + (long)oh * OW + ow] = warpsum[0] + warpsum[1];
    }
}

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor wp, torch::Tensor bias,
                                 int64_t stride, int64_t pad)
{
    const int B  = x.size(0);
    const int IC = x.size(1);
    const int Di = x.size(2);
    const int Hi = x.size(3);
    const int Wi = x.size(4);
    const int OC = wp.size(2);

    const int Do = (Di - 1) * (int)stride - 2 * (int)pad + K;
    const int Ho = (Hi - 1) * (int)stride - 2 * (int)pad + K;
    const int Wo = (Wi - 1) * (int)stride - 2 * (int)pad + K;

    auto Y = torch::empty({B, OC, Do, Ho, Wo}, x.options());

    const int ntw = (Wo + TILE - 1) / TILE;
    const int nth = (Ho + TILE - 1) / TILE;
    const int ntd = (Do + TILE - 1) / TILE;

    dim3 grid(ntw * nth * ntd, OC / OC_TILE, B);
    convt3d_s2k5_kernel<<<grid, 512>>>(
        x.data_ptr<float>(), wp.data_ptr<float>(), bias.data_ptr<float>(),
        Y.data_ptr<float>(), IC, Di, Hi, Wi, OC, Do, Ho, Wo, (int)pad);

    // fused pool1 (k=2,s=2) + pool2 (k=3,s=3) + channel sum
    const int D1 = (Do - 2) / 2 + 1;
    const int H1 = (Ho - 2) / 2 + 1;
    const int W1 = (Wo - 2) / 2 + 1;
    const int OD = (D1 - 3) / 3 + 1;
    const int OH = (H1 - 3) / 3 + 1;
    const int OW = (W1 - 3) / 3 + 1;

    auto O = torch::empty({B, 1, OD, OH, OW}, x.options());
    dim3 pgrid(OD * OH * OW, B, 1);
    pool6_sum_kernel<<<pgrid, 256>>>(
        Y.data_ptr<float>(), O.data_ptr<float>(), OC, Do, Ho, Wo, OD, OH, OW);

    return O;
}
"""

_cpp_src = "torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor wp, torch::Tensor bias, int64_t stride, int64_t pad);"

_fused = load_inline(
    name="convt3d_pool_sum_fused",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized Model:
      - ConvTranspose3d replaced by a custom tiled CUDA kernel that evaluates only
        parity-matched kernel taps (stride=2, k=5) with shared-memory blocking.
      - MaxPool3d(2) + MaxPool3d(3) + channel-sum fused into a single CUDA kernel
        exploiting that the composed pooling equals one non-overlapping 6x6x6 max-pool.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.max_pool1 = nn.MaxPool3d(kernel_size=2)
        self.max_pool2 = nn.MaxPool3d(kernel_size=3)
        self._packed_w = None
        self._packed_version = -1

    def _get_packed_weight(self, w):
        ver = w._version
        if (
            self._packed_w is None
            or self._packed_version != ver
            or self._packed_w.device != w.device
        ):
            # [IC, OC, K, K, K] -> [IC, K*K*K, OC]
            self._packed_w = w.permute(0, 2, 3, 4, 1).contiguous()
            self._packed_version = ver
        return self._packed_w

    def forward(self, x):
        ct = self.conv_transpose
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            and ct.kernel_size == (5, 5, 5)
            and ct.stride == (2, 2, 2)
            and ct.padding == (2, 2, 2)
            and ct.dilation == (1, 1, 1)
            and ct.output_padding == (0, 0, 0)
            and ct.groups == 1
            and ct.bias is not None
            and ct.in_channels % 4 == 0
            and ct.out_channels % 8 == 0
            and ct.out_channels == 64
            and ct.weight.is_cuda
            and ct.weight.dtype == torch.float32
        ):
            wp = self._get_packed_weight(ct.weight)
            return _fused.fused_forward_cuda(
                x, wp, ct.bias.contiguous(), ct.stride[0], ct.padding[0]
            )
        # generic fallback
        x = self.conv_transpose(x)
        x = self.max_pool1(x)
        x = self.max_pool2(x)
        x = torch.sum(x, dim=1, keepdim=True)
        return x
```