```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__ float fast_swish(float v) {
    return __fdividef(v, 1.0f + __expf(-v));
}

// Fused kernel for the target configuration:
//   ConvTranspose3d(ic=3, oc=16, k=3, s=2, p=1, op=1) -> MaxPool3d(2,2,0)
//   -> softmax(dim=1) -> subtract per-channel -> Swish -> max(dim=1)
//
// Each thread produces FOUR adjacent output voxels along W (pooled positions
// ow..ow+3).  All four outputs in a pool phase share the same kernel taps, so
// every shared-memory float4 weight load feeds 16 FMAs (4 outputs x 4
// channels), cutting shared-memory and issue pressure per FMA roughly in
// half compared to the two-output version.  Input activations are loaded as
// aligned float4 vectors in (ad,ah) slices just-in-time, and accumulators
// only cover one channel quarter at a time, which keeps the kernel within
// 128 registers so four 128-thread blocks reside on each SM.
__global__ void __launch_bounds__(128, 4) 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, int hin_tiles)
{
    // Weights are transposed on the fly from the PyTorch ConvTranspose3d
    // layout [ic][oc][kd][kh][kw] into [k][ic][oc] for float4 access over oc.
    __shared__ float4 sw[W4_SIZE];
    __shared__ float4 sb[OC4];
    __shared__ float4 ss[OC4];

    float* sws = reinterpret_cast<float*>(sw);
    float* sbs = reinterpret_cast<float*>(sb);
    float* sss = reinterpret_cast<float*>(ss);

    int tid = threadIdx.y * blockDim.x + threadIdx.x;
    for (int i = tid; i < W_SIZE; i += 128) {
        int ic = i / (OC * K3);
        int rem = i - ic * (OC * K3);
        int oc = rem / K3;
        int koff = rem - oc * K3;
        sws[(koff * IC + ic) * OC + oc] = __ldg(w + i);
    }
    if (tid < OC) {
        sbs[tid] = __ldg(bias + tid);
        sss[tid] = __ldg(sub + tid);
    }
    __syncthreads();

    // 3D grid removes all integer divisions from the per-thread index math.
    int n  = blockIdx.z;
    int od = blockIdx.y / hin_tiles;
    int oh = (blockIdx.y - od * hin_tiles) * 16 + threadIdx.y;
    int ow = (blockIdx.x * 8 + threadIdx.x) * 4;
    if (oh >= Hin || ow >= Win) return;

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

    float4 pooled[4][OC4];  // running pool max for outputs m=0..3, channel quarters
    int phase = 0;

    #pragma unroll
    for (int ad = 0; ad < 2; ++ad) {
        bool vd = (ad == 0) || (od + 1 < Din);
        int xd = vd ? (od + ad) : od;

        #pragma unroll
        for (int ah = 0; ah < 2; ++ah) {
            bool vh = (ah == 0) || (oh + 1 < Hin);
            int xh = vh ? (oh + ah) : oh;
            bool vslice = vd && vh;

            // Load x[ic][xd][xh][ow .. ow+4] (float4 + one scalar) per ic.
            float xs[5 * IC];
            long rowoff = ((long)xd * Hin + xh) * Win + ow;
            bool vw4 = (ow + 4) < Win;
            int o4 = vw4 ? 4 : (Win - 1 - ow);
            #pragma unroll
            for (int ic = 0; ic < IC; ++ic) {
                const float* p = xn + (long)ic * spatial + rowoff;
                float4 v = *reinterpret_cast<const float4*>(p);
                float t = __ldg(p + o4);
                xs[0 * IC + ic] = vslice ? v.x : 0.0f;
                xs[1 * IC + ic] = vslice ? v.y : 0.0f;
                xs[2 * IC + ic] = vslice ? v.z : 0.0f;
                xs[3 * IC + ic] = vslice ? v.w : 0.0f;
                xs[4 * IC + ic] = (vslice && vw4) ? t : 0.0f;
            }

            // All pool phases (i,j,k) that consume this (ad,ah) input slice.
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                if (i == 0 && ad > 0) continue;
                int kd = (i == 0) ? 1 : (2 - 2 * ad);

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

                    #pragma unroll
                    for (int k = 0; k < 2; ++k) {
                        // One channel quarter at a time keeps only four
                        // float4 accumulators live.
                        #pragma unroll
                        for (int g = 0; g < OC4; ++g) {
                            float4 a0, a1, a2, a3;
                            bool init = true;

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

                                #pragma unroll
                                for (int ic = 0; ic < IC; ++ic) {
                                    float x0 = xs[(awt + 0) * IC + ic];
                                    float x1 = xs[(awt + 1) * IC + ic];
                                    float x2 = xs[(awt + 2) * IC + ic];
                                    float x3 = xs[(awt + 3) * IC + ic];
                                    float4 wv = sw[(koff * IC + ic) * OC4 + g];
                                    if (init) {
                                        a0 = make_float4(x0 * wv.x, x0 * wv.y, x0 * wv.z, x0 * wv.w);
                                        a1 = make_float4(x1 * wv.x, x1 * wv.y, x1 * wv.z, x1 * wv.w);
                                        a2 = make_float4(x2 * wv.x, x2 * wv.y, x2 * wv.z, x2 * wv.w);
                                        a3 = make_float4(x3 * wv.x, x3 * wv.y, x3 * wv.z, x3 * wv.w);
                                        init = false;
                                    } else {
                                        a0.x = fmaf(x0, wv.x, a0.x);
                                        a0.y = fmaf(x0, wv.y, a0.y);
                                        a0.z = fmaf(x0, wv.z, a0.z);
                                        a0.w = fmaf(x0, wv.w, a0.w);
                                        a1.x = fmaf(x1, wv.x, a1.x);
                                        a1.y = fmaf(x1, wv.y, a1.y);
                                        a1.z = fmaf(x1, wv.z, a1.z);
                                        a1.w = fmaf(x1, wv.w, a1.w);
                                        a2.x = fmaf(x2, wv.x, a2.x);
                                        a2.y = fmaf(x2, wv.y, a2.y);
                                        a2.z = fmaf(x2, wv.z, a2.z);
                                        a2.w = fmaf(x2, wv.w, a2.w);
                                        a3.x = fmaf(x3, wv.x, a3.x);
                                        a3.y = fmaf(x3, wv.y, a3.y);
                                        a3.z = fmaf(x3, wv.z, a3.z);
                                        a3.w = fmaf(x3, wv.w, a3.w);
                                    }
                                }
                            }

                            // Bias is constant across pool phases, so it is
                            // added after the running max instead of inside
                            // every phase.
                            if (phase == 0) {
                                pooled[0][g] = a0;
                                pooled[1][g] = a1;
                                pooled[2][g] = a2;
                                pooled[3][g] = a3;
                            } else {
                                float4 p;
                                p = pooled[0][g];
                                p.x = fmaxf(p.x, a0.x); p.y = fmaxf(p.y, a0.y);
                                p.z = fmaxf(p.z, a0.z); p.w = fmaxf(p.w, a0.w);
                                pooled[0][g] = p;
                                p = pooled[1][g];
                                p.x = fmaxf(p.x, a1.x); p.y = fmaxf(p.y, a1.y);
                                p.z = fmaxf(p.z, a1.z); p.w = fmaxf(p.w, a1.w);
                                pooled[1][g] = p;
                                p = pooled[2][g];
                                p.x = fmaxf(p.x, a2.x); p.y = fmaxf(p.y, a2.y);
                                p.z = fmaxf(p.z, a2.z); p.w = fmaxf(p.w, a2.w);
                                pooled[2][g] = p;
                                p = pooled[3][g];
                                p.x = fmaxf(p.x, a3.x); p.y = fmaxf(p.y, a3.y);
                                p.z = fmaxf(p.z, a3.z); p.w = fmaxf(p.w, a3.w);
                                pooled[3][g] = p;
                            }
                        }
                        ++phase;
                    }
                }
            }
        }
    }

    // Epilogue per output: add bias, stable softmax over 16 channels,
    // subtract per-channel constant, Swish, then max over channels.
    float res[4];
    #pragma unroll
    for (int m = 0; m < 4; ++m) {
        float e[OC];
        float mx = -INFINITY;
        #pragma unroll
        for (int g = 0; g < OC4; ++g) {
            float4 b = sb[g];
            e[g * 4 + 0] = pooled[m][g].x + b.x;
            e[g * 4 + 1] = pooled[m][g].y + b.y;
            e[g * 4 + 2] = pooled[m][g].z + b.z;
            e[g * 4 + 3] = pooled[m][g].w + b.w;
            mx = fmaxf(mx, fmaxf(fmaxf(e[g * 4 + 0], e[g * 4 + 1]),
                                 fmaxf(e[g * 4 + 2], e[g * 4 + 3])));
        }

        float sum = 0.0f;
        #pragma unroll
        for (int c = 0; c < OC; ++c) {
            e[c] = __expf(e[c] - mx);
            sum += e[c];
        }
        float inv = __fdividef(1.0f, sum);

        float r = -INFINITY;
        #pragma unroll
        for (int g = 0; g < OC4; ++g) {
            float4 s = ss[g];
            r = fmaxf(r, fast_swish(e[g * 4 + 0] * inv - s.x));
            r = fmaxf(r, fast_swish(e[g * 4 + 1] * inv - s.y));
            r = fmaxf(r, fast_swish(e[g * 4 + 2] * inv - s.z));
            r = fmaxf(r, fast_swish(e[g * 4 + 3] * inv - s.w));
        }
        res[m] = r;
    }

    // ow is a multiple of four and Win % 4 == 0, so the store is a
    // 16-byte-aligned float4.
    long base = (((long)n * Din + od) * Hin + oh) * Win + ow;
    *reinterpret_cast<float4*>(out + base) = make_float4(res[0], res[1], res[2], res[3]);
}

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.");
    TORCH_CHECK(x.size(4) % 4 == 0, "input width must be a multiple of 4 for the fused kernel.");

    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;

    int hin_tiles = (Hin + 15) / 16;
    int w_tiles = (Win + 31) / 32;
    dim3 block(8, 16, 1);
    dim3 grid(w_tiles, Din * hin_tiles, N);
    auto stream = at::cuda::getCurrentCUDAStream();

    fused_convt_pool_softmax_sub_swish_max_kernel<<<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, hin_tiles
    );

    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_v6",
    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 x.size(4) % 4 == 0
            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
```