```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 two adjacent output voxels along W (pooled positions
// ow and ow+1).  The 16 output channels are processed one float4 quarter at
// a time inside the pool-phase loop, which keeps the live accumulators at
// two float4 registers and lets two thread blocks reside on each SM, so one
// block's prologue/epilogue overlaps the other's main loop.
__global__ void __launch_bounds__(256, 2) 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], 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);

    // 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 * 16 + threadIdx.x) * 2;

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

    // Issue the input loads before the shared-memory fill and barrier so the
    // memory latency overlaps the weight transpose instead of stalling after
    // the barrier.  Coordinates are clamped so out-of-range threads never
    // touch invalid addresses (their results are discarded).
    int ohc = min(oh, Hin - 1);
    int owc = min(ow, Win - 1);

    bool vd1 = (od + 1) < Din;
    bool vh1 = (oh + 1) < Hin;
    bool vw1 = (ow + 1) < Win;
    bool vw2 = (ow + 2) < Win;

    float xl[2 * 2 * 3 * IC];  // [ad][ah][aw][ic]

    #pragma unroll
    for (int ad = 0; ad < 2; ++ad) {
        bool okd = (ad == 0) || vd1;
        int xd = okd ? (od + ad) : od;
        #pragma unroll
        for (int ah = 0; ah < 2; ++ah) {
            bool okh = (ah == 0) || vh1;
            int xh = okh ? (ohc + ah) : ohc;
            #pragma unroll
            for (int aw = 0; aw < 3; ++aw) {
                bool okw = (aw == 0) || ((aw == 1) ? vw1 : vw2);
                int xw = okw ? (owc + aw) : owc;
                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) * 3 + aw) * IC + ic] = valid ? v : 0.0f;
                }
            }
        }
    }

    for (int i = threadIdx.y * blockDim.x + threadIdx.x; i < W_SIZE; i += blockDim.x * blockDim.y) {
        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);
    }

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

    if (oh >= Hin || ow >= Win) return;

    float4 pooled0[OC4];
    float4 pooled1[OC4];

    #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) {
                // Process the 16 output channels one float4 quarter at a time
                // so only two float4 accumulators stay live.
                #pragma unroll
                for (int g = 0; g < OC4; ++g) {
                    float4 acc0;
                    float4 acc1;

                    #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) * 3 + aw) * IC;

                                #pragma unroll
                                for (int ic = 0; ic < IC; ++ic) {
                                    float x0v = xl[xbase + ic];
                                    float x1v = xl[xbase + IC + ic];
                                    float4 wv = sw[(koff * IC + ic) * OC4 + g];
                                    if (ad == 0 && ah == 0 && aw == 0 && ic == 0) {
                                        acc0 = make_float4(x0v * wv.x, x0v * wv.y, x0v * wv.z, x0v * wv.w);
                                        acc1 = make_float4(x1v * wv.x, x1v * wv.y, x1v * wv.z, x1v * wv.w);
                                    } else {
                                        acc0.x = fmaf(x0v, wv.x, acc0.x);
                                        acc0.y = fmaf(x0v, wv.y, acc0.y);
                                        acc0.z = fmaf(x0v, wv.z, acc0.z);
                                        acc0.w = fmaf(x0v, wv.w, acc0.w);
                                        acc1.x = fmaf(x1v, wv.x, acc1.x);
                                        acc1.y = fmaf(x1v, wv.y, acc1.y);
                                        acc1.z = fmaf(x1v, wv.z, acc1.z);
                                        acc1.w = fmaf(x1v, wv.w, acc1.w);
                                    }
                                }
                            }
                        }
                    }

                    // Bias is constant across pool phases, so it is added
                    // after the running max instead of inside every phase.
                    if (i == 0 && j == 0 && k == 0) {
                        pooled0[g] = acc0;
                        pooled1[g] = acc1;
                    } else {
                        pooled0[g].x = fmaxf(pooled0[g].x, acc0.x);
                        pooled0[g].y = fmaxf(pooled0[g].y, acc0.y);
                        pooled0[g].z = fmaxf(pooled0[g].z, acc0.z);
                        pooled0[g].w = fmaxf(pooled0[g].w, acc0.w);
                        pooled1[g].x = fmaxf(pooled1[g].x, acc1.x);
                        pooled1[g].y = fmaxf(pooled1[g].y, acc1.y);
                        pooled1[g].z = fmaxf(pooled1[g].z, acc1.z);
                        pooled1[g].w = fmaxf(pooled1[g].w, acc1.w);
                    }
                }
            }
        }
    }

    // Add the per-channel bias (hoisted out of the pool-phase loop).
    float4 p0[OC4];
    float4 p1[OC4];
    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        float4 b = sb[g];
        p0[g].x = pooled0[g].x + b.x;
        p0[g].y = pooled0[g].y + b.y;
        p0[g].z = pooled0[g].z + b.z;
        p0[g].w = pooled0[g].w + b.w;
        p1[g].x = pooled1[g].x + b.x;
        p1[g].y = pooled1[g].y + b.y;
        p1[g].z = pooled1[g].z + b.z;
        p1[g].w = pooled1[g].w + b.w;
    }

    // Stable softmax across the 16 channels for both outputs.
    float mx0 = p0[0].x, mx1 = p1[0].x;
    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        if (g > 0) {
            mx0 = fmaxf(mx0, p0[g].x);
            mx1 = fmaxf(mx1, p1[g].x);
        }
        mx0 = fmaxf(mx0, p0[g].y);  mx1 = fmaxf(mx1, p1[g].y);
        mx0 = fmaxf(mx0, p0[g].z);  mx1 = fmaxf(mx1, p1[g].z);
        mx0 = fmaxf(mx0, p0[g].w);  mx1 = fmaxf(mx1, p1[g].w);
    }

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

    float sum0 = 0.0f, sum1 = 0.0f;
    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        sum0 += p0[g].x;  sum1 += p1[g].x;
        sum0 += p0[g].y;  sum1 += p1[g].y;
        sum0 += p0[g].z;  sum1 += p1[g].z;
        sum0 += p0[g].w;  sum1 += p1[g].w;
    }
    float inv0 = __fdividef(1.0f, sum0);
    float inv1 = __fdividef(1.0f, sum1);

    // Subtract, Swish, then max across channels.
    float res0 = -INFINITY, res1 = -INFINITY;
    #pragma unroll
    for (int g = 0; g < OC4; ++g) {
        float4 s = ss[g];
        float v;
        v = p0[g].x * inv0 - s.x;
        res0 = fmaxf(res0, fast_swish(v));
        v = p1[g].x * inv1 - s.x;
        res1 = fmaxf(res1, fast_swish(v));
        v = p0[g].y * inv0 - s.y;
        res0 = fmaxf(res0, fast_swish(v));
        v = p1[g].y * inv1 - s.y;
        res1 = fmaxf(res1, fast_swish(v));
        v = p0[g].z * inv0 - s.z;
        res0 = fmaxf(res0, fast_swish(v));
        v = p1[g].z * inv1 - s.z;
        res1 = fmaxf(res1, fast_swish(v));
        v = p0[g].w * inv0 - s.w;
        res0 = fmaxf(res0, fast_swish(v));
        v = p1[g].w * inv1 - s.w;
        res1 = fmaxf(res1, fast_swish(v));
    }

    // Adjacent even/odd W offsets guarantee 8-byte alignment for float2.
    long base = (((long)n * Din + od) * Hin + oh) * Win + ow;
    float2 r = make_float2(res0, res1);
    *reinterpret_cast<float2*>(out + base) = r;
}

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) % 2 == 0, "input width must be even 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 / 2) + 15) / 16;
    dim3 block(16, 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_v5",
    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) % 2 == 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
```