```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:
#   ConvTranspose3d (stride=2, pad=1, k=3, CI=3, CO=16)
#   + LogSumExp over channel dim
#   + HardSwish (x * sigmoid(x + 3) / 6)
#   + scalar bias subtraction
#   + clamp([-1, 1])
#
# The transposed convolution result for all 16 output channels of a given
# spatial voxel is accumulated in registers, so the huge (N, 16, D, H, W)
# intermediate tensor never touches global memory. Each thread produces two
# horizontally adjacent output voxels (even/odd parity), which makes the
# kernel-tap schedule uniform across a warp and keeps input loads coalesced.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

template <int CO>
__device__ __forceinline__ float lse_hswish_bias_clamp(float (&vals)[CO], float sb) {
    // numerically stable logsumexp over channels
    float m = vals[0];
    #pragma unroll
    for (int co = 1; co < CO; ++co) m = fmaxf(m, vals[co]);
    float s = 0.f;
    #pragma unroll
    for (int co = 0; co < CO; ++co) s += expf(vals[co] - m);
    const float lse = m + logf(s);
    // hardswish as written in the reference: x * sigmoid(x + 3) / 6
    const float sig = 1.f / (1.f + expf(-(lse + 3.f)));
    float y = (lse * sig) / 6.f - sb;
    return fminf(fmaxf(y, -1.f), 1.f);
}

template <int CI, int CO>
__global__ __launch_bounds__(256)
void fused_convt3d_lse_hswish_clamp_kernel(
    const float* __restrict__ x,          // (N, CI, ID, IH, IW)
    const float* __restrict__ w,          // (CI, CO, 3, 3, 3)
    const float* __restrict__ conv_bias,  // (CO)
    const float* __restrict__ scalar_bias,// (1)
    float* __restrict__ out,              // (N, 1, OD, OH, OW)
    const int N, const int ID, const int IH, const int IW,
    const int OD, const int OH, const int OW)
{
    constexpr int K = 3;
    constexpr int WSIZE = CI * CO * K * K * K;

    __shared__ float ws[WSIZE];
    __shared__ float cb[CO];

    for (int i = threadIdx.x; i < WSIZE; i += blockDim.x) ws[i] = w[i];
    if (threadIdx.x < CO) cb[threadIdx.x] = conv_bias[threadIdx.x];
    __syncthreads();

    const int OW2 = (OW + 1) >> 1;   // one thread handles a pair of adjacent voxels
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = N * OD * OH * OW2;
    if (idx >= total) return;

    const int wp = idx % OW2;
    int t = idx / OW2;
    const int oh = t % OH; t /= OH;
    const int od = t % OD;
    const int n  = t / OD;

    const bool has1 = (wp * 2 + 1) < OW;   // second (odd) voxel exists

    float vals0[CO];
    float vals1[CO];
    #pragma unroll
    for (int co = 0; co < CO; ++co) {
        vals0[co] = cb[co];
        vals1[co] = cb[co];
    }

    // out[o] = sum_{i,k} in[i] * w[k], with o = 2*i - 1 + k  (stride=2, pad=1)
    // => i = (o + 1 - k) / 2, valid when (o + 1 - k) is even and in range.
    const int kd0 = (od + 1) & 1;
    const int kh0 = (oh + 1) & 1;

    const float* xn = x + static_cast<long long>(n) * (CI * ID * IH * IW);

    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        const float* xc = xn + ci * (ID * IH * IW);
        const float* wc = ws + ci * (CO * 27);
        for (int kd = kd0; kd < K; kd += 2) {
            const int id = (od + 1 - kd) >> 1;
            if (static_cast<unsigned>(id) >= static_cast<unsigned>(ID)) continue;
            const float* xd = xc + id * (IH * IW);
            const float* wd = wc + kd * 9;
            for (int kh = kh0; kh < K; kh += 2) {
                const int ih = (oh + 1 - kh) >> 1;
                if (static_cast<unsigned>(ih) >= static_cast<unsigned>(IH)) continue;
                const float* xh = xd + ih * IW;
                const float* wh = wd + kh * 3;

                // kw = 1 contributes to the even voxel (ow = 2*wp), iw = wp
                {
                    const float xv = __ldg(xh + wp);
                    const float* wv = wh + 1;
                    #pragma unroll
                    for (int co = 0; co < CO; ++co)
                        vals0[co] = fmaf(xv, wv[co * 27], vals0[co]);
                }
                if (has1) {
                    // kw = 0 contributes to the odd voxel, iw = wp + 1
                    {
                        const float xv = __ldg(xh + wp + 1);
                        const float* wv = wh;
                        #pragma unroll
                        for (int co = 0; co < CO; ++co)
                            vals1[co] = fmaf(xv, wv[co * 27], vals1[co]);
                    }
                    // kw = 2 contributes to the odd voxel, iw = wp
                    {
                        const float xv = __ldg(xh + wp);
                        const float* wv = wh + 2;
                        #pragma unroll
                        for (int co = 0; co < CO; ++co)
                            vals1[co] = fmaf(xv, wv[co * 27], vals1[co]);
                    }
                }
            }
        }
    }

    const float sb = __ldg(scalar_bias);
    const float y0 = lse_hswish_bias_clamp<CO>(vals0, sb);

    const int out_idx = ((n * OD + od) * OH + oh) * OW + wp * 2;
    out[out_idx] = y0;
    if (has1) {
        const float y1 = lse_hswish_bias_clamp<CO>(vals1, sb);
        out[out_idx + 1] = y1;
    }
}

}  // namespace

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor scalar_bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && conv_bias.is_cuda() && scalar_bias.is_cuda(),
                "all tensors must be CUDA tensors");
    TORCH_CHECK(x.dim() == 5, "input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat &&
                conv_bias.scalar_type() == at::kFloat && scalar_bias.scalar_type() == at::kFloat,
                "float32 tensors required");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = conv_bias.contiguous();
    auto bs = scalar_bias.contiguous();

    const int N  = static_cast<int>(xc.size(0));
    const int CI = static_cast<int>(xc.size(1));
    const int ID = static_cast<int>(xc.size(2));
    const int IH = static_cast<int>(xc.size(3));
    const int IW = static_cast<int>(xc.size(4));
    const int CO = static_cast<int>(wc.size(1));

    TORCH_CHECK(CI == 3 && CO == 16, "unsupported channel configuration");
    TORCH_CHECK(bs.numel() == 1, "scalar bias must contain exactly one element");

    // output size of ConvTranspose3d with stride=2, padding=1, kernel=3
    const int OD = (ID - 1) * 2 - 2 + 3;
    const int OH = (IH - 1) * 2 - 2 + 3;
    const int OW = (IW - 1) * 2 - 2 + 3;

    auto out = torch::empty({N, 1, OD, OH, OW}, xc.options());

    const int OW2 = (OW + 1) / 2;
    const int total = N * OD * OH * OW2;
    const int block = 256;
    const int grid = (total + block - 1) / block;

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_convt3d_lse_hswish_clamp_kernel<3, 16><<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(), bs.data_ptr<float>(),
        out.data_ptr<float>(), N, ID, IH, IW, OD, OH, OW);

    return out;
}
"""

_cpp_src = ("torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
            "torch::Tensor conv_bias, torch::Tensor scalar_bias);")

_fused = load_inline(
    name="fused_convt3d_lse_hswish_clamp",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: a single fused CUDA kernel performs the 3D transposed
    convolution, LogSumExp, HardSwish, bias subtraction and clamp without
    materializing the large (N, C_out, D, H, W) intermediate tensor.
    """
    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 forward(self, x):
        ct = self.conv_transpose
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and ct.bias is not None
            and ct.groups == 1
            and ct.in_channels == 3
            and ct.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)
        ):
            return _fused.fused_forward(x, ct.weight, ct.bias, self.bias)

        # generic fallback (reference implementation)
        x = ct(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
```