```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 computes four pooled outputs on a 2x2 tile in (H, W) at a
// fixed depth.  All four outputs share the same kernel tap weights, so every
// shared-memory float4 weight load feeds 16 FMAs, halving shared-memory
// traffic per FMA compared to computing two outputs per thread.
__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, 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 the output-channel dimension.
    __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.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();

    // 3D grid removes 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;  // base pooled H
    int ow = (blockIdx.x * 16 + threadIdx.x) * 2;               // base pooled W

    bool vh0 = oh < Hin;
    bool vh1 = (oh + 1) < Hin;
    bool vw1 = (ow + 1) < Win;
    if (!(vh0 && ow < Win)) return;  // whole 2x2 tile is out of range

    // Input neighborhood needed by the 2x2 pooled outputs:
    //   d in [od, od+1], h in [oh, oh+2], w in [ow, ow+2].
    bool vd1 = (od + 1) < Din;
    bool vh2 = (oh + 2) < Hin;
    bool vw2 = (ow + 2) < Win;
    bool vhb[3] = {true, vh1, vh2};
    bool vwb[3] = {true, vw1, vw2};

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

    float xl[2 * 3 * 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 < 3; ++ah) {
            bool okh = vhb[ah];
            int xh = okh ? (oh + ah) : oh;
            #pragma unroll
            for (int aw = 0; aw < 3; ++aw) {
                bool okw = vwb[aw];
                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 * 3 + ah) * 3 + aw) * IC + ic] = valid ? v : 0.0f;
                }
            }
        }
    }

    float4 pooled[2][2][OC4];  // [c (H offset)][b (W offset)][channel group]
    int phase = 0;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {        // pool phase in D
        #pragma unroll
        for (int j = 0; j < 2; ++j) {    // pool phase in H
            #pragma unroll
            for (int k = 0; k < 2; ++k) {  // pool phase in W
                float4 acc[2][2][OC4];

                // First tap (ad=ah=aw=0, ic=0) initializes the accumulators
                // with multiplies, avoiding explicit zeroing.
                {
                    int kd = (i == 0) ? 1 : 2;
                    int kh = (j == 0) ? 1 : 2;
                    int kw = (k == 0) ? 1 : 2;
                    int koff = kd * 9 + kh * 3 + kw;
                    float x00 = xl[0 * IC + 0];
                    float x01 = xl[1 * IC + 0];
                    float x10 = xl[3 * IC + 0];
                    float x11 = xl[4 * IC + 0];
                    const float4* wp = sw + (koff * IC) * OC4;
                    #pragma unroll
                    for (int g = 0; g < OC4; ++g) {
                        float4 wv = wp[g];
                        acc[0][0][g] = make_float4(x00 * wv.x, x00 * wv.y, x00 * wv.z, x00 * wv.w);
                        acc[0][1][g] = make_float4(x01 * wv.x, x01 * wv.y, x01 * wv.z, x01 * wv.w);
                        acc[1][0][g] = make_float4(x10 * wv.x, x10 * wv.y, x10 * wv.z, x10 * wv.w);
                        acc[1][1][g] = make_float4(x11 * wv.x, x11 * wv.y, x11 * wv.z, x11 * wv.w);
                    }
                }

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

                            #pragma unroll
                            for (int ic = 0; ic < IC; ++ic) {
                                if (ad == 0 && ah == 0 && aw == 0 && ic == 0) continue;
                                int idx = xbase + ic;
                                float x00 = xl[idx];
                                float x01 = xl[idx + IC];
                                float x10 = xl[idx + 3 * IC];
                                float x11 = xl[idx + 4 * IC];
                                const float4* wp = sw + (koff * IC + ic) * OC4;

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

                // Bias is constant across pool phases, so it is added after
                // the running max instead of inside every phase.
                if (phase == 0) {
                    #pragma unroll
                    for (int c = 0; c < 2; ++c)
                        #pragma unroll
                        for (int b = 0; b < 2; ++b)
                            #pragma unroll
                            for (int g = 0; g < OC4; ++g)
                                pooled[c][b][g] = acc[c][b][g];
                } else {
                    #pragma unroll
                    for (int c = 0; c < 2; ++c)
                        #pragma unroll
                        for (int b = 0; b < 2; ++b)
                            #pragma unroll
                            for (int g = 0; g < OC4; ++g) {
                                pooled[c][b][g].x = fmaxf(pooled[c][b][g].x, acc[c][b][g].x);
                                pooled[c][b][g].y = fmaxf(pooled[c][b][g].y, acc[c][b][g].y);
                                pooled[c][b][g].z = fmaxf(pooled[c][b][g].z, acc[c][b][g].z);
                                pooled[c][b][g].w = fmaxf(pooled[c][b][g].w, acc[c][b][g].w);
                            }
                }
                ++phase;
            }
        }
    }

    // Epilogue per pooled output: add bias, stable softmax across the 16
    // channels, subtract, Swish, then max across channels.
    bool vout[2][2] = {{vh0, vh0 && vw1}, {vh1, vh1 && vw1}};

    #pragma unroll
    for (int c = 0; c < 2; ++c) {
        #pragma unroll
        for (int b = 0; b < 2; ++b) {
            if (!vout[c][b]) continue;

            float p[OC];
            #pragma unroll
            for (int g = 0; g < OC4; ++g) {
                float4 bb = sb[g];
                float4 v = pooled[c][b][g];
                p[g * 4 + 0] = v.x + bb.x;
                p[g * 4 + 1] = v.y + bb.y;
                p[g * 4 + 2] = v.z + bb.z;
                p[g * 4 + 3] = v.w + bb.w;
            }

            float mx = p[0];
            #pragma unroll
            for (int ch = 1; ch < OC; ++ch) mx = fmaxf(mx, p[ch]);

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

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

            long offset = (((long)n * Din + od) * Hin + (oh + c)) * Win + (ow + b);
            out[offset] = 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.");
    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_v4",
    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
```