```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 v2:
#   ConvTranspose3d (k=3, s=2, p=1, op=1) -> MaxPool3d (2,2,0) -> Softmax(dim=1)
#   -> per-channel subtract -> Swish -> Max over channels
#
# Improvements over v1:
#  * 4 threads cooperate on each output voxel (one "quad" of consecutive warp
#    lanes). Each thread computes 4 of the 16 output channels, then the
#    softmax max/sum and the final swish-max are combined with __shfl_xor.
#    This gives 4x more warps, ~half the registers per thread, much shorter
#    dependency chains, and better latency hiding (v1 was latency-bound).
#  * The redundant expf in the softmax denominator is removed: the 4 local
#    exponentials are kept in registers and reused for the probabilities.
#  * Channels are assigned strided (co = 4*cc + slice) so the float4 shared
#    -memory weight reads from the 4 lanes of a quad land on distinct banks.
#  * Weights live in shared memory as [ci*CO+co][7] float4 (28 floats with
#    one zero pad), so all weight reads are aligned LDS.128 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 CPT 4
#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 Wd,
    const int Dp, const int Hp, const int Wp, const long long nvox)
{
    __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 float4* Ws4 = reinterpret_cast<const float4*>(Ws);

    const long long tidx = (long long)blockIdx.x * BLOCK + tid;
    const int slice = tid & 3;              // channel slice within quad
    const long long v = tidx >> 2;          // voxel index
    const bool active = (v < nvox);

    // Load the 2x2x2xCI input block this voxel depends on (zero otherwise).
    float xv[CI][2][2][2];
    #pragma unroll
    for (int ci = 0; ci < CI; ++ci)
        #pragma unroll
        for (int a = 0; a < 2; ++a)
            #pragma unroll
            for (int c = 0; c < 2; ++c)
                #pragma unroll
                for (int d = 0; d < 2; ++d) xv[ci][a][c][d] = 0.0f;

    if (active) {
        const int pw = (int)(v % Wp);
        const long long q1 = v / 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 * Wd;
        const long long DHW = (long long)D * HW;
        const float* xb = x + (long long)b * CI * DHW;

        #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;
                        if (ii < D && jj < H && kk < Wd) {
                            xv[ci][id][ih][iw] =
                                __ldg(xb + (long long)ci * DHW
                                      + (long long)ii * HW + (long long)jj * Wd + kk);
                        }
                    }
                }
            }
        }
    }

    // Conv-transpose 2x2x2 block + max-pool for this thread's 4 channels.
    // Per-dim kernel mapping: k=0 -> (delta=1,i=1), k=1 -> (delta=0,i=0),
    //                         k=2 -> (delta=1,i=0)
    float m[CPT];
    #pragma unroll
    for (int cc = 0; cc < CPT; ++cc) {
        const int co = (cc << 2) + slice;
        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* wb = Ws4 + (ci * CO + co) * 7;
            #pragma unroll
            for (int t = 0; t < 7; ++t) {
                const float4 wv = wb[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]);
                    }
                }
            }
        }
        const float bc = bs[co];
        float mm = acc[0] + bc;
        #pragma unroll
        for (int a = 1; a < 8; ++a) mm = fmaxf(mm, acc[a] + bc);
        m[cc] = mm;
    }

    // Quad-reduced numerically-stable softmax over the 16 channels.
    float mx = m[0];
    #pragma unroll
    for (int i = 1; i < CPT; ++i) mx = fmaxf(mx, m[i]);
    mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 1));
    mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 2));

    float e[CPT];
    float sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < CPT; ++i) { e[i] = expf(m[i] - mx); sum += e[i]; }
    sum += __shfl_xor_sync(0xffffffffu, sum, 1);
    sum += __shfl_xor_sync(0xffffffffu, sum, 2);
    const float inv = 1.0f / sum;

    // Subtract, swish, and quad-reduced max over channels.
    float res = -FLT_MAX;
    #pragma unroll
    for (int i = 0; i < CPT; ++i) {
        const float p = e[i] * inv;
        const float y = p - ss[(i << 2) + slice];
        const float z = y / (1.0f + expf(-y));   // swish: y * sigmoid(y)
        res = fmaxf(res, z);
    }
    res = fmaxf(res, __shfl_xor_sync(0xffffffffu, res, 1));
    res = fmaxf(res, __shfl_xor_sync(0xffffffffu, res, 2));

    if (active && slice == 0) out[v] = 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 nvox = (long long)B * Dp * Hp * Wp;
    if (nvox == 0) return out;

    const long long total = nvox * 4;   // 4 threads per voxel
    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, nvox);
    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_v2",
    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. Four threads cooperate per output voxel (4 channels each) with
    warp-shuffle reductions for the softmax statistics and the final max.
    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
```