```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 implementation of: MaxPool3d(k=2) -> MaxPool3d(k=3) -> sum(dim=1, keepdim=True)
#
# Key insight: MaxPool3d(kernel=2, stride=2) followed by MaxPool3d(kernel=3, stride=3)
# is exactly equivalent to a single non-overlapping MaxPool3d(kernel=6, stride=6), since
# each second-level window [3b, 3b+2] of the first pooling's output covers input rows
# [6b, 6b+5]. Combined with the (linear) channel sum, the whole tail of the network is:
#   out[b, 0, d, h, w] = sum_c max_{i,j,k in [0,6)} x[b, c, 6d+i, 6h+j, 6w+k]
# This is implemented as two lightweight kernels (a per-channel 6x6x6 max-pool, then a
# channel reduction), eliminating two large intermediate tensors and three kernel launches.

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

__global__ void pool6_max_kernel(
    const float* __restrict__ x,
    float* __restrict__ temp,
    int C, int D, int H, int W,
    int D2, int H2, int W2, int total)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int w  = idx % W2;
    int h  = (idx / W2) % H2;
    int d  = (idx / (W2 * H2)) % D2;
    int bc = idx / (W2 * H2 * D2);

    const float* base = x + (long)bc * ((long)D * H * W);
    int z0 = d * 6;
    int y0 = h * 6;
    int x0 = w * 6;

    float m = -FLT_MAX;
    #pragma unroll
    for (int i = 0; i < 6; ++i) {
        const float* pz = base + ((long)(z0 + i) * H + y0) * W + x0;
        #pragma unroll
        for (int j = 0; j < 6; ++j) {
            const float* py = pz + (long)j * W;
            #pragma unroll
            for (int k = 0; k < 6; ++k) {
                m = fmaxf(m, py[k]);
            }
        }
    }
    temp[idx] = m;
}

__global__ void channel_sum_kernel(
    const float* __restrict__ temp,
    float* __restrict__ out,
    int C, int N, int total)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int n = idx % N;
    int b = idx / N;
    const float* base = temp + (long)b * C * N + n;
    float s = 0.0f;
    for (int c = 0; c < C; ++c) {
        s += base[(long)c * N];
    }
    out[idx] = s;
}

torch::Tensor fused_pool_sum(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "input 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), matching MaxPool3d(2) then MaxPool3d(3) output sizes
    const int D2 = (D / 2) / 3;
    const int H2 = (H / 2) / 3;
    const int W2 = (W / 2) / 3;
    TORCH_CHECK(D2 > 0 && H2 > 0 && W2 > 0, "spatial size too small");

    auto temp = torch::empty({B, C, D2, H2, W2}, xc.options());
    auto out  = torch::empty({B, 1, D2, H2, W2}, xc.options());

    const int threads = 256;
    const int total1 = B * C * D2 * H2 * W2;
    const int blocks1 = (total1 + threads - 1) / threads;
    pool6_max_kernel<<<blocks1, threads>>>(
        xc.data_ptr<float>(), temp.data_ptr<float>(),
        C, D, H, W, D2, H2, W2, total1);

    const int N = D2 * H2 * W2;
    const int total2 = B * N;
    const int blocks2 = (total2 + threads - 1) / threads;
    channel_sum_kernel<<<blocks2, threads>>>(
        temp.data_ptr<float>(), out.data_ptr<float>(), C, N, total2);

    return out;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN ConvTranspose3d and replaces the
    max-pool -> max-pool -> channel-sum tail with a fused CUDA operator.
    """
    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_op = fused_pool_sum_ext

    def forward(self, x):
        x = self.conv_transpose(x)
        if x.is_cuda and x.dtype == torch.float32:
            return self.fused_op.fused_pool_sum(x)
        # Fallback (identical math): the two non-overlapping pools compose into
        # a single kernel=6, stride=6 max pool, followed by the channel sum.
        x = F.max_pool3d(x, kernel_size=6, stride=6)
        return torch.sum(x, dim=1, keepdim=True)
```