```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 -> MaxPool3d -> Softmax(dim=1) -> Subtract -> Swish -> Max(dim=1)
#
# For the target configuration (kernel=3, stride=2, pad=1, output_padding=1,
# pool 2/2/0) the transposed convolution output is exactly 2x the input in each
# spatial dim, and the non-overlapping 2x2x2 max-pool brings it back to the input
# spatial size. Hence for every pooled output voxel we can compute the 2x2x2 conv
# outputs on the fly (each is used exactly once, so nothing is wasted) directly
# from a 2x2x2 input tile per input channel, entirely in registers:
#   even output coord  o = 2m   gets input[m]   * w[k=1]
#   odd  output coord  o = 2m+1 gets input[m]   * w[k=2] + input[m+1] * w[k=0]
# The convolution weights are broadcast-identical for every thread, so they are
# kept in __constant__ memory (fast uniform broadcast path). No 268 MB
# intermediate tensor is ever materialized.

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

#define MAX_W_FLOATS 16384   // 64 KB budget for constant-memory weights
#define MAX_C 1024

__constant__ float cW[MAX_W_FLOATS];
__constant__ float cBias[MAX_C];
__constant__ float cSub[MAX_C];

template <int IC, int OC>
__global__ void fused_convt_pool_softmax_swish_cmax_kernel(
    const float* __restrict__ x,   // (B, IC, D, H, W)
    float* __restrict__ out,       // (B, D, H, W)
    int B, int D, int H, int W)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * D * H * W;
    if (idx >= total) return;

    int ow = (int)(idx % W);
    int oy = (int)((idx / W) % H);
    int oz = (int)((idx / ((long long)W * H)) % D);
    int b  = (int)(idx / ((long long)W * H * D));

    const int plane = D * H * W;
    const float* xp = x + (size_t)b * IC * plane + (size_t)oz * (H * W) + (size_t)oy * W + ow;

    bool z1 = (oz + 1) < D;
    bool y1 = (oy + 1) < H;
    bool x1 = (ow + 1) < W;

    // 2x2x2 input tile per input channel (zero-filled outside the boundary,
    // which exactly matches "no contribution" for the transposed convolution).
    float X[IC][2][2][2];
    #pragma unroll
    for (int ic = 0; ic < IC; ++ic) {
        const float* p0 = xp + (size_t)ic * plane;
        X[ic][0][0][0] = p0[0];
        X[ic][0][0][1] = x1 ? p0[1] : 0.0f;
        X[ic][0][1][0] = y1 ? p0[W] : 0.0f;
        X[ic][0][1][1] = (y1 && x1) ? p0[W + 1] : 0.0f;
        const float* p1 = p0 + (H * W);
        X[ic][1][0][0] = z1 ? p1[0] : 0.0f;
        X[ic][1][0][1] = (z1 && x1) ? p1[1] : 0.0f;
        X[ic][1][1][0] = (z1 && y1) ? p1[W] : 0.0f;
        X[ic][1][1][1] = (z1 && y1 && x1) ? p1[W + 1] : 0.0f;
    }

    // For each of the 8 positions of the pool window compute all OC conv outputs
    // and take the per-channel max (fused MaxPool3d).
    float pooled[OC];
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) pooled[oc] = -FLT_MAX;

    #pragma unroll
    for (int ld = 0; ld < 2; ++ld) {
      #pragma unroll
      for (int lh = 0; lh < 2; ++lh) {
        #pragma unroll
        for (int lw = 0; lw < 2; ++lw) {
            float acc[OC];
            #pragma unroll
            for (int oc = 0; oc < OC; ++oc) acc[oc] = cBias[oc];
            #pragma unroll
            for (int ic = 0; ic < IC; ++ic) {
                #pragma unroll
                for (int a = 0; a <= ld; ++a) {
                    const int zoff = ld ? a : 0;
                    const int kd   = ld ? (a ? 0 : 2) : 1;
                    #pragma unroll
                    for (int e = 0; e <= lh; ++e) {
                        const int yoff = lh ? e : 0;
                        const int kh   = lh ? (e ? 0 : 2) : 1;
                        #pragma unroll
                        for (int f = 0; f <= lw; ++f) {
                            const int xoff = lw ? f : 0;
                            const int kw   = lw ? (f ? 0 : 2) : 1;
                            const float xv = X[ic][zoff][yoff][xoff];
                            const int kbase = ic * (OC * 27) + kd * 9 + kh * 3 + kw;
                            #pragma unroll
                            for (int oc = 0; oc < OC; ++oc) {
                                acc[oc] = fmaf(xv, cW[kbase + oc * 27], acc[oc]);
                            }
                        }
                    }
                }
            }
            #pragma unroll
            for (int oc = 0; oc < OC; ++oc) pooled[oc] = fmaxf(pooled[oc], acc[oc]);
        }
      }
    }

    // Softmax across channels (numerically stabilized).
    float mx = pooled[0];
    #pragma unroll
    for (int oc = 1; oc < OC; ++oc) mx = fmaxf(mx, pooled[oc]);
    float sum = 0.0f;
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) {
        pooled[oc] = expf(pooled[oc] - mx);
        sum += pooled[oc];
    }
    const float inv = 1.0f / sum;

    // Subtract, Swish (x * sigmoid(x)), and max across channels.
    float outv = -FLT_MAX;
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) {
        const float p = pooled[oc] * inv - cSub[oc];
        const float sw = p / (1.0f + expf(-p));
        outv = fmaxf(outv, sw);
    }
    out[idx] = outv;
}

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

    const int B  = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int D  = (int)x.size(2);
    const int H  = (int)x.size(3);
    const int W  = (int)x.size(4);
    const int OC = (int)weight.size(1);

    TORCH_CHECK(weight.numel() <= MAX_W_FLOATS, "weight too large for constant memory");
    TORCH_CHECK(OC <= MAX_C, "too many output channels");

    auto stream = at::cuda::getCurrentCUDAStream();

    cudaMemcpyToSymbolAsync(cW, weight.data_ptr<float>(), weight.numel() * sizeof(float),
                            0, cudaMemcpyDeviceToDevice, stream.stream());
    if (bias.defined() && bias.numel() == OC) {
        cudaMemcpyToSymbolAsync(cBias, bias.data_ptr<float>(), OC * sizeof(float),
                                0, cudaMemcpyDeviceToDevice, stream.stream());
    } else {
        static float zeros[MAX_C] = {0.0f};
        cudaMemcpyToSymbolAsync(cBias, zeros, OC * sizeof(float),
                                0, cudaMemcpyHostToDevice, stream.stream());
    }
    cudaMemcpyToSymbolAsync(cSub, sub.data_ptr<float>(), OC * sizeof(float),
                            0, cudaMemcpyDeviceToDevice, stream.stream());

    auto out = torch::empty({B, D, H, W}, x.options());

    const long long total = (long long)B * D * H * W;
    const int block = 256;
    const long long grid = (total + block - 1) / block;

    if (IC == 3 && OC == 16) {
        fused_convt_pool_softmax_swish_cmax_kernel<3, 16><<<(unsigned)grid, block, 0, stream.stream()>>>(
            x.data_ptr<float>(), out.data_ptr<float>(), B, D, H, W);
    } else {
        TORCH_CHECK(false, "unsupported channel configuration");
    }
    return out;
}
"""

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

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


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))

        def t3(v):
            return (v, v, v) if isinstance(v, int) else tuple(v)

        self._fast = (
            t3(kernel_size) == (3, 3, 3)
            and t3(stride) == (2, 2, 2)
            and t3(padding) == (1, 1, 1)
            and t3(output_padding) == (1, 1, 1)
            and t3(pool_kernel_size) == (2, 2, 2)
            and t3(pool_stride) == (2, 2, 2)
            and t3(pool_padding) == (0, 0, 0)
            and in_channels == 3
            and out_channels == 16
        )

    def forward(self, x):
        if (
            self._fast
            and x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
        ):
            x = x.contiguous()
            w = self.conv_transpose.weight.contiguous()
            b = self.conv_transpose.bias
            s = self.subtract.contiguous()
            return fused_ext.fused_forward(x, w, b, s)

        # 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
```