```
# 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 that performs, in a single pass:
#   ConvTranspose3d (ic=3, oc=16, k=3, stride=2, pad=1, out_pad=1)
#   -> MaxPool3d (k=2, stride=2, pad=0)
#   -> Softmax over channel dim
#   -> Per-channel subtract
#   -> Swish (x * sigmoid(x))
#   -> Max over channel dim
#
# Because the pool windows are disjoint and stride=2, every conv-transpose
# output element is consumed exactly once, so the whole pipeline can be
# computed without materializing any intermediate tensor. For conv transpose
# with stride=2, pad=1, k=3, an output position y receives contributions from
# input x = (y + 1 - k) / 2 for valid integer x:
#   - y even: one tap  (k=1, x=y/2)
#   - y odd : two taps (k=2, x=(y-1)/2) and (k=0, x=(y+1)/2, if in range)

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

#define IC 3
#define OC 16
#define K3 27
#define WSIZE (IC * OC * K3)

__global__ void 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)
{
    __shared__ float sw[WSIZE];
    __shared__ float sb[OC];
    __shared__ float ss[OC];

    for (int i = threadIdx.x; i < WSIZE; i += blockDim.x) sw[i] = w[i];
    if (threadIdx.x < OC) {
        sb[threadIdx.x] = bias[threadIdx.x];
        ss[threadIdx.x] = sub[threadIdx.x];
    }
    __syncthreads();

    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long spatial = (long)Din * Hin * Win;
    long total = (long)N * spatial;
    if (idx >= total) return;

    int ow = (int)(idx % Win);
    int oh = (int)((idx / Win) % Hin);
    int od = (int)((idx / ((long)Win * Hin)) % Din);
    int n  = (int)(idx / spatial);

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

    float pooled[OC];
    #pragma unroll
    for (int c = 0; c < OC; c++) pooled[c] = -INFINITY;

    // 2x2x2 max-pool window over conv-transpose outputs, computed on the fly.
    #pragma unroll
    for (int i = 0; i < 2; i++) {
        int ntd = (i == 0) ? 1 : ((od + 1 < Din) ? 2 : 1);
        #pragma unroll
        for (int j = 0; j < 2; j++) {
            int nth = (j == 0) ? 1 : ((oh + 1 < Hin) ? 2 : 1);
            #pragma unroll
            for (int k = 0; k < 2; k++) {
                int ntw = (k == 0) ? 1 : ((ow + 1 < Win) ? 2 : 1);

                float acc[OC];
                #pragma unroll
                for (int c = 0; c < OC; c++) acc[c] = sb[c];

                for (int a = 0; a < ntd; a++) {
                    int xd = (i == 0) ? od : (od + a);
                    int kd = (i == 0) ? 1 : (2 - 2 * a);
                    for (int b = 0; b < nth; b++) {
                        int xh = (j == 0) ? oh : (oh + b);
                        int kh = (j == 0) ? 1 : (2 - 2 * b);
                        const float* xrc = xn + ((long)xd * Hin + xh) * Win;
                        for (int cc = 0; cc < ntw; cc++) {
                            int xw = (k == 0) ? ow : (ow + cc);
                            int kw = (k == 0) ? 1 : (2 - 2 * cc);
                            int koff = kd * 9 + kh * 3 + kw;
                            #pragma unroll
                            for (int ic = 0; ic < IC; ic++) {
                                float xv = xrc[(long)ic * spatial + xw];
                                const float* wp = sw + ic * (OC * K3) + koff;
                                #pragma unroll
                                for (int oc = 0; oc < OC; oc++) {
                                    acc[oc] = fmaf(xv, wp[oc * K3], acc[oc]);
                                }
                            }
                        }
                    }
                }
                #pragma unroll
                for (int c = 0; c < OC; c++) pooled[c] = fmaxf(pooled[c], acc[c]);
            }
        }
    }

    // Softmax over channels (numerically stable).
    float mx = pooled[0];
    #pragma unroll
    for (int c = 1; c < OC; c++) mx = fmaxf(mx, pooled[c]);
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < OC; c++) {
        pooled[c] = expf(pooled[c] - mx);
        sum += pooled[c];
    }
    float inv = 1.f / sum;

    // Subtract, Swish, Max over channels.
    float res = -INFINITY;
    #pragma unroll
    for (int c = 0; c < OC; c++) {
        float v = pooled[c] * inv - ss[c];
        v = v / (1.f + expf(-v));
        res = fmaxf(res, v);
    }
    out[idx] = 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(), "inputs must be CUDA tensors");
    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 block = 256;
    long grid = (total + block - 1) / block;
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_convt_pool_softmax_sub_swish_max_kernel<<<(unsigned int)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);
    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",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. The entire operator chain
    (ConvTranspose3d -> MaxPool3d -> Softmax -> Subtract -> Swish -> Max)
    is fused into a single custom CUDA kernel when the configuration matches
    the supported specialization; otherwise it falls back to the eager 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
        )

    def forward(self, x):
        if self._use_fused and x.is_cuda and x.dtype == torch.float32 and self.conv_transpose.bias is not None:
            return fused_op.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.subtract)

        # Fallback eager path
        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
```