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

# Let cuDNN autotune the best algorithm for the (fixed) transposed-conv shape.
torch.backends.cudnn.benchmark = True

# Fused kernel: MaxPool3d(k=2,s=2) -> MaxPool3d(k=3,s=3) -> sum(dim=1, keepdim=True)
#
# The two non-overlapping max pools compose exactly into a single non-overlapping
# 6x6x6 max pool with stride 6:
#   pool1[i]   = max(x[2i], x[2i+1])
#   pool2[m]   = max(pool1[3m], pool1[3m+1], pool1[3m+2])
#              = max_{u in [0,6)} x[6m+u]
# so:
#   out[b,0,m,n,p] = sum_c max_{0<=u,v,w<6} x[b, c, 6m+u, 6n+v, 6p+w]
# This reads the conv output exactly once (no intermediate pooled tensors).
fused_pool_chansum_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void fused_pool_chansum_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    int C, int D, int H, int W,
    int oD, int oH, int oW)
{
    const int ow = blockIdx.x;
    const int oh = blockIdx.y;
    const int od = blockIdx.z % oD;
    const int b  = blockIdx.z / oD;

    const int base_d = od * 6;
    const int base_h = oh * 6;
    const int base_w = ow * 6;

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

    // Each thread handles one (or more, strided) channels: compute the max over
    // its 6x6x6 window, then block-reduce the per-channel maxes with a sum.
    float local = 0.0f;
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        const float* x0 = xb + (long)c * DHW
                            + (long)base_d * HW
                            + (long)base_h * W
                            + base_w;
        // Six independent max accumulators to expose instruction-level parallelism.
        float m0 = -FLT_MAX, m1 = -FLT_MAX, m2 = -FLT_MAX;
        float m3 = -FLT_MAX, m4 = -FLT_MAX, m5 = -FLT_MAX;
        #pragma unroll
        for (int u = 0; u < 6; ++u) {
            #pragma unroll
            for (int v = 0; v < 6; ++v) {
                const float* row = x0 + (long)u * HW + (long)v * W;
                m0 = fmaxf(m0, __ldg(row + 0));
                m1 = fmaxf(m1, __ldg(row + 1));
                m2 = fmaxf(m2, __ldg(row + 2));
                m3 = fmaxf(m3, __ldg(row + 3));
                m4 = fmaxf(m4, __ldg(row + 4));
                m5 = fmaxf(m5, __ldg(row + 5));
            }
        }
        local += fmaxf(fmaxf(fmaxf(m0, m1), fmaxf(m2, m3)), fmaxf(m4, m5));
    }

    // Block reduction (sum) across threads.
    __shared__ float sdata[32];
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int wid  = tid >> 5;

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        local += __shfl_down_sync(0xffffffffu, local, offset);

    if (lane == 0) sdata[wid] = local;
    __syncthreads();

    if (wid == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        float v = (lane < nwarps) ? sdata[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, offset);
        if (lane == 0)
            out[((long)b * oD + od) * ((long)oH * oW) + (long)oh * oW + ow] = v;
    }
}

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

    auto xc = x.contiguous();

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

    // Output sizes of MaxPool3d(2) then MaxPool3d(3) (stride=kernel, no padding).
    const int p1D = (D - 2) / 2 + 1;
    const int p1H = (H - 2) / 2 + 1;
    const int p1W = (W - 2) / 2 + 1;
    const int oD = (p1D - 3) / 3 + 1;
    const int oH = (p1H - 3) / 3 + 1;
    const int oW = (p1W - 3) / 3 + 1;

    if (oD <= 0 || oH <= 0 || oW <= 0 || B <= 0) {
        return torch::empty({B, 1, std::max(oD, 0), std::max(oH, 0), std::max(oW, 0)},
                            xc.options());
    }

    auto out = torch::empty({B, 1, oD, oH, oW}, xc.options());

    dim3 grid(oW, oH, oD * B);
    int threads = ((C + 31) / 32) * 32;
    if (threads < 32) threads = 32;
    if (threads > 256) threads = 256;

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_pool_chansum_kernel<<<grid, threads, 0, stream.stream()>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(),
        C, D, H, W, oD, oH, oW);

    return out;
}
"""

fused_pool_chansum_cpp_source = (
    "torch::Tensor fused_pool_chansum_cuda(torch::Tensor x);"
)

fused_pool_chansum = load_inline(
    name="fused_pool_chansum_ext",
    cpp_sources=fused_pool_chansum_cpp_source,
    cuda_sources=fused_pool_chansum_source,
    functions=["fused_pool_chansum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the 3D transposed convolution stays on cuDNN, while the two
    max pooling layers and the channel sum are fused into a single custom CUDA
    kernel (a composed 6x6x6 max pool + channel reduction, read-once).
    """
    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
        )

    def forward(self, x):
        x = self.conv_transpose(x)
        if x.is_cuda and x.dtype == torch.float32:
            return fused_pool_chansum.fused_pool_chansum_cuda(x)
        # Fallback path (e.g. CPU tensors): original ops.
        x = F.max_pool3d(x, kernel_size=2)
        x = F.max_pool3d(x, kernel_size=3)
        return torch.sum(x, dim=1, keepdim=True)
```