```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 (v3):
#   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])
#
# Key ideas:
#   * The (N, 16, D, H, W) transposed-conv result is never materialized; each
#     thread accumulates all 16 channels for two horizontally adjacent output
#     voxels (even/odd parity) directly in registers.
#   * All convolution weights / biases live in __constant__ memory. Because a
#     warp follows a uniform kernel-tap schedule, the weights are fetched via
#     constant-cache broadcasts (often folded directly into FMA operands),
#     completely eliminating the shared-memory load traffic that dominated the
#     previous versions.
#   * The tap schedule is specialized on the (od, oh) parity with warp-uniform
#     branches, so every constant-memory index is a compile-time immediate.
#   * A template specialization for ID=16, IH=32, IW=32 makes all address
#     arithmetic and thread decomposition compile-time (shifts / magic mults).
#   * Fast SFU transcendentals (__expf/__logf) in the epilogue.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

namespace {

constexpr int CI = 3;
constexpr int CO = 16;

__constant__ float CW[CI * CO * 27];  // conv weight layout [ci][co][kd][kh][kw]
__constant__ float CB[CO];            // conv bias
__constant__ float SB[1];             // scalar bias

__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 = __frcp_rn(1.f + __expf(-(lse + 3.f)));
    const float y = fmaf(lse, sig * 0.16666666666666666f, -sb);
    return fminf(fmaxf(y, -1.f), 1.f);
}

// Accumulate the transposed-conv taps for a pair of adjacent voxels.
//   even voxel: (od, oh, 2*wp)     gathers (id, ih, wp)   with kw = 1
//   odd voxel:  (od, oh, 2*wp + 1) gathers (id, ih, wp+1) with kw = 0
//                                    and (id, ih, wp)     with kw = 2
// PD/PH encode the parity of od/oh: parity 0 -> tap index 1 only,
// parity 1 -> tap indices {0, 2}. All loops are fully unrolled so every
// constant-memory offset is an immediate.
template <int PD, int PH>
__device__ __forceinline__ void compute_pair(
    const float* __restrict__ xn,
    const int ID, const int IH, const int IW,
    const int od, const int oh, const int wp, const bool has1,
    float (&v0)[CO], float (&v1)[CO])
{
    constexpr int KDA = PD ? 0 : 1;
    constexpr int KDB = PD ? 2 : 1;
    constexpr int NKD = PD ? 2 : 1;
    constexpr int KHA = PH ? 0 : 1;
    constexpr int KHB = PH ? 2 : 1;
    constexpr int NKH = PH ? 2 : 1;

    const int HW  = IH * IW;
    const int DHW = ID * HW;

    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        const float* xc = xn + ci * DHW;
        #pragma unroll
        for (int kdi = 0; kdi < NKD; ++kdi) {
            const int kd = (kdi == 0) ? KDA : KDB;
            const int id = (od + 1 - kd) >> 1;   // always in [0, ID)
            const float* xd = xc + id * HW;
            #pragma unroll
            for (int khi = 0; khi < NKH; ++khi) {
                const int kh = (khi == 0) ? KHA : KHB;
                const int ih = (oh + 1 - kh) >> 1;   // always in [0, IH)
                const float* p = xd + ih * IW + wp;
                const float x0 = __ldg(p);
                const float* wpv = CW + ci * (CO * 27) + kd * 9 + kh * 3;
                #pragma unroll
                for (int co = 0; co < CO; ++co)
                    v0[co] = fmaf(x0, wpv[co * 27 + 1], v0[co]);
                if (has1) {
                    const float x1 = __ldg(p + 1);
                    #pragma unroll
                    for (int co = 0; co < CO; ++co)
                        v1[co] = fmaf(x1, wpv[co * 27], v1[co]);
                    #pragma unroll
                    for (int co = 0; co < CO; ++co)
                        v1[co] = fmaf(x0, wpv[co * 27 + 2], v1[co]);
                }
            }
        }
    }
}

__device__ __forceinline__ void body(const float* __restrict__ x,
                                     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 OW2 = IW;   // one thread covers voxels (2*wp, 2*wp+1)

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * OD * OH * OW2) 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 + 1) < IW;

    float v0[CO], v1[CO];
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        const float b = CB[c];
        v0[c] = b;
        v1[c] = b;
    }

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

    const int pd = od & 1;
    const int ph = oh & 1;
    if (pd == 0) {
        if (ph == 0) compute_pair<0, 0>(xn, ID, IH, IW, od, oh, wp, has1, v0, v1);
        else         compute_pair<0, 1>(xn, ID, IH, IW, od, oh, wp, has1, v0, v1);
    } else {
        if (ph == 0) compute_pair<1, 0>(xn, ID, IH, IW, od, oh, wp, has1, v0, v1);
        else         compute_pair<1, 1>(xn, ID, IH, IW, od, oh, wp, has1, v0, v1);
    }

    const float sb = SB[0];
    const long long oi =
        (static_cast<long long>(n * OD + od) * OH + oh) * OW + 2 * wp;
    const float y0 = epilogue(v0, sb);
    if (has1) {
        const float y1 = epilogue(v1, sb);
        *reinterpret_cast<float2*>(out + oi) = make_float2(y0, y1);
    } else {
        out[oi] = y0;
    }
}

template <int ID, int IH, int IW>
__global__ __launch_bounds__(256)
void fused_kernel_spec(const float* __restrict__ x, float* __restrict__ out,
                       const int N) {
    body(x, out, N, ID, IH, IW);
}

__global__ __launch_bounds__(256)
void fused_kernel_gen(const float* __restrict__ x, float* __restrict__ out,
                      const int N, const int ID, const int IH, const int IW) {
    body(x, 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");

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

    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 stream = at::cuda::getCurrentCUDAStream();

    // stage weights / biases in constant memory (device-to-device, async)
    cudaMemcpyToSymbolAsync(CW, wc.data_ptr<float>(), sizeof(float) * CI * CO * 27,
                            0, cudaMemcpyDeviceToDevice, stream.stream());
    cudaMemcpyToSymbolAsync(CB, bc.data_ptr<float>(), sizeof(float) * CO,
                            0, cudaMemcpyDeviceToDevice, stream.stream());
    cudaMemcpyToSymbolAsync(SB, bs.data_ptr<float>(), sizeof(float),
                            0, cudaMemcpyDeviceToDevice, stream.stream());

    const int total = N * OD * OH * IW;
    const int grid = (total + 255) / 256;

    if (ID == 16 && IH == 32 && IW == 32) {
        fused_kernel_spec<16, 32, 32><<<grid, 256, 0, stream.stream()>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(), N);
    } else {
        fused_kernel_gen<<<grid, 256, 0, stream.stream()>>>(
            xc.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_v3",
    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. Weights
    are staged in constant memory so they are consumed as broadcast FMA
    operands with no shared/global load traffic in the inner loop.
    """
    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
```