```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
from torch.utils.cpp_extension import load_inline

# Fully fused CUDA kernel:
#   ConvTranspose3d (k=3, s=2, p=1) -> LogSumExp(dim=1) -> HardSwish
#   -> bias subtraction -> clamp(-1, 1)
#
# Key idea: with stride=2, pad=1, k=3, every output voxel of the transposed
# convolution only depends on up to 2x2x2 input voxels (selected purely by the
# parity of the output coordinates). Each thread computes ALL 16 output channels
# of one spatial voxel in registers, immediately reduces them with a numerically
# stable logsumexp, applies the pointwise chain, and writes a single float.
# The huge intermediate conv-transpose tensor (N, 16, Do, Ho, Wo) never touches
# global memory.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define C_OUT 16
#define C_IN 3
#define KD 3
#define K3 (KD * KD * KD)
#define WSIZE (C_IN * C_OUT * K3)

__global__ __launch_bounds__(256)
void fused_convt3d_lse_hswish_clamp_kernel(
    const float* __restrict__ input,     // [N, C_IN, Di, Hi, Wi]
    const float* __restrict__ weight,    // [C_IN, C_OUT, KD, KD, KD]
    const float* __restrict__ conv_bias, // [C_OUT]
    const float* __restrict__ sub_bias,  // [1]
    float* __restrict__ out,             // [N, 1, Do, Ho, Wo]
    const int N, const int Di, const int Hi, const int Wi,
    const int Do, const int Ho, const int Wo)
{
    // Stage all weights + conv bias into shared memory (all threads read the
    // same addresses, so this is broadcast-friendly).
    __shared__ float s_w[WSIZE + C_OUT];
    for (int i = threadIdx.x; i < WSIZE + C_OUT; i += blockDim.x) {
        s_w[i] = (i < WSIZE) ? weight[i] : conv_bias[i - WSIZE];
    }
    __syncthreads();

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = N * Do * Ho * Wo;
    if (idx >= total) return;

    const int ow = idx % Wo;
    const int t0 = idx / Wo;
    const int oh = t0 % Ho;
    const int t1 = t0 / Ho;
    const int od = t1 % Do;
    const int n  = t1 / Do;

    const int HiWi   = Hi * Wi;
    const int DiHiWi = Di * HiWi;
    const float* in_n = input + n * (C_IN * DiHiWi);

    // Accumulators for all output channels of this spatial voxel.
    float acc[C_OUT];
    #pragma unroll
    for (int oc = 0; oc < C_OUT; ++oc) acc[oc] = s_w[WSIZE + oc];

    // out[o] = sum_{o = i*2 - 1 + k} w[k] * in[i]  =>  i = (o + 1 - k) / 2,
    // valid only when (o + 1 - k) is even. Parity selects 1 or 2 taps per dim.
    #pragma unroll
    for (int kd = 0; kd < KD; ++kd) {
        const int td = od + 1 - kd;
        if (td & 1) continue;
        const int id = td >> 1;
        if (id < 0 || id >= Di) continue;
        #pragma unroll
        for (int kh = 0; kh < KD; ++kh) {
            const int th = oh + 1 - kh;
            if (th & 1) continue;
            const int ih = th >> 1;
            if (ih < 0 || ih >= Hi) continue;
            #pragma unroll
            for (int kw = 0; kw < KD; ++kw) {
                const int tw = ow + 1 - kw;
                if (tw & 1) continue;
                const int iw = tw >> 1;
                if (iw < 0 || iw >= Wi) continue;

                const int kidx  = (kd * KD + kh) * KD + kw;
                const int ibase = id * HiWi + ih * Wi + iw;
                #pragma unroll
                for (int ic = 0; ic < C_IN; ++ic) {
                    const float v = in_n[ic * DiHiWi + ibase];
                    const float* wp = &s_w[ic * (C_OUT * K3) + kidx];
                    #pragma unroll
                    for (int oc = 0; oc < C_OUT; ++oc) {
                        acc[oc] = fmaf(v, wp[oc * K3], acc[oc]);
                    }
                }
            }
        }
    }

    // Numerically stable logsumexp over the channel accumulators.
    float m = acc[0];
    #pragma unroll
    for (int oc = 1; oc < C_OUT; ++oc) m = fmaxf(m, acc[oc]);
    float s = 0.f;
    #pragma unroll
    for (int oc = 0; oc < C_OUT; ++oc) s += __expf(acc[oc] - m);
    const float lse = m + __logf(s);

    // HardSwish: x * sigmoid(x + 3) / 6
    const float sig = 1.f / (1.f + __expf(-(lse + 3.f)));
    float y = lse * sig / 6.f;
    // Bias subtraction and clamp.
    y -= sub_bias[0];
    out[idx] = fminf(fmaxf(y, -1.f), 1.f);
}

torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor sub_bias) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");

    auto x  = input.contiguous();
    auto w  = weight.contiguous();
    auto cb = conv_bias.contiguous();
    auto sb = sub_bias.contiguous();

    const int N  = x.size(0);
    const int Di = x.size(2);
    const int Hi = x.size(3);
    const int Wi = x.size(4);
    // ConvTranspose3d output size with stride=2, pad=1, k=3, dilation=1, output_padding=0.
    const int Do = (Di - 1) * 2 - 2 + 3;
    const int Ho = (Hi - 1) * 2 - 2 + 3;
    const int Wo = (Wi - 1) * 2 - 2 + 3;

    auto out = at::empty({N, 1, Do, Ho, Wo}, x.options());
    const int total = N * Do * Ho * Wo;
    if (total == 0) return out;

    const int block = 256;
    const int grid = (total + block - 1) / block;
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_convt3d_lse_hswish_clamp_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), cb.data_ptr<float>(),
        sb.data_ptr<float>(), out.data_ptr<float>(),
        N, Di, Hi, Wi, Do, Ho, Wo);
    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor sub_bias);"
)

_fused_op = load_inline(
    name="fused_convt3d_lse_hswish_clamp",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the entire forward pass (ConvTranspose3d + LogSumExp +
    HardSwish + bias subtraction + clamp) is executed by a single fused CUDA
    kernel whenever the architecture matches the specialized configuration
    (in_channels=3, out_channels=16, kernel=3, stride=2, padding=1). Otherwise
    it falls back to the original PyTorch operator sequence.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.bias = nn.Parameter(torch.randn(1, 1, 1, 1))

        ct = self.conv_transpose
        self._fast = (
            in_channels == 3
            and out_channels == 16
            and tuple(ct.kernel_size) == (3, 3, 3)
            and tuple(ct.stride) == (2, 2, 2)
            and tuple(ct.padding) == (1, 1, 1)
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
            and ct.groups == 1
            and ct.bias is not None
            and ct.padding_mode == "zeros"
        )

    def forward(self, x):
        if self._fast and x.is_cuda and x.dtype == torch.float32:
            return _fused_op.fused_forward(
                x,
                self.conv_transpose.weight,
                self.conv_transpose.bias,
                self.bias.reshape(1),
            )
        # Generic fallback (original operator sequence).
        x = self.conv_transpose(x)
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = x * torch.sigmoid(x + 3) / 6
        x = x - self.bias
        x = torch.clamp(x, min=-1, max=1)
        return x
```