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

# Enable cuDNN autotuning for the (fixed-shape) transposed convolution
torch.backends.cudnn.benchmark = True

# Fused kernel:
#   MaxPool3d(k=2,s=2) followed by MaxPool3d(k=3,s=3) is mathematically
#   equivalent to a single non-overlapping max over 6x6x6 windows
#   (max of maxes, windows are aligned because stride == kernel size).
#   We fuse that 6x6x6 max-pool together with the final sum over channels
#   into a single kernel, eliminating all intermediate tensors and making
#   exactly one pass over the conv-transpose output.
fused_pool_sum_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <float.h>

__global__ void __launch_bounds__(256) fused_pool6_channel_sum_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    int C, int D, int H, int W,
    int D2, int H2, int W2)
{
    // One thread block per output voxel.
    int out_idx = blockIdx.x;
    int w2 = out_idx % W2;
    int tmp = out_idx / W2;
    int h2 = tmp % H2;
    tmp /= H2;
    int d2 = tmp % D2;
    int b  = tmp / D2;

    const int d0 = d2 * 6;
    const int h0 = h2 * 6;
    const int w0 = w2 * 6;
    const int nd = min(d0 + 6, D) - d0;
    const int nh = min(h0 + 6, H) - h0;
    const int nw = min(w0 + 6, W) - w0;
    const int nhnw = nh * nw;
    const int region = nd * nhnw;

    extern __shared__ float smax[];  // one per-channel max

    const int lane   = threadIdx.x & 31;
    const int warp   = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;

    const int HW  = H * W;
    const int DHW = D * H * W;
    const float* xb = x + (long)b * C * DHW;

    // Each warp computes the windowed max for a subset of channels.
    for (int c = warp; c < C; c += nwarps) {
        const float* base = xb + (long)c * DHW;
        float m = -FLT_MAX;
        for (int s = lane; s < region; s += 32) {
            int dd = s / nhnw;
            int r  = s - dd * nhnw;
            int hh = r / nw;
            int ww = r - hh * nw;
            float v = base[(d0 + dd) * HW + (h0 + hh) * W + (w0 + ww)];
            m = fmaxf(m, v);
        }
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, off));
        if (lane == 0) smax[c] = m;
    }
    __syncthreads();

    // First warp reduces the per-channel maxima into the final sum.
    if (warp == 0) {
        float s = 0.0f;
        for (int c = lane; c < C; c += 32) s += smax[c];
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            s += __shfl_xor_sync(0xffffffffu, s, off);
        if (lane == 0) out[out_idx] = s;
    }
}

torch::Tensor fused_pool_sum(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 5, "x must be 5D (N, C, D, H, W)");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");

    auto xc = x.contiguous();

    const int B = xc.size(0);
    const int C = xc.size(1);
    const int D = xc.size(2);
    const int H = xc.size(3);
    const int W = xc.size(4);

    // floor(floor(D/2)/3) == floor(D/6); matches the two stacked pools.
    const int D2 = D / 6;
    const int H2 = H / 6;
    const int W2 = W / 6;

    auto out = torch::empty({B, 1, D2, H2, W2}, xc.options());
    const long total = (long)B * D2 * H2 * W2;
    if (total == 0) return out;

    const int threads = 256;
    const size_t shmem = (size_t)C * sizeof(float);
    auto stream = at::cuda::getCurrentCUDAStream();

    fused_pool6_channel_sum_kernel<<<(int)total, threads, shmem, stream>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(),
        C, D, H, W, D2, H2, W2);

    return out;
}
"""

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

fused_pool_sum = load_inline(
    name="fused_pool6_channel_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: ConvTranspose3d (cuDNN) followed by a single fused
    CUDA kernel that performs the equivalent of MaxPool3d(2) + MaxPool3d(3)
    (combined into one 6x6x6 max pool) and the channel sum in one pass.
    """
    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_pool_sum = fused_pool_sum

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.fused_pool_sum.fused_pool_sum(x)
        return x
```