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

# Fully fused CUDA kernel that performs the entire pipeline in a single launch:
#   ConvTranspose3d (k=3, s=2, p=1, out_pad=1)
#   + MaxPool3d (k=2, s=2, p=0)
#   + Softmax over channels (dim=1)
#   + per-channel subtract
#   + Swish (x * sigmoid(x))
#   + Max over channels
# This eliminates all intermediate global-memory traffic (the conv-transpose
# output alone is ~1GB, which is never written to DRAM).
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define CIN 3
#define COUT 16
#define KSIZE 3
#define PAD 1
#define NEG_INF -3.402823466e+38f

__global__ void fused_pipeline_kernel(
    const float* __restrict__ x,      // [N, CIN, Din, Hin, Win]
    const float* __restrict__ w,      // [CIN, COUT, 3, 3, 3]
    const float* __restrict__ bias,   // [COUT]
    const float* __restrict__ sub,    // [COUT]
    float* __restrict__ out,          // [N, Dp, Hp, Wp]
    const int N,
    const int Din, const int Hin, const int Win,
    const int Dp, const int Hp, const int Wp)
{
    __shared__ float ws[CIN][COUT][KSIZE][KSIZE][KSIZE];
    __shared__ float bs[COUT];
    __shared__ float ss[COUT];

    const int tid = threadIdx.x;
    float* ws_flat = &ws[0][0][0][0][0];
    for (int i = tid; i < CIN * COUT * KSIZE * KSIZE * KSIZE; i += blockDim.x) {
        ws_flat[i] = w[i];
    }
    if (tid < COUT) {
        bs[tid] = bias[tid];
        ss[tid] = sub[tid];
    }
    __syncthreads();

    const long idx = (long)blockIdx.x * blockDim.x + tid;
    const long total = (long)N * Dp * Hp * Wp;
    if (idx >= total) return;

    const int wp = (int)(idx % Wp);
    const int hp = (int)((idx / Wp) % Hp);
    const int dp = (int)((idx / ((long)Wp * Hp)) % Dp);
    const int n  = (int)(idx / ((long)Wp * Hp * Dp));

    // Max-pooled conv-transpose values, one per output channel.
    float pooled[COUT];
    #pragma unroll
    for (int o = 0; o < COUT; ++o) pooled[o] = NEG_INF;

    // Each pooled element is the max over a 2x2x2 window of conv-transpose outputs.
    #pragma unroll
    for (int dz = 0; dz < 2; ++dz) {
        const int d = 2 * dp + dz;
        #pragma unroll
        for (int dy = 0; dy < 2; ++dy) {
            const int h = 2 * hp + dy;
            #pragma unroll
            for (int dx = 0; dx < 2; ++dx) {
                const int wpos = 2 * wp + dx;

                // ConvTranspose3d output at (d, h, wpos) for all output channels:
                // out[o] = bias[o] + sum_i sum_{kd,kh,kw} w[i][o][kd][kh][kw]
                //          * x[n][i][(d+PAD-kd)/2][(h+PAD-kh)/2][(wpos+PAD-kw)/2]
                // where (pos + PAD - k) must be divisible by stride (2) and in range.
                float acc[COUT];
                #pragma unroll
                for (int o = 0; o < COUT; ++o) acc[o] = 0.0f;

                #pragma unroll
                for (int i = 0; i < CIN; ++i) {
                    const float* xb = x + (long)(n * CIN + i) * Din * Hin * Win;
                    #pragma unroll
                    for (int kd = 0; kd < KSIZE; ++kd) {
                        const int nd = d + PAD - kd;
                        if ((nd & 1) != 0) continue;
                        const int di = nd >> 1;
                        if (di < 0 || di >= Din) continue;
                        #pragma unroll
                        for (int kh = 0; kh < KSIZE; ++kh) {
                            const int nh = h + PAD - kh;
                            if ((nh & 1) != 0) continue;
                            const int hi = nh >> 1;
                            if (hi < 0 || hi >= Hin) continue;
                            #pragma unroll
                            for (int kw = 0; kw < KSIZE; ++kw) {
                                const int nw = wpos + PAD - kw;
                                if ((nw & 1) != 0) continue;
                                const int wi = nw >> 1;
                                if (wi < 0 || wi >= Win) continue;
                                const float xv = xb[(di * Hin + hi) * Win + wi];
                                #pragma unroll
                                for (int o = 0; o < COUT; ++o) {
                                    acc[o] = fmaf(xv, ws[i][o][kd][kh][kw], acc[o]);
                                }
                            }
                        }
                    }
                }

                #pragma unroll
                for (int o = 0; o < COUT; ++o) {
                    pooled[o] = fmaxf(pooled[o], acc[o] + bs[o]);
                }
            }
        }
    }

    // Softmax across channels (numerically stable).
    float mx = pooled[0];
    #pragma unroll
    for (int o = 1; o < COUT; ++o) mx = fmaxf(mx, pooled[o]);
    float sum = 0.0f;
    #pragma unroll
    for (int o = 0; o < COUT; ++o) {
        pooled[o] = expf(pooled[o] - mx);
        sum += pooled[o];
    }
    const float inv_sum = 1.0f / sum;

    // Subtract + Swish + Max across channels.
    float result = NEG_INF;
    #pragma unroll
    for (int o = 0; o < COUT; ++o) {
        const float y = pooled[o] * inv_sum - ss[o];
        const float z = y / (1.0f + expf(-y));
        result = fmaxf(result, z);
    }

    out[idx] = result;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sub) {
    TORCH_CHECK(x.dim() == 5, "x must be 5D");
    TORCH_CHECK(x.size(1) == CIN, "in_channels must be 3");
    TORCH_CHECK(weight.size(0) == CIN && weight.size(1) == COUT, "bad weight shape");

    const int N   = x.size(0);
    const int Din = x.size(2);
    const int Hin = x.size(3);
    const int Win = x.size(4);

    // ConvTranspose3d output dims for k=3, s=2, p=1, output_padding=1.
    const int Dout = (Din - 1) * 2 - 2 * PAD + KSIZE + 1;
    const int Hout = (Hin - 1) * 2 - 2 * PAD + KSIZE + 1;
    const int Wout = (Win - 1) * 2 - 2 * PAD + KSIZE + 1;
    TORCH_CHECK(Dout % 2 == 0 && Hout % 2 == 0 && Wout % 2 == 0, "dims must be even");

    // MaxPool3d with k=2, s=2, p=0 halves the spatial dims.
    const int Dp = Dout / 2;
    const int Hp = Hout / 2;
    const int Wp = Wout / 2;

    auto out = torch::empty({N, Dp, Hp, Wp}, x.options());

    const long total = (long)N * Dp * Hp * Wp;
    if (total == 0) return out;

    const int block = 128;
    const unsigned int grid = (unsigned int)((total + block - 1) / block);

    fused_pipeline_kernel<<<grid, block>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        sub.data_ptr<float>(),
        out.data_ptr<float>(),
        N, Din, Hin, Win, Dp, Hp, Wp);

    return out;
}
"""

_fused_cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sub);"
)

_fused_pipeline = load_inline(
    name="fused_convt_pool_softmax_swish_max",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    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._fast_path = (
            in_channels == 3 and out_channels == 16 and kernel_size == 3
            and stride == 2 and padding == 1 and output_padding == 1
            and pool_kernel_size == 2 and pool_stride == 2 and pool_padding == 0
            and self.conv_transpose.bias is not None
        )

    def forward(self, x):
        if (self._fast_path and x.is_cuda and x.dtype == torch.float32
                and self.conv_transpose.weight.is_cuda):
            return _fused_pipeline.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                self.conv_transpose.bias.contiguous(),
                self.subtract.contiguous(),
            )
        # Generic fallback (identical to the reference implementation).
        x = self.conv_transpose(x)
        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
```