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


# ---------------------------------------------------------------------------
# Fully fused CUDA kernel:
#   ConvTranspose3d -> MaxPool3d -> Softmax(dim=1) -> Subtract -> Swish
#   -> Max(dim=1)
#
# Key idea: with pool_stride == conv stride, each pooled output element only
# needs a small (BD x BH x BW) block of input voxels (per input channel).
# Each thread computes one final output element: it loads its input block
# into registers, evaluates every conv-transpose output inside its pooling
# window on the fly (weights staged in shared memory), takes the window max,
# then performs an online softmax / subtract / swish / channel-max entirely
# in registers.  The huge intermediate conv-transpose tensor is never
# materialized in global memory.
# ---------------------------------------------------------------------------

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

#define CI @CI@
#define CO @CO@
#define KD @KD@
#define KH @KH@
#define KW @KW@
#define PKD @PKD@
#define PKH @PKH@
#define PKW @PKW@
#define BD @BD@
#define BH @BH@
#define BW @BW@

// IREL_D[dd][k]: input offset (relative to the pooled index) of the input
// voxel that kernel tap k contributes to for pool-window offset dd, or -1.
__constant__ int IREL_D[PKD][KD] = @TAB_D@;
__constant__ int IREL_H[PKH][KH] = @TAB_H@;
__constant__ int IREL_W[PKW][KW] = @TAB_W@;

__global__ void fused_ct3d_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,
    const int B, const int D, const int H, const int W,
    const int PD, const int PH, const int PW)
{
    extern __shared__ float smem[];
    float* sw = smem;                          // CI*CO*KD*KH*KW weights
    float* sb = sw + CI * CO * KD * KH * KW;   // CO bias
    float* ss = sb + CO;                       // CO subtract

    for (int t = threadIdx.x; t < CI * CO * KD * KH * KW; t += blockDim.x) {
        sw[t] = w[t];
    }
    for (int t = threadIdx.x; t < CO; t += blockDim.x) {
        sb[t] = bias[t];
        ss[t] = sub[t];
    }
    __syncthreads();

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = B * PD * PH * PW;
    if (idx >= total) return;

    const int pw = idx % PW;
    const int ph = (idx / PW) % PH;
    const int pd = (idx / (PW * PH)) % PD;
    const int b  = idx / (PW * PH * PD);

    // Load the input block this thread needs (zero padded at borders).
    float vin[CI][BD][BH][BW];
    #pragma unroll
    for (int ic = 0; ic < CI; ++ic) {
        #pragma unroll
        for (int di = 0; di < BD; ++di) {
            const int ii = pd + di;
            #pragma unroll
            for (int dj = 0; dj < BH; ++dj) {
                const int jj = ph + dj;
                #pragma unroll
                for (int dl = 0; dl < BW; ++dl) {
                    const int ll = pw + dl;
                    float v = 0.f;
                    if (ii < D && jj < H && ll < W) {
                        v = x[(((int64_t)b * CI + ic) * D + ii) * H * W
                              + (int64_t)jj * W + ll];
                    }
                    vin[ic][di][dj][dl] = v;
                }
            }
        }
    }

    // For each output channel: evaluate conv-transpose at every position of
    // the pooling window, then take the max (fused MaxPool3d).
    float pvals[CO];
    float pmax = -FLT_MAX;

    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        float pm = -FLT_MAX;
        #pragma unroll
        for (int dd = 0; dd < PKD; ++dd) {
            #pragma unroll
            for (int dh = 0; dh < PKH; ++dh) {
                #pragma unroll
                for (int dw = 0; dw < PKW; ++dw) {
                    float acc = sb[c];
                    #pragma unroll
                    for (int k = 0; k < KD; ++k) {
                        const int ir = IREL_D[dd][k];
                        if (ir < 0) continue;
                        #pragma unroll
                        for (int m = 0; m < KH; ++m) {
                            const int jr = IREL_H[dh][m];
                            if (jr < 0) continue;
                            #pragma unroll
                            for (int n = 0; n < KW; ++n) {
                                const int lr = IREL_W[dw][n];
                                if (lr < 0) continue;
                                #pragma unroll
                                for (int ic = 0; ic < CI; ++ic) {
                                    acc = fmaf(vin[ic][ir][jr][lr],
                                               sw[(((ic * CO + c) * KD + k) * KH + m) * KW + n],
                                               acc);
                                }
                            }
                        }
                    }
                    pm = fmaxf(pm, acc);
                }
            }
        }
        pvals[c] = pm;
        pmax = fmaxf(pmax, pm);
    }

    // Softmax across channels (numerically stable).
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        const float e = expf(pvals[c] - pmax);
        pvals[c] = e;
        sum += e;
    }
    const float inv_sum = 1.f / sum;

    // Subtract, Swish (x * sigmoid(x)), and max across channels.
    float outmax = -FLT_MAX;
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        const float y = pvals[c] * inv_sum - ss[c];
        const float s = y / (1.f + expf(-y));
        outmax = fmaxf(outmax, s);
    }

    out[idx] = outmax;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                            torch::Tensor sub, int64_t PD, int64_t PH, int64_t PW)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

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

    const int B = (int)xc.size(0);
    const int D = (int)xc.size(2);
    const int H = (int)xc.size(3);
    const int W = (int)xc.size(4);

    auto out = at::empty({B, PD, PH, PW}, xc.options());

    const int total = B * (int)PD * (int)PH * (int)PW;
    if (total == 0) return out;

    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    const size_t shmem = (size_t)(CI * CO * KD * KH * KW + 2 * CO) * sizeof(float);

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_ct3d_pool_softmax_sub_swish_max_kernel<<<blocks, threads, shmem, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
        sc.data_ptr<float>(), out.data_ptr<float>(),
        B, D, H, W, (int)PD, (int)PH, (int)PW);

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ",
                cudaGetErrorString(err));
    return out;
}
"""

_CPP_SRC = ("torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, "
            "torch::Tensor bias, torch::Tensor sub, int64_t PD, int64_t PH, "
            "int64_t PW);")


def _triple(v):
    if isinstance(v, (list, tuple)):
        assert len(v) == 3
        return (int(v[0]), int(v[1]), int(v[2]))
    v = int(v)
    return (v, v, v)


def _irel_table(stride, pad, ksize, pool_pad, pool_k):
    """table[dd][k] = input offset relative to the pooled output index for
    window offset dd and kernel tap k (-1 if the tap does not land there)."""
    tab = []
    for dd in range(pool_k):
        row = []
        for kk in range(ksize):
            t = dd - pool_pad + pad - kk
            if t >= 0 and t % stride == 0:
                row.append(t // stride)
            else:
                row.append(-1)
        tab.append(row)
    return tab


def _fmt_table(tab):
    return "{" + ", ".join("{" + ", ".join(str(v) for v in row) + "}"
                           for row in tab) + "}"


_EXT_CACHE = {}


def _get_extension(cfg):
    if cfg in _EXT_CACHE:
        return _EXT_CACHE[cfg]

    (ci, co, kd, kh, kw, sd, sh, sw_, pd_, ph, pw_, od, oh, ow,
     pkd, pkh, pkw, psd, psh, psw, ppd, pph, ppw) = cfg

    tabs = [
        _irel_table(sd, pd_, kd, ppd, pkd),
        _irel_table(sh, ph, kh, pph, pkh),
        _irel_table(sw_, pw_, kw, ppw, pkw),
    ]
    blks = []
    for tab in tabs:
        m = 0
        for row in tab:
            for v in row:
                if v > m:
                    m = v
        blks.append(m + 1)

    src = _CUDA_SRC
    for tok, val in (
        ("@CI@", str(ci)), ("@CO@", str(co)),
        ("@PKD@", str(pkd)), ("@PKH@", str(pkh)), ("@PKW@", str(pkw)),
        ("@KD@", str(kd)), ("@KH@", str(kh)), ("@KW@", str(kw)),
        ("@BD@", str(blks[0])), ("@BH@", str(blks[1])), ("@BW@", str(blks[2])),
        ("@TAB_D@", _fmt_table(tabs[0])),
        ("@TAB_H@", _fmt_table(tabs[1])),
        ("@TAB_W@", _fmt_table(tabs[2])),
    ):
        src = src.replace(tok, val)

    name = ("fused_ct3d_pool_sm_swish_max_ci%d_co%d_k%d%d%d_s%d%d%d_p%d%d%d_"
            "o%d%d%d_pk%d%d%d_ps%d%d%d_pp%d%d%d" % cfg)

    mod = load_inline(
        name=name,
        cpp_sources=_CPP_SRC,
        cuda_sources=src,
        functions=["fused_forward"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
    )
    _EXT_CACHE[cfg] = mod
    return mod


class ModelNew(nn.Module):
    """
    Optimized version of Model.  A single fused CUDA kernel performs:
        ConvTranspose3d -> MaxPool3d -> Softmax(dim=1) -> Subtract -> Swish
        -> Max(dim=1)
    without ever materializing the large intermediate tensors.
    Falls back to the reference PyTorch implementation for configurations the
    fused kernel does not support.
    """

    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._k = _triple(kernel_size)
        self._s = _triple(stride)
        self._p = _triple(padding)
        self._o = _triple(output_padding)
        self._pk = _triple(pool_kernel_size)
        self._ps = _triple(pool_stride)
        self._pp = _triple(pool_padding)

        cfg = ((in_channels, out_channels) + self._k + self._s + self._p +
               self._o + self._pk + self._ps + self._pp)

        # Conditions required by the fused kernel.
        self._use_fused = (
            self._ps == self._s and
            self._pp == (0, 0, 0) and
            all(self._p[d] + 0 - (self._k[d] - 1) > -self._s[d]
                for d in range(3))
        )
        self._ext = _get_extension(cfg) if self._use_fused else None

    def _pooled_dims(self, D, H, W):
        sizes = (D, H, W)
        out = []
        for d in range(3):
            ct = ((sizes[d] - 1) * self._s[d] - 2 * self._p[d] +
                  self._k[d] + self._o[d])
            out.append((ct + 2 * self._pp[d] - self._pk[d]) // self._ps[d] + 1)
        return out

    def forward(self, x):
        if (self._use_fused and x.is_cuda and
                x.dtype == torch.float32 and x.dim() == 5):
            PD, PH, PW = self._pooled_dims(
                int(x.shape[2]), int(x.shape[3]), int(x.shape[4]))
            if PD > 0 and PH > 0 and PW > 0:
                return self._ext.fused_forward(
                    x, self.conv_transpose.weight, self.conv_transpose.bias,
                    self.subtract, PD, PH, PW)

        # Reference fallback.
        x = self.conv_transpose(x)
        x = F.max_pool3d(x, self._pk, self._ps, self._pp)
        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
```