```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 (v2):
#   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])
#
# Optimizations over v1:
#   * Each thread computes a 2x2 tile of output voxels (height x width
#     parities), so every loaded input value is reused across all kernel taps
#     that need it (4 loads feed 144 FMAs per (ci, kd)).
#   * Weights are staged in shared memory in a transposed [tap][co] layout so
#     the 16 channel weights of a tap are fetched with 4 float4 (128-bit)
#     shared loads instead of 16 scalar ones, cutting shared-memory
#     instruction count 4x.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr int CI = 3;
constexpr int CO = 16;
constexpr int K  = 3;
constexpr int WTAP_FLOATS = CI * CO * K * K * K;  // 1296

__device__ __forceinline__ float lse_hswish_bias_clamp(const float* v, float sb) {
    float m = v[0];
    #pragma unroll
    for (int c = 1; c < CO; ++c) m = fmaxf(m, v[c]);
    float s = 0.f;
    #pragma unroll
    for (int c = 0; c < CO; ++c) s += expf(v[c] - m);
    const float lse = m + logf(s);
    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);
}

// Accumulate one (kh, kw) tap for all 16 output channels.
#define APPLY_TAP(ACC, XV, KD, KH, KW)                                              \
    do {                                                                            \
        const float4* wtp = ws4 + ((((ci * K + (KD)) * K + (KH)) * K + (KW)) << 2); \
        _Pragma("unroll")                                                           \
        for (int j = 0; j < 4; ++j) {                                               \
            const float4 wv = wtp[j];                                               \
            ACC[4*j+0] = fmaf((XV), wv.x, ACC[4*j+0]);                              \
            ACC[4*j+1] = fmaf((XV), wv.y, ACC[4*j+1]);                              \
            ACC[4*j+2] = fmaf((XV), wv.z, ACC[4*j+2]);                              \
            ACC[4*j+3] = fmaf((XV), wv.w, ACC[4*j+3]);                              \
        }                                                                           \
    } while (0)

__global__ __launch_bounds__(256)
void fused_convt3d_lse_hswish_clamp_kernel_v2(
    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)
{
    __shared__ __align__(16) float ws[WTAP_FLOATS];
    __shared__ __align__(16) float cb[CO];

    // transpose weights into [tap][co] layout for float4 access
    for (int i = threadIdx.x; i < WTAP_FLOATS; i += blockDim.x) {
        const int ci  = i / (CO * 27);
        const int rem = i - ci * (CO * 27);
        const int t   = rem >> 4;   // tap index 0..26
        const int co  = rem & 15;
        ws[i] = w[ci * (CO * 27) + co * 27 + t];
    }
    if (threadIdx.x < CO) cb[threadIdx.x] = conv_bias[threadIdx.x];
    __syncthreads();

    const float4* ws4 = reinterpret_cast<const float4*>(ws);

    const int OH2 = (OH + 1) >> 1;
    const int OW2 = (OW + 1) >> 1;

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = N * OD * OH2 * OW2;
    if (idx >= total) return;

    const int wp = idx % OW2;
    int t0 = idx / OW2;
    const int hp = t0 % OH2; t0 /= OH2;
    const int od = t0 % OD;
    const int n  = t0 / OD;

    const bool has_w1 = (2 * wp + 1) < OW;   // odd-width voxel / iw = wp+1 valid
    const bool has_h1 = (2 * hp + 1) < OH;   // odd-height voxel / ih = hp+1 valid

    float aEE[CO], aEO[CO], aOE[CO], aOO[CO];
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        const float b = cb[c];
        aEE[c] = b; aEO[c] = b; aOE[c] = b; aOO[c] = b;
    }

    const float* xn = x + static_cast<long long>(n) * (CI * ID * IH * IW);
    const int kd0 = (od + 1) & 1;   // od even -> kd=1 only; od odd -> kd in {0,2}

    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        const float* xc = xn + ci * (ID * IH * IW);
        for (int kd = kd0; kd < K; kd += 2) {
            const int id = (od + 1 - kd) >> 1;   // always in [0, ID) for OD = 2*ID-1
            const float* row0 = xc + id * (IH * IW) + hp * IW;
            const float* row1 = row0 + IW;

            const float x00 = __ldg(row0 + wp);
            const float x01 = has_w1 ? __ldg(row0 + wp + 1) : 0.f;
            float x10 = 0.f, x11 = 0.f;
            if (has_h1) {
                x10 = __ldg(row1 + wp);
                x11 = has_w1 ? __ldg(row1 + wp + 1) : 0.f;
            }

            // oh even (=2hp) uses kh=1, ih=hp; oh odd (=2hp+1) uses kh=0 (ih=hp+1)
            // and kh=2 (ih=hp). ow even (=2wp) uses kw=1, iw=wp; ow odd (=2wp+1)
            // uses kw=0 (iw=wp+1) and kw=2 (iw=wp).
            APPLY_TAP(aEE, x00, kd, 1, 1);
            APPLY_TAP(aEO, x01, kd, 1, 0);
            APPLY_TAP(aEO, x00, kd, 1, 2);
            APPLY_TAP(aOE, x10, kd, 0, 1);
            APPLY_TAP(aOO, x11, kd, 0, 0);
            APPLY_TAP(aOO, x10, kd, 0, 2);
            APPLY_TAP(aOE, x00, kd, 2, 1);
            APPLY_TAP(aOO, x01, kd, 2, 0);
            APPLY_TAP(aOO, x00, kd, 2, 2);
        }
    }

    const float sb = __ldg(scalar_bias);
    const int out_base = ((n * OD + od) * OH + 2 * hp) * OW + 2 * wp;

    out[out_base] = lse_hswish_bias_clamp(aEE, sb);
    if (has_w1) {
        out[out_base + 1] = lse_hswish_bias_clamp(aEO, sb);
    }
    if (has_h1) {
        const int ob2 = out_base + OW;
        out[ob2] = lse_hswish_bias_clamp(aOE, sb);
        if (has_w1) {
            out[ob2 + 1] = lse_hswish_bias_clamp(aOO, sb);
        }
    }
}

#undef APPLY_TAP

}  // namespace

torch::Tensor fused_forward_v2(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 CIv = 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 COv = static_cast<int>(wc.size(1));

    TORCH_CHECK(CIv == 3 && COv == 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 OH2 = (OH + 1) / 2;
    const int OW2 = (OW + 1) / 2;
    const int total = N * OD * OH2 * OW2;
    const int block = 256;
    const int grid = (total + block - 1) / block;

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_convt3d_lse_hswish_clamp_kernel_v2<<<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_v2(torch::Tensor x, torch::Tensor weight, "
            "torch::Tensor conv_bias, torch::Tensor scalar_bias);")

_fused = load_inline(
    name="fused_convt3d_lse_hswish_clamp_v2",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward_v2"],
    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_v2(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
```