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

torch.backends.cudnn.benchmark = True

# Fused CUDA implementation of: MaxPool3d -> Softmax(dim=1) -> Subtract -> Swish -> Max(dim=1)
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math_constants.h>

#define BLOCK_SIZE 256

// Fast path: 2x2x2 pool, stride 2, no padding, W even (float2 loads), C templated (1..32).
// One thread per output voxel; loops over channels, fusing:
// maxpool(2x2x2) -> softmax over channels -> subtract -> swish -> max over channels.
template <int C>
__global__ void fused_pool2_softmax_sub_swish_cmax_kernel(
    const float* __restrict__ x, const float* __restrict__ sub,
    float* __restrict__ out, int N, int D, int H, int W, long total) {
    const int PD = D >> 1;
    const int PH = H >> 1;
    const int PW = W >> 1;
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int pw = (int)(idx % PW);
    long t = idx / PW;
    int ph = (int)(t % PH); t /= PH;
    int pd = (int)(t % PD); t /= PD;
    int n  = (int)t;

    const size_t HW = (size_t)H * (size_t)W;
    const float* base = x + (size_t)n * (size_t)C * (size_t)D * HW
                      + (size_t)(2 * pd) * HW + (size_t)(2 * ph) * (size_t)W + (size_t)(2 * pw);

    float e[C];
    float maxv = -CUDART_INF_F;
#pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* p = base + (size_t)c * (size_t)D * HW;
        float2 r0 = *reinterpret_cast<const float2*>(p);
        float2 r1 = *reinterpret_cast<const float2*>(p + W);
        float2 r2 = *reinterpret_cast<const float2*>(p + HW);
        float2 r3 = *reinterpret_cast<const float2*>(p + HW + W);
        float m = fmaxf(fmaxf(fmaxf(r0.x, r0.y), fmaxf(r1.x, r1.y)),
                        fmaxf(fmaxf(r2.x, r2.y), fmaxf(r3.x, r3.y)));
        e[c] = m;
        maxv = fmaxf(maxv, m);
    }

    float sum = 0.0f;
#pragma unroll
    for (int c = 0; c < C; ++c) {
        e[c] = expf(e[c] - maxv);
        sum += e[c];
    }
    const float inv_sum = 1.0f / sum;

    float res = -CUDART_INF_F;
#pragma unroll
    for (int c = 0; c < C; ++c) {
        float s = e[c] * inv_sum - sub[c];      // softmax -> subtract
        float v = s / (1.0f + expf(-s));        // swish: x * sigmoid(x)
        res = fmaxf(res, v);                    // max over channels
    }
    out[idx] = res;
}

// Generic fallback kernels (any pool config / any channel count).
__global__ void maxpool3d_kernel(
    const float* __restrict__ x, float* __restrict__ pooled,
    int N, int C, int D, int H, int W,
    int PD, int PH, int PW,
    int kD, int kH, int kW,
    int sD, int sH, int sW,
    int pD, int pH, int pW, long total) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int pw = (int)(idx % PW);
    long t = idx / PW;
    int ph = (int)(t % PH); t /= PH;
    int pd = (int)(t % PD); t /= PD;
    int c  = (int)(t % C);  t /= C;
    int n  = (int)t;

    const float* base = x + (((size_t)n * C + c) * (size_t)D) * (size_t)H * (size_t)W;
    float m = -CUDART_INF_F;
    for (int kd = 0; kd < kD; ++kd) {
        int d = pd * sD + kd - pD;
        if (d < 0 || d >= D) continue;
        for (int kh = 0; kh < kH; ++kh) {
            int h = ph * sH + kh - pH;
            if (h < 0 || h >= H) continue;
            const float* row = base + ((size_t)d * H + h) * (size_t)W;
            int w0 = pw * sW - pW;
            for (int kw = 0; kw < kW; ++kw) {
                int w = w0 + kw;
                if (w < 0 || w >= W) continue;
                m = fmaxf(m, row[w]);
            }
        }
    }
    pooled[idx] = m;
}

__global__ void softmax_sub_swish_cmax_kernel(
    const float* __restrict__ pooled, const float* __restrict__ sub,
    float* __restrict__ out, int C, long spatial, long total) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    long n = idx / spatial;
    long s0 = idx - n * spatial;
    const float* p = pooled + ((size_t)n * C) * (size_t)spatial + s0;

    float maxv = -CUDART_INF_F;
    for (int c = 0; c < C; ++c) maxv = fmaxf(maxv, p[(size_t)c * spatial]);
    float sum = 0.0f;
    for (int c = 0; c < C; ++c) sum += expf(p[(size_t)c * spatial] - maxv);
    const float inv_sum = 1.0f / sum;

    float res = -CUDART_INF_F;
    for (int c = 0; c < C; ++c) {
        float s = expf(p[(size_t)c * spatial] - maxv) * inv_sum - sub[c];
        float v = s / (1.0f + expf(-s));
        res = fmaxf(res, v);
    }
    out[idx] = res;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor sub,
                            int64_t kD, int64_t kH, int64_t kW,
                            int64_t sD, int64_t sH, int64_t sW,
                            int64_t pD, int64_t pH, int64_t pW) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    auto subc = sub.contiguous();

    const int N = (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);

    const int PD = (int)((D + 2 * pD - kD) / sD + 1);
    const int PH = (int)((H + 2 * pH - kH) / sH + 1);
    const int PW = (int)((W + 2 * pW - kW) / sW + 1);

    auto out = at::empty({N, PD, PH, PW}, xc.options());
    const long total = (long)N * PD * PH * PW;
    if (total == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const bool fast_shape = (kD == 2 && kH == 2 && kW == 2 &&
                             sD == 2 && sH == 2 && sW == 2 &&
                             pD == 0 && pH == 0 && pW == 0 && (W % 2 == 0));

    if (fast_shape && C >= 1 && C <= 32) {
        const long blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
        #define CASE_C(CV) \
            case CV: \
                fused_pool2_softmax_sub_swish_cmax_kernel<CV><<<blocks, BLOCK_SIZE, 0, stream>>>( \
                    xc.data_ptr<float>(), subc.data_ptr<float>(), out.data_ptr<float>(), \
                    N, D, H, W, total); \
                break;
        switch (C) {
            CASE_C(1)  CASE_C(2)  CASE_C(3)  CASE_C(4)  CASE_C(5)  CASE_C(6)  CASE_C(7)  CASE_C(8)
            CASE_C(9)  CASE_C(10) CASE_C(11) CASE_C(12) CASE_C(13) CASE_C(14) CASE_C(15) CASE_C(16)
            CASE_C(17) CASE_C(18) CASE_C(19) CASE_C(20) CASE_C(21) CASE_C(22) CASE_C(23) CASE_C(24)
            CASE_C(25) CASE_C(26) CASE_C(27) CASE_C(28) CASE_C(29) CASE_C(30) CASE_C(31) CASE_C(32)
            default: break;
        }
        #undef CASE_C
        return out;
    }

    // Generic two-kernel fallback.
    auto pooled = at::empty({N, C, PD, PH, PW}, xc.options());
    const long ptotal = (long)N * C * PD * PH * PW;
    maxpool3d_kernel<<<(ptotal + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), pooled.data_ptr<float>(),
        N, C, D, H, W, PD, PH, PW,
        (int)kD, (int)kH, (int)kW, (int)sD, (int)sH, (int)sW, (int)pD, (int)pH, (int)pW,
        ptotal);
    const long spatial = (long)PD * PH * PW;
    softmax_sub_swish_cmax_kernel<<<(total + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(
        pooled.data_ptr<float>(), subc.data_ptr<float>(), out.data_ptr<float>(),
        C, spatial, total);
    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor sub, "
    "int64_t kD, int64_t kH, int64_t kW, int64_t sD, int64_t sH, int64_t sW, "
    "int64_t pD, int64_t pH, int64_t pW);"
)

_ext = load_inline(
    name="fused_pool_softmax_swish_cmax",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


def _triple(v):
    if isinstance(v, (tuple, list)):
        return tuple(int(x) for x in v)
    return (int(v), int(v), int(v))


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The ConvTranspose3d is kept as-is (cuDNN), while the entire tail of the
    network -- MaxPool3d -> Softmax(dim=1) -> Subtract -> Swish -> Max(dim=1) --
    is executed by a single fused custom CUDA kernel (with a generic CUDA
    fallback for non-standard pool configurations).
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding,
                 output_padding, pool_kernel_size, pool_stride, pool_padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, output_padding=output_padding)
        self.max_pool = nn.MaxPool3d(
            kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding)
        self.subtract = nn.Parameter(torch.randn(out_channels))

        self.pool_k = _triple(pool_kernel_size)
        self.pool_s = _triple(pool_stride)
        self.pool_p = _triple(pool_padding)

    def forward(self, x):
        x = self.conv_transpose(x)
        if x.is_cuda and x.dtype == torch.float32:
            return _ext.fused_forward(
                x, self.subtract,
                self.pool_k[0], self.pool_k[1], self.pool_k[2],
                self.pool_s[0], self.pool_s[1], self.pool_s[2],
                self.pool_p[0], self.pool_p[1], self.pool_p[2])
        # Reference fallback for non-CUDA / non-float32 inputs.
        x = self.max_pool(x)
        x = torch.softmax(x, dim=1)
        x = x - self.subtract.view(1, -1, 1, 1, 1)
        x = torch.sigmoid(x) * x
        x = torch.max(x, dim=1)[0]
        return x
```