```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.nn.modules.utils import _triple
from torch.utils.cpp_extension import load_inline

# Fused CUDA kernel implementing:
#   MaxPool3d -> Softmax(dim=1) -> Subtract(per-channel) -> Swish -> Max(dim=1)
# in a single pass over the ConvTranspose3d output.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define MAX_C 128

// Fast path: 2x2x2 max pool, stride 2, no padding, even spatial dims, compile-time channels.
template <int C>
__global__ void fused_pool2_softmax_swish_max_kernel(
    const float* __restrict__ x,    // [N, C, D, H, W]
    const float* __restrict__ sub,  // [C]
    float* __restrict__ out,        // [N, D/2, H/2, W/2]
    const int N,
    const int D, const int H, const int W,
    const int D2, const int H2, const int W2)
{
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    const long long total = (long long)N * D2 * H2 * W2;
    if (idx >= total) return;

    const int w2 = (int)(idx % W2);
    const long long q1 = idx / W2;
    const int h2 = (int)(q1 % H2);
    const long long q2 = q1 / H2;
    const int d2 = (int)(q2 % D2);
    const int n  = (int)(q2 / D2);

    const long long HW = (long long)H * W;
    const long long row = (long long)(d2 * 2) * HW + (long long)(h2 * 2) * W + (w2 * 2);

    float vals[C];
    float mx = -INFINITY;

    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* base = x + ((long long)n * C + c) * D * HW + row;
        const float2 v00 = *reinterpret_cast<const float2*>(base);
        const float2 v01 = *reinterpret_cast<const float2*>(base + W);
        const float2 v10 = *reinterpret_cast<const float2*>(base + HW);
        const float2 v11 = *reinterpret_cast<const float2*>(base + HW + W);
        const float m0 = fmaxf(fmaxf(v00.x, v00.y), fmaxf(v01.x, v01.y));
        const float m1 = fmaxf(fmaxf(v10.x, v10.y), fmaxf(v11.x, v11.y));
        const float m = fmaxf(m0, m1);
        vals[c] = m;
        mx = fmaxf(mx, m);
    }

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

    float best = -INFINITY;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float p = vals[c] / sum;
        const float y = p - sub[c];
        const float sw = y / (1.0f + expf(-y));  // Swish: x * sigmoid(x)
        best = fmaxf(best, sw);
    }
    out[idx] = best;
}

// Generic path: arbitrary pooling params, runtime channel count.
__global__ void fused_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ sub,
    float* __restrict__ out,
    const int N, const int C,
    const int D, const int H, const int W,
    const int D2, const int H2, const int W2,
    const int kD, const int kH, const int kW,
    const int sD, const int sH, const int sW,
    const int pD, const int pH, const int pW)
{
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    const long long total = (long long)N * D2 * H2 * W2;
    if (idx >= total) return;

    const int w2 = (int)(idx % W2);
    const long long q1 = idx / W2;
    const int h2 = (int)(q1 % H2);
    const long long q2 = q1 / H2;
    const int d2 = (int)(q2 % D2);
    const int n  = (int)(q2 / D2);

    const int d_start = d2 * sD - pD;
    const int h_start = h2 * sH - pH;
    const int w_start = w2 * sW - pW;

    const long long HW = (long long)H * W;

    float vals[MAX_C];
    float mx = -INFINITY;

    for (int c = 0; c < C; ++c) {
        const float* base = x + ((long long)n * C + c) * D * HW;
        float m = -INFINITY;
        for (int kd = 0; kd < kD; ++kd) {
            const int d = d_start + kd;
            if (d < 0 || d >= D) continue;
            for (int kh = 0; kh < kH; ++kh) {
                const int h = h_start + kh;
                if (h < 0 || h >= H) continue;
                const float* rowp = base + (long long)d * HW + (long long)h * W;
                for (int kw = 0; kw < kW; ++kw) {
                    const int w = w_start + kw;
                    if (w < 0 || w >= W) continue;
                    m = fmaxf(m, rowp[w]);
                }
            }
        }
        vals[c] = m;
        mx = fmaxf(mx, m);
    }

    float sum = 0.0f;
    for (int c = 0; c < C; ++c) {
        const float e = expf(vals[c] - mx);
        vals[c] = e;
        sum += e;
    }

    float best = -INFINITY;
    for (int c = 0; c < C; ++c) {
        const float p = vals[c] / sum;
        const float y = p - sub[c];
        const float sw = y / (1.0f + expf(-y));
        best = fmaxf(best, sw);
    }
    out[idx] = best;
}

torch::Tensor fused_postop_cuda(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() && sub.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && sub.scalar_type() == at::kFloat, "float32 only");
    TORCH_CHECK(x.dim() == 5, "x must be 5D (N, C, D, H, W)");

    auto xc = x.contiguous();
    auto sc = 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);

    TORCH_CHECK(sc.numel() == C, "subtract must have one value per channel");

    const int D2 = (int)((D + 2 * pD - kD) / sD + 1);
    const int H2 = (int)((H + 2 * pH - kH) / sH + 1);
    const int W2 = (int)((W + 2 * pW - kW) / sW + 1);
    TORCH_CHECK(D2 > 0 && H2 > 0 && W2 > 0, "invalid pooled output shape");

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

    const int threads = 256;
    const unsigned int blocks = (unsigned int)((total + threads - 1) / threads);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float* xp = xc.data_ptr<float>();
    const float* sp = sc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    const bool fast = (kD == 2 && kH == 2 && kW == 2 &&
                       sD == 2 && sH == 2 && sW == 2 &&
                       pD == 0 && pH == 0 && pW == 0 &&
                       (D % 2 == 0) && (H % 2 == 0) && (W % 2 == 0));
    if (fast) {
        switch (C) {
            case 8:
                fused_pool2_softmax_swish_max_kernel<8><<<blocks, threads, 0, stream>>>(
                    xp, sp, op, N, D, H, W, D2, H2, W2);
                return out;
            case 16:
                fused_pool2_softmax_swish_max_kernel<16><<<blocks, threads, 0, stream>>>(
                    xp, sp, op, N, D, H, W, D2, H2, W2);
                return out;
            case 32:
                fused_pool2_softmax_swish_max_kernel<32><<<blocks, threads, 0, stream>>>(
                    xp, sp, op, N, D, H, W, D2, H2, W2);
                return out;
            default:
                break;
        }
    }

    TORCH_CHECK(C <= MAX_C, "channel count exceeds MAX_C");
    fused_generic_kernel<<<blocks, threads, 0, stream>>>(
        xp, sp, op, N, C, D, H, W, D2, H2, W2,
        (int)kD, (int)kH, (int)kW,
        (int)sD, (int)sH, (int)sW,
        (int)pD, (int)pH, (int)pW);
    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_postop_cuda(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);"
)

fused_postop = load_inline(
    name="fused_postop_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_postop_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
        - ConvTranspose3d (cuDNN via F.conv_transpose3d)
        - Fused custom CUDA kernel for:
            MaxPool3d -> Softmax(dim=1) -> Subtract -> Swish -> Max(dim=1)
    """
    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.out_channels = out_channels
        self.pool_kernel_size = _triple(pool_kernel_size)
        self.pool_stride = _triple(pool_stride)
        self.pool_padding = _triple(pool_padding)

    def forward(self, x):
        x = F.conv_transpose3d(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            self.conv_transpose.stride,
            self.conv_transpose.padding,
            self.conv_transpose.output_padding,
            self.conv_transpose.groups,
            self.conv_transpose.dilation,
        )
        if x.is_cuda and x.dtype == torch.float32 and self.out_channels <= 128:
            k = self.pool_kernel_size
            s = self.pool_stride
            p = self.pool_padding
            return fused_postop.fused_postop_cuda(
                x, self.subtract,
                k[0], k[1], k[2],
                s[0], s[1], s[2],
                p[0], p[1], p[2],
            )
        # Fallback to the original sequence of PyTorch ops.
        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
```