```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 v3:
#   ConvTranspose3d (k=3, s=2, p=1, op=1) -> MaxPool3d (2,2,0) -> Softmax(dim=1)
#   -> per-channel subtract -> Swish -> Max over channels
#
# Optimizations over the previous fused version:
#  * Division-free thread mapping: a 3D grid gives (b, pd, ph) directly from
#    block indices and (pw) from the thread x-index, so the kernel contains no
#    integer div/mod at all and uses 32-bit addressing throughout.
#  * The softmax + swish epilogue no longer dominates: exponentials use the
#    __expf intrinsic (MUFU-based, ~1e-6 relative error, far inside the
#    tolerance that already permits TF32), reciprocals/divisions use
#    __fdividef, and the per-channel exponentials computed for the softmax
#    denominator are kept in registers and reused for the probabilities
#    (the reference recomputed them, doubling the transcendental count).
#  * Bias is folded after the pool max (bc + max8 instead of 8 adds + max),
#    keeping the FP32 pipe (the binding resource) as lean as possible.

_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

__global__ __launch_bounds__(256) 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 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.y * 32 + threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += 256) {
        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 int pw = blockIdx.x * 32 + threadIdx.x;
    const int ph = blockIdx.y * 8  + threadIdx.y;
    if (pw >= Wp || ph >= Hp) return;
    const int pd = blockIdx.z % Dp;
    const int b  = blockIdx.z / Dp;

    const int HW  = H * W;
    const int DHW = D * HW;
    const float* xb = x + b * CI * DHW;

    // 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(xb + ci * DHW + ii * HW + 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)
    const float4* Ws4 = reinterpret_cast<const float4*>(Ws);
    float e[CO];
    float mx = -FLT_MAX;
    #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 float4* wbase = Ws4 + (ci * CO + co) * 7;
            #pragma unroll
            for (int t = 0; t < 7; ++t) {
                const float4 wv = wbase[t];
                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]);
                    }
                }
            }
        }
        float m = acc[0];
        #pragma unroll
        for (int a = 1; a < 8; ++a) m = fmaxf(m, acc[a]);
        m += bs[co];
        e[co] = m;              // temporarily holds the pooled logit
        mx = fmaxf(mx, m);
    }

    // Softmax across channels: reuse the exponentials for the probabilities.
    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        e[c] = __expf(e[c] - mx);
        sum += e[c];
    }
    const float inv = __fdividef(1.0f, sum);

    // Subtract, swish (y * sigmoid(y)) and max over channels.
    float res = -FLT_MAX;
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        const float y = fmaf(e[c], inv, -ss[c]);
        const float z = __fdividef(y, 1.0f + __expf(-y));
        res = fmaxf(res, z);
    }

    out[(b * Dp + pd) * (Hp * Wp) + ph * Wp + pw] = 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 dim3 block(32, 8, 1);
    const dim3 grid((unsigned)((Wp + 31) / 32),
                    (unsigned)((Hp + 7) / 8),
                    (unsigned)(B * Dp));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_convT_pool_softmax_swish_cmax_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        sub.data_ptr<float>(), out.data_ptr<float>(),
        (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_v3",
    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:
            B, D, H, W = x.shape[0], 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
            # 32-bit kernel indexing + 3D grid limits
            if (x.numel() < 2**31 and B * Dp <= 65535
                    and B * D * H * W < 2**31 and Dp * Hp * Wp * B < 2**31):
                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
```