```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

# Fused CUDA kernel:
#   y = ConvTranspose3d(x)                 (cin=3, cout=16, k=3, stride=2, pad=1)
#   z = logsumexp(y, dim=channel)          (online max/sum, fully in registers)
#   z = z * sigmoid(z + 3) / 6             (hardswish)
#   z = clamp(z - bias, -1, 1)
# Each thread produces one output spatial voxel: it evaluates all 16 output
# channels of the transposed convolution in registers (each loaded input tap is
# reused across all 16 channels), reduces them with logsumexp, applies the
# elementwise tail and writes a single float. The huge (N,16,Do,Ho,Wo)
# intermediate tensor is never materialized in global memory.
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define CIN 3
#define COUT 16
#define KTAP 27
#define WWORDS 1296  // 3*16*3*3*3

__global__ __launch_bounds__(256)
void fused_convt3d_lse_hswish_clamp_kernel(
        const float* __restrict__ x,         // (N, 3, Di, Hi, Wi)
        const float* __restrict__ weight,    // (3, 16, 3, 3, 3)
        const float* __restrict__ conv_bias, // (16)
        const float* __restrict__ sub_bias,  // scalar
        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)
{
    extern __shared__ float smem[];
    float* sw = smem;            // convolution weights
    float* sb = smem + WWORDS;   // convolution bias

    for (int i = threadIdx.x; i < WWORDS; i += blockDim.x) sw[i] = weight[i];
    if (threadIdx.x < COUT) sb[threadIdx.x] = conv_bias[threadIdx.x];
    __syncthreads();

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

    const int w = idx % Wo;
    int t = idx / Wo;
    const int h = t % Ho;
    t /= Ho;
    const int d = t % Do;
    const int n = t / Do;

    // For stride=2, pad=1, k=3 an output coord o gets contributions from
    // input i and tap k where o = 2*i + k - 1:
    //   o even -> k=1, i=o/2
    //   o odd  -> k=0, i=(o+1)/2  and  k=2, i=(o-1)/2
    int nkd, kds[2], ids[2];
    if ((d & 1) == 0) { nkd = 1; kds[0] = 1; ids[0] = d >> 1; }
    else { nkd = 2; kds[0] = 0; ids[0] = (d + 1) >> 1; kds[1] = 2; ids[1] = (d - 1) >> 1; }

    int nkh, khs[2], ihs[2];
    if ((h & 1) == 0) { nkh = 1; khs[0] = 1; ihs[0] = h >> 1; }
    else { nkh = 2; khs[0] = 0; ihs[0] = (h + 1) >> 1; khs[1] = 2; ihs[1] = (h - 1) >> 1; }

    int nkw, kws[2], iws[2];
    if ((w & 1) == 0) { nkw = 1; kws[0] = 1; iws[0] = w >> 1; }
    else { nkw = 2; kws[0] = 0; iws[0] = (w + 1) >> 1; kws[1] = 2; iws[1] = (w - 1) >> 1; }

    const int plane  = Hi * Wi;
    const int cplane = Di * plane;
    const float* xn = x + (long)n * (CIN * cplane);

    float vals[COUT];
    #pragma unroll
    for (int oc = 0; oc < COUT; ++oc) vals[oc] = sb[oc];

    for (int a = 0; a < nkd; ++a) {
        const int idoff = ids[a] * plane;
        const int wd = kds[a] * 9;
        for (int b = 0; b < nkh; ++b) {
            const int ioff = idoff + ihs[b] * Wi;
            const int wh = wd + khs[b] * 3;
            for (int c = 0; c < nkw; ++c) {
                const int xoff = ioff + iws[c];
                const float x0 = __ldg(xn + xoff);
                const float x1 = __ldg(xn + cplane + xoff);
                const float x2 = __ldg(xn + 2 * cplane + xoff);
                const float* wp = sw + (wh + kws[c]);
                #pragma unroll
                for (int oc = 0; oc < COUT; ++oc) {
                    vals[oc] = fmaf(x0, wp[oc * KTAP],
                              fmaf(x1, wp[oc * KTAP + COUT * KTAP],
                              fmaf(x2, wp[oc * KTAP + 2 * COUT * KTAP], vals[oc])));
                }
            }
        }
    }

    // logsumexp over the 16 channels (max-shifted)
    float m = vals[0];
    #pragma unroll
    for (int oc = 1; oc < COUT; ++oc) m = fmaxf(m, vals[oc]);
    float s = 0.f;
    #pragma unroll
    for (int oc = 0; oc < COUT; ++oc) s += __expf(vals[oc] - m);
    const float lse = m + __logf(s);

    // hardswish, bias subtraction, clamp
    const float sig = 1.f / (1.f + __expf(-(lse + 3.f)));
    float y = lse * sig / 6.f;
    y -= __ldg(sub_bias);
    y = fminf(fmaxf(y, -1.f), 1.f);

    out[idx] = y;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor sub_bias)
{
    TORCH_CHECK(x.is_cuda() && x.dim() == 5, "expected 5D CUDA input");
    const c10::cuda::OptionalCUDAGuard guard(x.device());

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = conv_bias.contiguous();
    auto sc = sub_bias.contiguous();

    const int N  = (int)xc.size(0);
    const int Di = (int)xc.size(2);
    const int Hi = (int)xc.size(3);
    const int Wi = (int)xc.size(4);
    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 = torch::empty({N, 1, Do, Ho, Wo}, xc.options());
    const int total = N * Do * Ho * Wo;
    if (total == 0) return out;

    const int block = 256;
    const int grid = (total + block - 1) / block;
    const size_t smem = (WWORDS + COUT) * sizeof(float);

    fused_convt3d_lse_hswish_clamp_kernel<<<grid, block, smem,
        at::cuda::getCurrentCUDAStream()>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            sc.data_ptr<float>(), out.data_ptr<float>(),
            N, Di, Hi, Wi, Do, Ho, Wo);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fused_cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor sub_bias);"
)

_fused_ext = load_inline(
    name="fused_convt3d_lse_hswish_v1",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. The ConvTranspose3d, channel LogSumExp,
    HardSwish, bias subtraction and clamp are fused into a single custom CUDA
    kernel, eliminating the large (N, out_channels, Do, Ho, Wo) intermediate.
    Keeps identical parameters/state_dict as Model and falls back to the
    reference PyTorch ops for unsupported configurations.
    """
    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))

        def _triple(v):
            return tuple(v) if isinstance(v, (tuple, list)) else (v, v, v)

        self._fused_ok = (
            in_channels == 3 and out_channels == 16
            and _triple(kernel_size) == (3, 3, 3)
            and _triple(stride) == (2, 2, 2)
            and _triple(padding) == (1, 1, 1)
            and self.conv_transpose.bias is not None
        )

    def forward(self, x):
        if (self._fused_ok and x.is_cuda and x.dtype == torch.float32
                and self.conv_transpose.weight.is_cuda):
            return _fused_ext.fused_forward(
                x, self.conv_transpose.weight, self.conv_transpose.bias, self.bias)
        # Generic fallback: reference implementation.
        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
```