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

# Fused CUDA kernel:
#   MaxPool3d(kernel_size=2) -> MaxPool3d(kernel_size=3) -> sum(dim=1, keepdim=True)
#
# Key observation: maxpool(k=2, s=2) followed by maxpool(k=3, s=3) is exactly
# equivalent to a single maxpool with k=6, s=6 (nested floor property), i.e.
# output[b,c,t] = max over conv_out[b,c,6t:6t+6] in each dim. The trailing
# elements (indices 60..62 for D=63) are dropped exactly like the original
# sequential pooling does. We then fuse the channel-sum into the same kernel,
# so the whole post-conv pipeline reads the conv output exactly once and
# writes only the tiny (N,1,10,10,10) result, eliminating two large
# intermediate tensors (~250MB of extra traffic).

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

__global__ void fused_pool_sum_kernel(const float* __restrict__ x,
                                      float* __restrict__ out,
                                      int N, int C, int D, int T) {
    const int lane = threadIdx.x & 31;
    const int gid  = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
    const int total = N * T * T * T;
    if (gid >= total) return;

    // decode gid -> (b, d, h, w)
    int w = gid % T;
    int r = gid / T;
    int h = r % T;
    r /= T;
    int d = r % T;
    int b = r / T;

    const long SD = (long)D * D * D;
    const long SH = (long)D * D;

    // Precompute this lane's offsets inside the 6x6x6 window (216 voxels,
    // 32 lanes -> lanes 0..23 handle 7 voxels, lanes 24..31 handle 6).
    int offs[7];
    #pragma unroll
    for (int v = 0; v < 7; ++v) {
        int idx = v * 32 + lane;
        int i  = idx / 36;
        int rr = idx - i * 36;
        int j  = rr / 6;
        int k  = rr - j * 6;
        offs[v] = (int)(i * SH + j * D + k);
    }
    const int nv = (lane < 24) ? 7 : 6;

    const float* xb = x + (long)b * C * SD
                        + (long)(d * 6) * SH
                        + (long)(h * 6) * D
                        + (long)(w * 6);

    const float NEG_INF = __int_as_float(0xff800000);
    float sum = 0.f;

    #pragma unroll 4
    for (int c = 0; c < C; ++c) {
        const float* xc = xb + (long)c * SD;
        float m = NEG_INF;
        #pragma unroll
        for (int v = 0; v < 7; ++v) {
            if (v < nv) m = fmaxf(m, __ldg(xc + offs[v]));
        }
        // warp-wide max reduction
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1)
            m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o));
        sum += m;
    }

    if (lane == 0) out[gid] = sum;
}

torch::Tensor fused_pool_sum(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    auto xc = x.contiguous();

    const int N = xc.size(0);
    const int C = xc.size(1);
    const int D = xc.size(2);
    TORCH_CHECK(xc.size(3) == D && xc.size(4) == D, "expected cubic volume");

    // floor(floor(D/2)/3) == floor(D/6); window [6t, 6t+6) always in range
    const int T = D / 6;
    TORCH_CHECK(T > 0, "volume too small");

    auto out = torch::empty({N, 1, T, T, T}, xc.options());

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

    fused_pool_sum_kernel<<<blocks, threads>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(), N, C, D, T);

    return out;
}
"""

fused_pool_sum_cpp_source = "torch::Tensor fused_pool_sum(torch::Tensor x);"

fused_pool_sum_mod = load_inline(
    name="fused_pool_sum",
    cpp_sources=fused_pool_sum_cpp_source,
    cuda_sources=fused_pool_sum_cuda_source,
    functions=["fused_pool_sum"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN ConvTranspose3d (the dominant compute),
    and replaces MaxPool3d(2) + MaxPool3d(3) + channel-sum with a single
    fused custom CUDA kernel (mathematically equivalent to maxpool k=6,s=6
    followed by the sum), removing two large intermediate tensors and two
    kernel launches.
    """
    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.fused = fused_pool_sum_mod

    def forward(self, x):
        x = self.conv_transpose(x)
        if x.is_cuda and x.dtype == torch.float32:
            return self.fused.fused_pool_sum(x)
        # Fallback (CPU / non-fp32): original PyTorch ops
        x = F.max_pool3d(x, kernel_size=2)
        x = F.max_pool3d(x, kernel_size=3)
        x = torch.sum(x, dim=1, keepdim=True)
        return x
```