```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 kernel:
#   y = ConvTranspose3d(x, w, b)        (kernel=3, stride=2, padding=1)
#   z = logsumexp(y, dim=1)             (over the 16 output channels)
#   z = z * sigmoid(z + 3) / 6          (hardswish)
#   z = z - bias
#   z = clamp(z, -1, 1)
#
# Strategy: one thread per output spatial voxel. Each thread gathers its input
# contributions directly (pull formulation of transposed convolution), keeps all
# 16 output-channel accumulators in registers, then performs the channel-wise
# LogSumExp + elementwise epilogue in registers. This avoids materializing the
# full (N, 16, D', H', W') intermediate tensor entirely (~1 GB of traffic saved).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define OUT_C 16

__global__ void fused_convt3d_lse_hswish_bias_clamp_kernel(
    const float* __restrict__ x,          // (N, C_in, D, H, W)
    const float* __restrict__ wt,         // (C_in, 27, OUT_C) transposed weights
    const float* __restrict__ conv_bias,  // (OUT_C)
    const float* __restrict__ extra_bias, // (1,)
    float* __restrict__ out,              // (N, 1, D_out, H_out, W_out)
    const int N, const int C_in,
    const int D, const int H, const int W,
    const int D_out, const int H_out, const int W_out)
{
    // Stage transposed weights + conv bias in shared memory (16B aligned).
    extern __shared__ float4 sw4[];
    float* sw = reinterpret_cast<float*>(sw4);
    const int wsize = C_in * 27 * OUT_C;
    for (int i = threadIdx.x; i < wsize; i += blockDim.x) sw[i] = wt[i];
    if (threadIdx.x < OUT_C) sw[wsize + threadIdx.x] = conv_bias[threadIdx.x];
    __syncthreads();
    const float* sbias = sw + wsize;

    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    const long long total = (long long)N * D_out * H_out * W_out;
    if (idx >= total) return;

    const int ow = (int)(idx % W_out);
    const int oh = (int)((idx / W_out) % H_out);
    const int od = (int)((idx / ((long long)W_out * H_out)) % D_out);
    const int n  = (int)(idx / ((long long)W_out * H_out * D_out));

    // Valid kernel taps per dim: output o pulls from input i = (o + pad - k)/stride
    // with stride=2, pad=1, kernel=3 -> k valid iff (o + 1 - k) is even and in range.
    int klist[3][2];
    int ilist[3][2];
    int cnt[3];

    #define COMPUTE_TAPS(dim, o, L)                                    \
        {                                                              \
            int c = 0;                                                 \
            _Pragma("unroll")                                          \
            for (int k = 0; k < 3; ++k) {                              \
                int t = (o) + 1 - k;                                   \
                if ((t & 1) == 0) {                                    \
                    int i = t >> 1;                                    \
                    if (i >= 0 && i < (L)) {                           \
                        klist[dim][c] = k; ilist[dim][c] = i; ++c;     \
                    }                                                  \
                }                                                      \
            }                                                          \
            cnt[dim] = c;                                              \
        }

    COMPUTE_TAPS(0, od, D)
    COMPUTE_TAPS(1, oh, H)
    COMPUTE_TAPS(2, ow, W)
    #undef COMPUTE_TAPS

    float acc[OUT_C];
    #pragma unroll
    for (int oc = 0; oc < OUT_C; ++oc) acc[oc] = sbias[oc];

    const long long xCstride = (long long)D * H * W;
    const long long xNstride = (long long)C_in * xCstride;

    for (int a = 0; a < cnt[0]; ++a) {
        const int kd = klist[0][a];
        const int id = ilist[0][a];
        for (int b = 0; b < cnt[1]; ++b) {
            const int kh = klist[1][b];
            const int ih = ilist[1][b];
            for (int c = 0; c < cnt[2]; ++c) {
                const int kw = klist[2][c];
                const int iw = ilist[2][c];
                const int koff = kd * 9 + kh * 3 + kw;
                const long long xbase = n * xNstride + ((long long)id * H + ih) * W + iw;
                for (int ic = 0; ic < C_in; ++ic) {
                    const float xv = __ldg(x + xbase + (long long)ic * xCstride);
                    const float4* wp = reinterpret_cast<const float4*>(sw + (ic * 27 + koff) * OUT_C);
                    const float4 w0 = wp[0];
                    const float4 w1 = wp[1];
                    const float4 w2 = wp[2];
                    const float4 w3 = wp[3];
                    acc[0]  = fmaf(xv, w0.x, acc[0]);
                    acc[1]  = fmaf(xv, w0.y, acc[1]);
                    acc[2]  = fmaf(xv, w0.z, acc[2]);
                    acc[3]  = fmaf(xv, w0.w, acc[3]);
                    acc[4]  = fmaf(xv, w1.x, acc[4]);
                    acc[5]  = fmaf(xv, w1.y, acc[5]);
                    acc[6]  = fmaf(xv, w1.z, acc[6]);
                    acc[7]  = fmaf(xv, w1.w, acc[7]);
                    acc[8]  = fmaf(xv, w2.x, acc[8]);
                    acc[9]  = fmaf(xv, w2.y, acc[9]);
                    acc[10] = fmaf(xv, w2.z, acc[10]);
                    acc[11] = fmaf(xv, w2.w, acc[11]);
                    acc[12] = fmaf(xv, w3.x, acc[12]);
                    acc[13] = fmaf(xv, w3.y, acc[13]);
                    acc[14] = fmaf(xv, w3.z, acc[14]);
                    acc[15] = fmaf(xv, w3.w, acc[15]);
                }
            }
        }
    }

    // LogSumExp across the 16 channels (numerically stable).
    float m = acc[0];
    #pragma unroll
    for (int oc = 1; oc < OUT_C; ++oc) m = fmaxf(m, acc[oc]);
    float s = 0.0f;
    #pragma unroll
    for (int oc = 0; oc < OUT_C; ++oc) s += expf(acc[oc] - m);
    const float lse = m + logf(s);

    // HardSwish: x * sigmoid(x + 3) / 6
    const float sig = 1.0f / (1.0f + expf(-(lse + 3.0f)));
    float y = (lse * sig) / 6.0f;
    // Bias subtraction + clamp.
    y -= extra_bias[0];
    y = fminf(fmaxf(y, -1.0f), 1.0f);
    out[idx] = y;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor extra_bias)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
    TORCH_CHECK(weight.size(1) == OUT_C, "out_channels must be 16");
    TORCH_CHECK(weight.size(2) == 3 && weight.size(3) == 3 && weight.size(4) == 3,
                "kernel_size must be 3x3x3");

    auto xc  = x.contiguous();
    auto wtc = weight.permute({0, 2, 3, 4, 1}).contiguous(); // (C_in, 3,3,3, C_out)
    auto bc  = conv_bias.contiguous();
    auto ebc = extra_bias.contiguous();

    const int N    = xc.size(0);
    const int C_in = xc.size(1);
    const int D    = xc.size(2);
    const int H    = xc.size(3);
    const int W    = xc.size(4);

    // stride=2, padding=1, kernel=3
    const int D_out = (D - 1) * 2 - 2 + 3;
    const int H_out = (H - 1) * 2 - 2 + 3;
    const int W_out = (W - 1) * 2 - 2 + 3;

    auto out = torch::empty({N, 1, D_out, H_out, W_out}, xc.options());

    const long long total = (long long)N * D_out * H_out * W_out;
    if (total == 0) return out;

    const int block = 256;
    const unsigned int grid = (unsigned int)((total + block - 1) / block);
    const size_t shmem = (size_t)(C_in * 27 * OUT_C + OUT_C) * sizeof(float);

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_convt3d_lse_hswish_bias_clamp_kernel<<<grid, block, shmem, stream>>>(
        xc.data_ptr<float>(), wtc.data_ptr<float>(), bc.data_ptr<float>(),
        ebc.data_ptr<float>(), out.data_ptr<float>(),
        N, C_in, D, H, W, D_out, H_out, W_out);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor extra_bias);"
)

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


class ModelNew(nn.Module):
    """
    Optimized model: ConvTranspose3d + LogSumExp + HardSwish + bias + clamp
    fused into a single custom CUDA kernel.
    """
    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))
        self._use_fused = (kernel_size == 3 and stride == 2 and padding == 1
                           and out_channels == 16)

    def forward(self, x):
        if self._use_fused 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)
        # Generic fallback path (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
```