```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])
#
# Optimizations:
#   * The (N, 16, D, H, W) transposed-conv result is never materialized; all 16
#     channels of each output voxel are accumulated in registers.
#   * Each thread produces FOUR horizontally adjacent output voxels (two
#     even/odd pairs). The same kernel weights apply to all four voxels, so
#     every weight fetched from shared memory feeds several FMAs, and the three
#     gathered input values (x0, x1, x2) are reused across voxels.
#   * For the common 16x32x32 input the kernel is specialized by output parity
#     (blockIdx.y selects the (depth, height) parity class). With the parity
#     known at compile time the kd/kh tap loops have compile-time trip counts,
#     everything is fully unrolled, and every shared-memory / global-memory
#     address becomes a base register plus an immediate offset.
#   * Weights are repacked once per call into a [ci][kd][kh][kw][co] layout so
#     the inner loop fetches them as broadcast float4 shared-memory loads.
#   * Fast SFU transcendentals (__expf / __logf / __fdividef) in the epilogue.
#   * A generic fallback kernel (same fusion, runtime indices) handles any
#     other spatial size.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

namespace {

constexpr int CI = 3;
constexpr int CO = 16;
constexpr int REPACK = CI * 27 * CO;  // [ci][kd][kh][kw][co] == 1296 floats

// Repack weights from ConvTranspose3d layout [CI][CO][3][3][3] into
// [CI][kd][kh][kw][co] so each (tap, kw) plane holds 16 contiguous co values.
__global__ void repack_kernel(const float* __restrict__ w, float* __restrict__ wr) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= REPACK) return;
    const int co = i & 15;
    int t = i >> 4;
    const int kw = t % 3; t /= 3;
    const int kh = t % 3; t /= 3;
    const int kd = t % 3;
    const int ci = t / 3;
    wr[i] = w[((ci * CO + co) * 27) + kd * 9 + kh * 3 + kw];
}

__device__ __forceinline__ float epilogue(const float (&v)[CO], const float sb) {
    // numerically stable logsumexp over channels
    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);
    // hardswish as in the reference: x * sigmoid(x + 3) / 6
    const float sig = __fdividef(1.f, 1.f + __expf(-(lse + 3.f)));
    const float y = fmaf(lse, sig * 0.16666666666666666f, -sb);
    return fminf(fmaxf(y, -1.f), 1.f);
}

// ---------------------------------------------------------------------------
// Specialized path for ID=16, IH=32, IW=32 (OD=31, OH=63, OW=63).
// blockIdx.y selects the parity class: PD = by & 1, PH = by >> 1.
//   even output index o = 2i   -> single tap k=1 at input i
//   odd  output index o = 2i+1 -> taps k=0 at input i+1 and k=2 at input i
// Each thread computes output voxels ow = 4q .. 4q+3 (q < 16):
//   a: ow=4q   even, kw=1 @ iw=2q
//   b: ow=4q+1 odd,  kw=0 @ iw=2q+1, kw=2 @ iw=2q
//   c: ow=4q+2 even, kw=1 @ iw=2q+1
//   d: ow=4q+3 odd,  kw=0 @ iw=2q+2, kw=2 @ iw=2q+1   (valid only for q < 15)
// ---------------------------------------------------------------------------

// One (kd, kh) tap for the 4-voxel group. xh points at the input row base,
// wplane at the 48-float [kw][co] plane in shared memory.
__device__ __forceinline__ void tapfma_spec(
    const float* __restrict__ xh, const float* wplane,
    const int wp, const int i1, const int i2,
    float* va, float* vb, float* vc, float* vd)
{
    const float x0 = __ldg(xh + wp);
    const float x1 = __ldg(xh + i1);
    const float x2 = __ldg(xh + i2);
    const float4* wq = reinterpret_cast<const float4*>(wplane);
    #pragma unroll
    for (int cg = 0; cg < 4; ++cg) {
        const float4 W0 = wq[cg];      // kw = 0, co = 4cg .. 4cg+3
        const float4 W1 = wq[4 + cg];  // kw = 1
        const float4 W2 = wq[8 + cg];  // kw = 2
        const int c0 = 4 * cg;
        va[c0+0] = fmaf(x0, W1.x, va[c0+0]);
        va[c0+1] = fmaf(x0, W1.y, va[c0+1]);
        va[c0+2] = fmaf(x0, W1.z, va[c0+2]);
        va[c0+3] = fmaf(x0, W1.w, va[c0+3]);
        vb[c0+0] = fmaf(x1, W0.x, vb[c0+0]);
        vb[c0+1] = fmaf(x1, W0.y, vb[c0+1]);
        vb[c0+2] = fmaf(x1, W0.z, vb[c0+2]);
        vb[c0+3] = fmaf(x1, W0.w, vb[c0+3]);
        vc[c0+0] = fmaf(x1, W1.x, vc[c0+0]);
        vc[c0+1] = fmaf(x1, W1.y, vc[c0+1]);
        vc[c0+2] = fmaf(x1, W1.z, vc[c0+2]);
        vc[c0+3] = fmaf(x1, W1.w, vc[c0+3]);
        vb[c0+0] = fmaf(x0, W2.x, vb[c0+0]);
        vb[c0+1] = fmaf(x0, W2.y, vb[c0+1]);
        vb[c0+2] = fmaf(x0, W2.z, vb[c0+2]);
        vb[c0+3] = fmaf(x0, W2.w, vb[c0+3]);
        vd[c0+0] = fmaf(x2, W0.x, vd[c0+0]);
        vd[c0+1] = fmaf(x2, W0.y, vd[c0+1]);
        vd[c0+2] = fmaf(x2, W0.z, vd[c0+2]);
        vd[c0+3] = fmaf(x2, W0.w, vd[c0+3]);
        vd[c0+0] = fmaf(x1, W2.x, vd[c0+0]);
        vd[c0+1] = fmaf(x1, W2.y, vd[c0+1]);
        vd[c0+2] = fmaf(x1, W2.z, vd[c0+2]);
        vd[c0+3] = fmaf(x1, W2.w, vd[c0+3]);
    }
}

template <int PH>
__device__ __forceinline__ void htaps_spec(
    const float* __restrict__ xd, const float* wd, const int ih,
    const int wp, const int i1, const int i2,
    float* va, float* vb, float* vc, float* vd)
{
    constexpr int IW = 32;
    if (PH == 0) {
        tapfma_spec(xd + ih * IW, wd + 48, wp, i1, i2, va, vb, vc, vd);  // kh=1
    } else {
        tapfma_spec(xd + (ih + 1) * IW, wd,        wp, i1, i2, va, vb, vc, vd);  // kh=0
        tapfma_spec(xd + ih * IW,       wd + 96,   wp, i1, i2, va, vb, vc, vd);  // kh=2
    }
}

template <int PD, int PH>
__device__ __forceinline__ void dtaps_spec(
    const float* __restrict__ xc, const float* wci, const int id, const int ih,
    const int wp, const int i1, const int i2,
    float* va, float* vb, float* vc, float* vd)
{
    constexpr int SLICE = 32 * 32;
    if (PD == 0) {
        htaps_spec<PH>(xc + id * SLICE, wci + 144, ih, wp, i1, i2, va, vb, vc, vd);  // kd=1
    } else {
        htaps_spec<PH>(xc + (id + 1) * SLICE, wci,       ih, wp, i1, i2, va, vb, vc, vd);  // kd=0
        htaps_spec<PH>(xc + id * SLICE,       wci + 288, ih, wp, i1, i2, va, vb, vc, vd);  // kd=2
    }
}

template <int PD, int PH>
__global__ __launch_bounds__(256)
void fused_kernel_spec(const float* __restrict__ x, const float* __restrict__ wr,
                       const float* __restrict__ cb, const float* __restrict__ sbp,
                       float* __restrict__ out, const int N) {
    constexpr int ID = 16, IH = 32, IW = 32;
    constexpr int OD = 31, OH = 63, OW = 63;
    constexpr int NQ = 16;
    constexpr int Dp = PD ? ID - 1 : ID;
    constexpr int Hp = PH ? IH - 1 : IH;

    __shared__ __align__(16) float ws[REPACK];
    for (int i = threadIdx.x; i < REPACK; i += 256) ws[i] = wr[i];
    __syncthreads();

    const int idx = blockIdx.x * 256 + threadIdx.x;
    if (idx >= N * (Dp * Hp * NQ)) return;

    const int q = idx & (NQ - 1);
    int t = idx >> 4;
    const int ih = t % Hp; t /= Hp;
    const int id = t % Dp;
    const int n  = t / Dp;
    const int od = 2 * id + PD;
    const int oh = 2 * ih + PH;

    const int wp = 2 * q;
    const int i1 = wp + 1;              // always <= IW-1 for q < 16
    const int i2 = min(wp + 2, IW - 1); // clamped only for q = 15 (vd discarded)
    const bool has_d = (q < NQ - 1);

    float va[CO], vb[CO], vc[CO], vd[CO];
    #pragma unroll
    for (int co = 0; co < CO; ++co) {
        const float b = __ldg(cb + co);
        va[co] = b; vb[co] = b; vc[co] = b; vd[co] = b;
    }

    const float* xn = x + static_cast<long long>(n) * (CI * ID * IH * IW);
    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        dtaps_spec<PD, PH>(xn + ci * (ID * IH * IW), ws + ci * 432,
                           id, ih, wp, i1, i2, va, vb, vc, vd);
    }

    const float sb = __ldg(sbp);
    const long long oi =
        (static_cast<long long>(n * OD + od) * OH + oh) * OW + 4 * q;
    out[oi]     = epilogue(va, sb);
    out[oi + 1] = epilogue(vb, sb);
    out[oi + 2] = epilogue(vc, sb);
    if (has_d) out[oi + 3] = epilogue(vd, sb);
}

// ---------------------------------------------------------------------------
// Generic fallback (any ID/IH/IW), same fusion with runtime indices.
// ---------------------------------------------------------------------------
__device__ __forceinline__ void body_gen(
    const float* __restrict__ x,
    const float* __restrict__ wr,
    const float* __restrict__ conv_bias,
    const float* __restrict__ scalar_bias,
    float* __restrict__ out,
    const int N, const int ID, const int IH, const int IW)
{
    const int OD = 2 * ID - 1;
    const int OH = 2 * IH - 1;
    const int OW = 2 * IW - 1;
    const int NQ = (IW + 1) >> 1;

    __shared__ __align__(16) float ws[REPACK];
    for (int i = threadIdx.x; i < REPACK; i += blockDim.x) ws[i] = wr[i];
    __syncthreads();

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * OD * OH * NQ) return;

    const int q  = idx % NQ;
    int t = idx / NQ;
    const int oh = t % OH; t /= OH;
    const int od = t % OD;
    const int n  = t / OD;

    const int wp = 2 * q;
    const int i1 = min(wp + 1, IW - 1);
    const int i2 = min(wp + 2, IW - 1);
    const bool has_p1 = (wp + 1) < IW;
    const bool has_d  = (wp + 2) < IW;

    float va[CO], vb[CO], vc[CO], vd[CO];
    #pragma unroll
    for (int co = 0; co < CO; ++co) {
        const float b = __ldg(conv_bias + co);
        va[co] = b; vb[co] = b; vc[co] = b; vd[co] = b;
    }

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

    const int kd0 = (od + 1) & 1;
    const int kh0 = (oh + 1) & 1;

    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        const float* xc = xn + ci * DHW;
        const float* wt = ws + ci * (27 * CO);
        for (int kd = kd0; kd < 3; kd += 2) {
            const int id = (od + 1 - kd) >> 1;
            const float* xd = xc + id * HW;
            const float* wd = wt + kd * (9 * CO);
            for (int kh = kh0; kh < 3; kh += 2) {
                const int ih = (oh + 1 - kh) >> 1;
                const float* xh = xd + ih * IW;
                const float4* wq = reinterpret_cast<const float4*>(wd + kh * (3 * CO));

                const float x0 = __ldg(xh + wp);
                const float x1 = __ldg(xh + i1);
                const float x2 = __ldg(xh + i2);

                #pragma unroll
                for (int cg = 0; cg < 4; ++cg) {
                    const float4 W0 = wq[cg];
                    const float4 W1 = wq[4 + cg];
                    const float4 W2 = wq[8 + cg];
                    const int c0 = 4 * cg;
                    va[c0+0] = fmaf(x0, W1.x, va[c0+0]);
                    va[c0+1] = fmaf(x0, W1.y, va[c0+1]);
                    va[c0+2] = fmaf(x0, W1.z, va[c0+2]);
                    va[c0+3] = fmaf(x0, W1.w, va[c0+3]);
                    vb[c0+0] = fmaf(x1, W0.x, vb[c0+0]);
                    vb[c0+1] = fmaf(x1, W0.y, vb[c0+1]);
                    vb[c0+2] = fmaf(x1, W0.z, vb[c0+2]);
                    vb[c0+3] = fmaf(x1, W0.w, vb[c0+3]);
                    vc[c0+0] = fmaf(x1, W1.x, vc[c0+0]);
                    vc[c0+1] = fmaf(x1, W1.y, vc[c0+1]);
                    vc[c0+2] = fmaf(x1, W1.z, vc[c0+2]);
                    vc[c0+3] = fmaf(x1, W1.w, vc[c0+3]);
                    vb[c0+0] = fmaf(x0, W2.x, vb[c0+0]);
                    vb[c0+1] = fmaf(x0, W2.y, vb[c0+1]);
                    vb[c0+2] = fmaf(x0, W2.z, vb[c0+2]);
                    vb[c0+3] = fmaf(x0, W2.w, vb[c0+3]);
                    vd[c0+0] = fmaf(x2, W0.x, vd[c0+0]);
                    vd[c0+1] = fmaf(x2, W0.y, vd[c0+1]);
                    vd[c0+2] = fmaf(x2, W0.z, vd[c0+2]);
                    vd[c0+3] = fmaf(x2, W0.w, vd[c0+3]);
                    vd[c0+0] = fmaf(x1, W2.x, vd[c0+0]);
                    vd[c0+1] = fmaf(x1, W2.y, vd[c0+1]);
                    vd[c0+2] = fmaf(x1, W2.z, vd[c0+2]);
                    vd[c0+3] = fmaf(x1, W2.w, vd[c0+3]);
                }
            }
        }
    }

    const float sb = __ldg(scalar_bias);
    const long long oi =
        (static_cast<long long>(n * OD + od) * OH + oh) * OW + 4 * q;

    out[oi] = epilogue(va, sb);
    if (has_p1) {
        out[oi + 1] = epilogue(vb, sb);
        out[oi + 2] = epilogue(vc, sb);
        if (has_d) out[oi + 3] = epilogue(vd, sb);
    }
}

__global__ __launch_bounds__(256)
void fused_kernel_gen(const float* __restrict__ x, const float* __restrict__ wr,
                      const float* __restrict__ cb, const float* __restrict__ sb,
                      float* __restrict__ out,
                      const int N, const int ID, const int IH, const int IW) {
    body_gen(x, wr, cb, sb, out, N, ID, IH, IW);
}

}  // 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");

    const c10::cuda::CUDAGuard device_guard(x.device());

    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 == CI && COv == CO, "unsupported channel configuration");
    TORCH_CHECK(wc.numel() == CI * CO * 27, "unsupported kernel size");
    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 = 2 * ID - 1;
    const int OH = 2 * IH - 1;
    const int OW = 2 * IW - 1;

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

    auto stream = at::cuda::getCurrentCUDAStream();
    cudaStream_t s = stream.stream();

    repack_kernel<<<(REPACK + 255) / 256, 256, 0, s>>>(
        wc.data_ptr<float>(), wr.data_ptr<float>());

    if (ID == 16 && IH == 32 && IW == 32) {
        // grid.y selects the (depth, height) parity class; grid.x covers the
        // largest class (even/even). Smaller classes retire extra threads.
        const int maxthreads = N * 16 * 32 * 16;
        dim3 grid((maxthreads + 255) / 256, 4);
        const float* xp = xc.data_ptr<float>();
        const float* wp = wr.data_ptr<float>();
        const float* bp = bc.data_ptr<float>();
        const float* sp = bs.data_ptr<float>();
        float* op = out.data_ptr<float>();
        fused_kernel_spec<0, 0><<<grid, 256, 0, s>>>(xp, wp, bp, sp, op, N);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        const int NQ = (IW + 1) / 2;
        const int total = N * OD * OH * NQ;
        const int grid = (total + 255) / 256;
        fused_kernel_gen<<<grid, 256, 0, s>>>(
            xc.data_ptr<float>(), wr.data_ptr<float>(), bc.data_ptr<float>(),
            bs.data_ptr<float>(), out.data_ptr<float>(), N, ID, IH, IW);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    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_v6",
    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. Each
    thread computes four adjacent output voxels so shared-memory weight loads
    and gathered input values are reused across voxels; for 16x32x32 inputs a
    parity-specialized fully-unrolled kernel is used.
    """
    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
```