```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 kernel: ConvTranspose3d + MaxPool3d + Softmax(dim=1) + Subtract
# + Swish + Max(dim=1) in a single pass, with no intermediate global tensors.
# Conv-transpose weights (small) live in __constant__ memory; each thread
# computes BPT=2 batch elements at one pooled spatial location so every
# constant-cache weight fetch feeds two FMAs.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math_constants.h>

#define MAX_W_FLOATS 8192
#define MAX_OC_FLOATS 128

__constant__ float c_weight[MAX_W_FLOATS];
__constant__ float c_bias[MAX_OC_FLOATS];
__constant__ float c_sub[MAX_OC_FLOATS];

template <int OC, int BPT>
__global__ void fused_convt_pool_softmax_swish_max_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    const int B, const int IC,
    const int D, const int H, const int W,
    const int OD, const int OH, const int OW,
    const int PD, const int PH, const int PW,
    const int KD, const int KH, const int KW,
    const int SD, const int SH, const int SW,
    const int PAD_D, const int PAD_H, const int PAD_W,
    const int PKD, const int PKH, const int PKW,
    const int PSD, const int PSH, const int PSW,
    const int PPD, const int PPH, const int PPW)
{
    const int nbp = (B + BPT - 1) / BPT;
    const int total = nbp * PD * PH * PW;
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;

    int pw = tid % PW;
    int t = tid / PW;
    int ph = t % PH; t /= PH;
    int pd = t % PD;
    int bp = t / PD;

    int b[BPT];
    #pragma unroll
    for (int i = 0; i < BPT; ++i) b[i] = bp * BPT + i;

    float biasReg[OC];
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) biasReg[oc] = c_bias[oc];

    float pooled[BPT][OC];
    #pragma unroll
    for (int i = 0; i < BPT; ++i)
        #pragma unroll
        for (int oc = 0; oc < OC; ++oc) pooled[i][oc] = -CUDART_INF_F;

    const int KHW = KH * KW;
    const int K3  = KD * KHW;

    // ---- ConvTranspose3d (gather form) fused with MaxPool3d ----
    for (int tz = 0; tz < PKD; ++tz) {
        const int od = pd * PSD + tz - PPD;
        if (od < 0 || od >= OD) continue;
        for (int ty = 0; ty < PKH; ++ty) {
            const int oh = ph * PSH + ty - PPH;
            if (oh < 0 || oh >= OH) continue;
            for (int tx = 0; tx < PKW; ++tx) {
                const int ow = pw * PSW + tx - PPW;
                if (ow < 0 || ow >= OW) continue;

                // one conv-transpose output voxel (all OC channels)
                float acc[BPT][OC];
                #pragma unroll
                for (int i = 0; i < BPT; ++i)
                    #pragma unroll
                    for (int oc = 0; oc < OC; ++oc) acc[i][oc] = biasReg[oc];

                for (int ic = 0; ic < IC; ++ic) {
                    const int icw = ic * OC * K3;
                    const int xic = ic * D * H * W;
                    for (int kd = 0; kd < KD; ++kd) {
                        const int tdd = od + PAD_D - kd;
                        if (tdd < 0 || (tdd % SD) != 0) continue;
                        const int id = tdd / SD;
                        if (id >= D) continue;
                        const int wr_d = icw + kd * KHW;
                        const int x_d = xic + id * H * W;
                        for (int kh = 0; kh < KH; ++kh) {
                            const int thh = oh + PAD_H - kh;
                            if (thh < 0 || (thh % SH) != 0) continue;
                            const int ih = thh / SH;
                            if (ih >= H) continue;
                            const int wr_h = wr_d + kh * KW;
                            const int x_h = x_d + ih * W;
                            for (int kw = 0; kw < KW; ++kw) {
                                const int tww = ow + PAD_W - kw;
                                if (tww < 0 || (tww % SW) != 0) continue;
                                const int iw = tww / SW;
                                if (iw >= W) continue;

                                const int xoff = x_h + iw;
                                float val[BPT];
                                #pragma unroll
                                for (int i = 0; i < BPT; ++i) {
                                    val[i] = (b[i] < B)
                                        ? __ldg(&x[b[i] * IC * D * H * W + xoff])
                                        : 0.0f;
                                }
                                const int wbase = wr_h + kw;
                                #pragma unroll
                                for (int oc = 0; oc < OC; ++oc) {
                                    const float wv = c_weight[wbase + oc * K3];
                                    #pragma unroll
                                    for (int i = 0; i < BPT; ++i)
                                        acc[i][oc] = fmaf(val[i], wv, acc[i][oc]);
                                }
                            }
                        }
                    }
                }

                #pragma unroll
                for (int i = 0; i < BPT; ++i)
                    #pragma unroll
                    for (int oc = 0; oc < OC; ++oc)
                        pooled[i][oc] = fmaxf(pooled[i][oc], acc[i][oc]);
            }
        }
    }

    // ---- Softmax(dim=1) + Subtract + Swish + Max(dim=1) ----
    #pragma unroll
    for (int i = 0; i < BPT; ++i) {
        if (b[i] >= B) continue;
        float mx = pooled[i][0];
        #pragma unroll
        for (int oc = 1; oc < OC; ++oc) mx = fmaxf(mx, pooled[i][oc]);
        float sum = 0.0f;
        #pragma unroll
        for (int oc = 0; oc < OC; ++oc) {
            pooled[i][oc] = expf(pooled[i][oc] - mx);
            sum += pooled[i][oc];
        }
        const float inv = 1.0f / sum;
        float res = -CUDART_INF_F;
        #pragma unroll
        for (int oc = 0; oc < OC; ++oc) {
            const float y  = pooled[i][oc] * inv - c_sub[oc];
            const float sg = 1.0f / (1.0f + expf(-y));
            res = fmaxf(res, sg * y);
        }
        out[((b[i] * PD + pd) * PH + ph) * PW + pw] = res;
    }
}

torch::Tensor fused_forward(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sub,
    int64_t sd, int64_t sh, int64_t sw,
    int64_t pad_d, int64_t pad_h, int64_t pad_w,
    int64_t opd, int64_t oph, int64_t opw,
    int64_t pkd, int64_t pkh, int64_t pkw,
    int64_t psd, int64_t psh, int64_t psw,
    int64_t ppd, int64_t pph, int64_t ppw)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "fused_forward expects a CUDA float32 tensor");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();
    auto sc = sub.contiguous();

    const int B  = (int)xc.size(0);
    const int IC = (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 OC = (int)wc.size(1);
    const int KD = (int)wc.size(2);
    const int KH = (int)wc.size(3);
    const int KW = (int)wc.size(4);

    const int OD = (D - 1) * (int)sd - 2 * (int)pad_d + (KD - 1) + (int)opd + 1;
    const int OH = (H - 1) * (int)sh - 2 * (int)pad_h + (KH - 1) + (int)oph + 1;
    const int OW = (W - 1) * (int)sw - 2 * (int)pad_w + (KW - 1) + (int)opw + 1;
    const int PD = (OD + 2 * (int)ppd - (int)pkd) / (int)psd + 1;
    const int PH = (OH + 2 * (int)pph - (int)pkh) / (int)psh + 1;
    const int PW = (OW + 2 * (int)ppw - (int)pkw) / (int)psw + 1;

    TORCH_CHECK(OD > 0 && OH > 0 && OW > 0 && PD > 0 && PH > 0 && PW > 0,
                "bad output geometry");
    TORCH_CHECK(wc.numel() <= MAX_W_FLOATS, "weights exceed constant-memory capacity");
    TORCH_CHECK(OC <= MAX_OC_FLOATS, "channels exceed constant-memory capacity");
    TORCH_CHECK(OC == 16, "fused kernel instantiated for out_channels == 16");

    auto out = torch::empty({B, PD, PH, PW}, xc.options());
    if (out.numel() == 0) return out;

    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemcpyToSymbolAsync(c_weight, wc.data_ptr<float>(),
                            wc.numel() * sizeof(float), 0,
                            cudaMemcpyDeviceToDevice, stream.stream());
    cudaMemcpyToSymbolAsync(c_bias, bc.data_ptr<float>(),
                            OC * sizeof(float), 0,
                            cudaMemcpyDeviceToDevice, stream.stream());
    cudaMemcpyToSymbolAsync(c_sub, sc.data_ptr<float>(),
                            OC * sizeof(float), 0,
                            cudaMemcpyDeviceToDevice, stream.stream());

    const int nbp = (B + 1) / 2;
    const int total = nbp * PD * PH * PW;
    const int block = 256;
    const int grid = (total + block - 1) / block;

    fused_convt_pool_softmax_swish_max_kernel<16, 2><<<grid, block, 0, stream.stream()>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(),
        B, IC, D, H, W,
        OD, OH, OW, PD, PH, PW,
        KD, KH, KW,
        (int)sd, (int)sh, (int)sw,
        (int)pad_d, (int)pad_h, (int)pad_w,
        (int)pkd, (int)pkh, (int)pkw,
        (int)psd, (int)psh, (int)psw,
        (int)ppd, (int)pph, (int)ppw);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward("
    "torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sub,"
    "int64_t sd, int64_t sh, int64_t sw,"
    "int64_t pad_d, int64_t pad_h, int64_t pad_w,"
    "int64_t opd, int64_t oph, int64_t opw,"
    "int64_t pkd, int64_t pkh, int64_t pkw,"
    "int64_t psd, int64_t psh, int64_t psw,"
    "int64_t ppd, int64_t pph, int64_t ppw);"
)

_ext = load_inline(
    name="fused_convt3d_pool_softmax_swish_max",
    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(t) for t in v)
    return (int(v), int(v), int(v))


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._s = _triple(self.conv_transpose.stride)
        self._p = _triple(self.conv_transpose.padding)
        self._op = _triple(self.conv_transpose.output_padding)
        self._dil = _triple(self.conv_transpose.dilation)
        self._pk = _triple(pool_kernel_size)
        self._ps = _triple(pool_stride if pool_stride is not None else pool_kernel_size)
        self._pp = _triple(pool_padding)

    def forward(self, x):
        w = self.conv_transpose.weight
        if (x.is_cuda and x.dtype == torch.float32
                and self._dil == (1, 1, 1)
                and w.size(1) == 16
                and w.numel() <= 8192):
            return _ext.fused_forward(
                x, w, self.conv_transpose.bias, self.subtract,
                self._s[0], self._s[1], self._s[2],
                self._p[0], self._p[1], self._p[2],
                self._op[0], self._op[1], self._op[2],
                self._pk[0], self._pk[1], self._pk[2],
                self._ps[0], self._ps[1], self._ps[2],
                self._pp[0], self._pp[1], self._pp[2])
        # generic fallback (identical semantics to the reference Model)
        y = self.conv_transpose(x)
        y = self.max_pool(y)
        y = torch.softmax(y, dim=1)
        y = y - self.subtract.view(1, -1, 1, 1, 1)
        y = torch.sigmoid(y) * y
        return torch.max(y, dim=1)[0]
```