```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:
#   ConvTranspose3d (k=3, s=2, p=1, op=1) -> MaxPool3d (2,2,0) -> Softmax(dim=1)
#   -> per-channel subtract -> Swish -> Max over channels
#
# Key ideas:
#  * With stride=2/pad=1/k=3/op=1, the transposed conv output coordinate is
#    o = 2*i - 1 + k. Hence each output voxel of a 2x2x2 pooling block only
#    depends on a 2x2x2 input voxel block with 27 (input,kernel) pairings per
#    channel (no full 3x3x3 per output voxel is needed).
#  * Everything (conv -> pool -> softmax -> sub -> swish -> channel-max) is
#    fused into a single kernel, so the huge intermediate conv output
#    (B x 16 x 2D x 2H x 2W) is never materialized in global memory.
#  * Weights are staged in shared memory, zero-padded to 28 floats per
#    (in_ch,out_ch) so they can be read with aligned float4 broadcasts.

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>

#define CI 3
#define CO 16
#define KPAD 28
#define BLOCK 256

__global__ __launch_bounds__(BLOCK) void fused_convT_pool_softmax_swish_cmax_kernel(
    const float* __restrict__ x,        // [B, CI, D, H, W]
    const float* __restrict__ weight,   // [CI, CO, 3, 3, 3]
    const float* __restrict__ bias,     // [CO]
    const float* __restrict__ sub,      // [CO]
    float* __restrict__ out,            // [B, Dp, Hp, Wp]
    const int B, const int D, const int H, const int W,
    const int Dp, const int Hp, const int Wp)
{
    __shared__ __align__(16) float Ws[CI * CO * KPAD];
    __shared__ float bs[CO];
    __shared__ float ss[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += BLOCK) {
        const int g = i / 27;
        const int k = i - g * 27;
        Ws[g * KPAD + k] = weight[i];
    }
    if (tid < CI * CO) Ws[tid * KPAD + 27] = 0.0f;
    if (tid < CO) { bs[tid] = bias[tid]; ss[tid] = sub[tid]; }
    __syncthreads();

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

    const int pw = (int)(idx % Wp);
    const long long q1 = idx / Wp;
    const int ph = (int)(q1 % Hp);
    const long long q2 = q1 / Hp;
    const int pd = (int)(q2 % Dp);
    const int b  = (int)(q2 / Dp);

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

    // Load the 2x2x2xCI input block this output voxel depends on.
    float xv[CI][2][2][2];
    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        #pragma unroll
        for (int id = 0; id < 2; ++id) {
            #pragma unroll
            for (int ih = 0; ih < 2; ++ih) {
                #pragma unroll
                for (int iw = 0; iw < 2; ++iw) {
                    const int ii = pd + id;
                    const int jj = ph + ih;
                    const int kk = pw + iw;
                    float v = 0.0f;
                    if (ii < D && jj < H && kk < W) {
                        v = __ldg(x + ((long long)b * CI + ci) * DHW
                                  + (long long)ii * HW + (long long)jj * W + kk);
                    }
                    xv[ci][id][ih][iw] = v;
                }
            }
        }
    }

    // For each output channel: compute the 2x2x2 conv-transpose block and pool it.
    // Per-dim mapping of kernel offset k -> (block offset delta, input offset):
    //   k=0 -> (delta=1, i_off=1), k=1 -> (delta=0, i_off=0), k=2 -> (delta=1, i_off=0)
    float pooled[CO];
    #pragma unroll
    for (int co = 0; co < CO; ++co) {
        float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
        #pragma unroll
        for (int ci = 0; ci < CI; ++ci) {
            const float* wbase = Ws + (ci * CO + co) * KPAD;
            #pragma unroll
            for (int t = 0; t < 7; ++t) {
                const float4 wv = *reinterpret_cast<const float4*>(wbase + (t << 2));
                const float wa[4] = {wv.x, wv.y, wv.z, wv.w};
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    const int k = (t << 2) + j;
                    if (k < 27) {
                        const int kd = k / 9;
                        const int kh = (k / 3) % 3;
                        const int kw = k % 3;
                        const int dd = (kd == 1) ? 0 : 1;
                        const int id = (kd == 0) ? 1 : 0;
                        const int dh = (kh == 1) ? 0 : 1;
                        const int ih = (kh == 0) ? 1 : 0;
                        const int dw = (kw == 1) ? 0 : 1;
                        const int iw = (kw == 0) ? 1 : 0;
                        acc[(dd << 2) | (dh << 1) | dw] =
                            fmaf(xv[ci][id][ih][iw], wa[j],
                                 acc[(dd << 2) | (dh << 1) | dw]);
                    }
                }
            }
        }
        const float bc = bs[co];
        float m = acc[0] + bc;
        #pragma unroll
        for (int a = 1; a < 8; ++a) m = fmaxf(m, acc[a] + bc);
        pooled[co] = m;
    }

    // Softmax across channels (numerically stable), subtract, swish, max.
    float mx = pooled[0];
    #pragma unroll
    for (int c = 1; c < CO; ++c) mx = fmaxf(mx, pooled[c]);

    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < CO; ++c) sum += expf(pooled[c] - mx);
    const float inv = 1.0f / sum;

    float res = -FLT_MAX;
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        const float p = expf(pooled[c] - mx) * inv;
        const float y = p - ss[c];
        const float z = y / (1.0f + expf(-y));   // swish: y * sigmoid(y)
        res = fmaxf(res, z);
    }

    out[idx] = res;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            torch::Tensor sub, int64_t Dp, int64_t Hp, int64_t Wp)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda() && sub.is_cuda(),
                "all tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.size(0) == CI && weight.size(1) == CO && weight.size(2) == 3,
                "unexpected weight shape");

    const int64_t B = x.size(0);
    const int64_t D = x.size(2);
    const int64_t H = x.size(3);
    const int64_t W = x.size(4);

    auto out = at::empty({B, Dp, Hp, Wp}, x.options());
    const long long total = (long long)B * Dp * Hp * Wp;
    if (total == 0) return out;

    const int blocks = (int)((total + BLOCK - 1) / BLOCK);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_convT_pool_softmax_swish_cmax_kernel<<<blocks, BLOCK, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        sub.data_ptr<float>(), out.data_ptr<float>(),
        (int)B, (int)D, (int)H, (int)W, (int)Dp, (int)Hp, (int)Wp);
    return out;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, "
    "torch::Tensor sub, int64_t Dp, int64_t Hp, int64_t Wp);"
)

_fused_ext = load_inline(
    name="fused_convT3d_pool_softmax_sub_swish_cmax",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Uses a single fully-fused CUDA kernel that
    performs ConvTranspose3d + MaxPool3d + Softmax + Subtract + Swish +
    channel-wise Max without materializing any intermediate tensor in global
    memory. Falls back to the original PyTorch ops for unsupported configs.
    """
    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))

        def _i(v):
            return v[0] if isinstance(v, (tuple, list)) else v

        cfg = (in_channels, out_channels, _i(kernel_size), _i(stride), _i(padding),
               _i(output_padding), _i(pool_kernel_size), _i(pool_stride), _i(pool_padding))
        self._fused_ok = (cfg == (3, 16, 3, 2, 1, 1, 2, 2, 0))
        self._zero_bias = None

    def forward(self, x):
        if self._fused_ok and x.is_cuda and x.dtype == torch.float32:
            D, H, W = x.shape[2], x.shape[3], x.shape[4]
            # conv-transpose output size: (n-1)*2 - 2 + 3 + 1 = 2n
            Dc = (D - 1) * 2 - 2 + 3 + 1
            Hc = (H - 1) * 2 - 2 + 3 + 1
            Wc = (W - 1) * 2 - 2 + 3 + 1
            # max-pool (k=2, s=2, p=0) output size
            Dp = (Dc - 2) // 2 + 1
            Hp = (Hc - 2) // 2 + 1
            Wp = (Wc - 2) // 2 + 1
            bias = self.conv_transpose.bias
            if bias is None:
                if self._zero_bias is None or self._zero_bias.device != x.device:
                    self._zero_bias = torch.zeros(16, device=x.device, dtype=torch.float32)
                bias = self._zero_bias
            return _fused_ext.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                bias.contiguous(),
                self.subtract.contiguous(),
                Dp, Hp, Wp)

        # Generic fallback: original PyTorch operator sequence.
        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
```