```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


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

#define IC 3
#define OC 16
#define OC4 4
#define K3 27
#define W4_SIZE (K3 * IC * OC4)
#define W_SIZE (K3 * IC * OC)

__device__ __forceinline__ float4 max4(float4 a, float4 b) {
    return make_float4(
        fmaxf(a.x, b.x),
        fmaxf(a.y, b.y),
        fmaxf(a.z, b.z),
        fmaxf(a.w, b.w)
    );
}

__device__ __forceinline__ float fast_swish(float v) {
    return __fdividef(v, 1.0f + __expf(-v));
}

__global__ void __launch_bounds__(256) fused_convt_pool_softmax_sub_swish_max_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    const float* __restrict__ sub,
    float* __restrict__ out,
    int N, int Din, int Hin, int Win)
{
    // Weights are transposed on the fly from PyTorch ConvTranspose3d layout
    // [ic][oc][kd][kh][kw] into [k][ic][oc], which permits float4 loads over oc.
    __shared__ float4 sw[W4_SIZE];
    __shared__ float4 sb[OC4];
    __shared__ float4 ss[OC4];

    float* sw_scalar = reinterpret_cast<float*>(sw);
    float* sb_scalar = reinterpret_cast<float*>(sb);
    float* ss_scalar = reinterpret_cast<float*>(ss);

    for (int i = threadIdx.x; i < W_SIZE; i += blockDim.x) {
        int ic = i / (OC * K3);
        int rem = i - ic * (OC * K3);
        int oc = rem / K3;
        int koff = rem - oc * K3;
        int dst = (koff * IC + ic) * OC + oc;
        sw_scalar[dst] = __ldg(w + i);
    }

    if (threadIdx.x < OC) {
        sb_scalar[threadIdx.x] = __ldg(bias + threadIdx.x);
        ss_scalar[threadIdx.x] = __ldg(sub + threadIdx.x);
    }
    __syncthreads();

    long spatial = (long)Din * Hin * Win;
    long total = (long)N * spatial;
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = (int)(idx % Win);
    int oh = (int)((idx / Win) % Hin);
    int od = (int)((idx / ((long)Win * Hin)) % Din);
    int n  = (int)(idx / spatial);

    const float* xn = x + (long)n * IC * spatial;

    // The whole 2x2x2 pool window only depends on the 2x2x2 input neighborhood
    // beginning at (od, oh, ow). Load each unique voxel once and reuse it for
    // every applicable transposed-convolution tap.
    bool vd = (od + 1) < Din;
    bool vh = (oh + 1) < Hin;
    bool vw = (ow + 1) < Win;

    float xl[2 * 2 * 2 * IC];

    #pragma unroll
    for (int ad = 0; ad < 2; ++ad) {
        bool okd = (ad == 0) || vd;
        int xd = okd ? (od + ad) : od;
        #pragma unroll
        for (int ah = 0; ah < 2; ++ah) {
            bool okh = (ah == 0) || vh;
            int xh = okh ? (oh + ah) : oh;
            #pragma unroll
            for (int aw = 0; aw < 2; ++aw) {
                bool okw = (aw == 0) || vw;
                int xw = okw ? (ow + aw) : ow;
                bool valid = okd && okh && okw;
                long off = ((long)xd * Hin + xh) * Win + xw;
                const float* xp = xn + off;

                #pragma unroll
                for (int ic = 0; ic < IC; ++ic) {
                    float v = __ldg(xp + (long)ic * spatial);
                    xl[((ad * 2 + ah) * 2 + aw) * IC + ic] = valid ? v : 0.0f;
                }
            }
        }
    }

    float4 pooled[OC4];
    int phase = 0;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int k = 0; k < 2; ++k) {
                float4 acc[OC4];
                #pragma unroll
                for (int g = 0; g < OC4; ++g) acc[g] = sb[g];

                #pragma unroll
                for (int ad = 0; ad < 2; ++ad) {
                    if (i == 0 && ad > 0) continue;
                    int kd = (i == 0) ? 1 : (2 - 2 * ad);

                    #pragma unroll
                    for (int ah = 0; ah < 2; ++ah) {
                        if (j == 0 && ah > 0) continue;
                        int kh = (j == 0) ? 1 : (2 - 2 * ah);

                        #pragma unroll
                        for (int aw = 0; aw < 2; ++aw) {
                            if (k == 0 && aw > 0) continue;
                            int kw = (k == 0) ? 1 : (2 - 2 * aw);
                            int koff = kd * 9 + kh * 3 + kw;
                            int xbase = ((ad * 2 + ah) * 2 + aw) * IC;

                            #pragma unroll
                            for (int ic = 0; ic < IC; ++ic) {
                                float xv = xl[xbase + ic];
                                const float4* wp = sw + (koff * IC + ic) * OC4;

                                #pragma unroll
                                for (int g = 0; g < OC4; ++g) {
                                    float4 wv = wp[g];
                                    acc[g].x = fmaf(xv, wv.x, acc[g].x);
                                    acc[g].y = fmaf(xv, wv.y, acc[g].y);
                                    acc[g].z = fmaf(xv, wv.z, acc[g].z);
                                    acc[g].w = fmaf(xv, wv.w, acc[g].w);
                                }
                            }
                        }
                    }
                }

                if (phase == 0) {
                    #pragma unroll
                    for (int g = 0; g < OC4; ++g) pooled[g] = acc[g];
                } else {
                    #pragma unroll
                    for (int g = 0; g < OC4; ++g) pooled[g] = max4(pooled[g], acc[g]);
                }
                ++phase;
            }
        }
    }

    // Stable softmax across the 16 channels.
    float mx = pooled[0].x;
    mx = fmaxf(mx, pooled[0].y);
    mx = fmaxf(mx, pooled[0].z);
    mx = fmaxf(mx, pooled[0].w);
    #pragma unroll
    for (int g = 1; g < OC4; ++g) {
        mx = fmaxf(mx, pooled[g].x);
        mx = fmaxf(mx, pooled[g].y);
        mx = fmaxf(mx, pooled[g].z);
        mx = fmaxf(mx, pooled[g].w);
    }

    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        pooled[g].x = __expf(pooled[g].x - mx);
        pooled[g].y = __expf(pooled[g].y - mx);
        pooled[g].z = __expf(pooled[g].z - mx);
        pooled[g].w = __expf(pooled[g].w - mx);
    }

    float sum = 0.0f;
    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        sum += pooled[g].x;
        sum += pooled[g].y;
        sum += pooled[g].z;
        sum += pooled[g].w;
    }
    float inv_sum = __fdividef(1.0f, sum);

    // Subtract, Swish, then max across channels.
    float res = -INFINITY;
    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        float v;
        v = pooled[g].x * inv_sum - ss[g].x;
        res = fmaxf(res, fast_swish(v));
        v = pooled[g].y * inv_sum - ss[g].y;
        res = fmaxf(res, fast_swish(v));
        v = pooled[g].z * inv_sum - ss[g].z;
        res = fmaxf(res, fast_swish(v));
        v = pooled[g].w * inv_sum - ss[g].w;
        res = fmaxf(res, fast_swish(v));
    }

    out[idx] = res;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sub) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && bias.is_cuda() && sub.is_cuda(), "All inputs must be CUDA tensors.");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32.");
    TORCH_CHECK(w.scalar_type() == at::kFloat, "weight must be float32.");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32.");
    TORCH_CHECK(sub.scalar_type() == at::kFloat, "subtract must be float32.");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = bias.contiguous();
    auto sc = sub.contiguous();

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

    auto out = torch::empty({N, Din, Hin, Win}, xc.options());
    long total = (long)N * Din * Hin * Win;
    if (total == 0) return out;

    const int block = 256;
    long grid = (total + block - 1) / block;
    auto stream = at::cuda::getCurrentCUDAStream();

    fused_convt_pool_softmax_sub_swish_max_kernel<<<(unsigned int)grid, block, 0, stream>>>(
        xc.data_ptr<float>(),
        wc.data_ptr<float>(),
        bc.data_ptr<float>(),
        sc.data_ptr<float>(),
        out.data_ptr<float>(),
        N, Din, Hin, Win
    );

    return out;
}
"""

cpp_source = "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sub);"

fused_op = load_inline(
    name="fused_convt_pool_softmax_sub_swish_max_v2",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Model with ConvTranspose3d, MaxPool3d, channel softmax,
    subtraction, Swish, and channel max fused into one CUDA kernel for the
    target configuration. Other configurations use the eager PyTorch path.
    """
    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.subtract = nn.Parameter(torch.randn(out_channels))

        self.pool_kernel_size = pool_kernel_size
        self.pool_stride = pool_stride
        self.pool_padding = pool_padding

        self._use_fused = (
            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
        )
        self._fused_forward = fused_op.fused_forward

    def forward(self, x):
        if (
            self._use_fused
            and x.is_cuda
            and x.dtype == torch.float32
            and self.conv_transpose.bias is not None
        ):
            return self._fused_forward(
                x,
                self.conv_transpose.weight,
                self.conv_transpose.bias,
                self.subtract,
            )

        x = self.conv_transpose(x)
        x = F.max_pool3d(
            x,
            kernel_size=self.pool_kernel_size,
            stride=self.pool_stride,
            padding=self.pool_padding,
        )
        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
```