```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 CUDA kernel:
#   ConvTranspose3d (groups=1, dilation=1) -> LogSumExp(dim=1) -> HardSwish
#   -> scalar bias subtraction -> clamp(-1, 1)
# This avoids materializing the huge (B, C_out, D, H, W) intermediate tensor,
# turning the whole chain into a single memory-light kernel.
# A second kernel fuses the post-processing ops (used as a fallback path).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

template <int OC>
__global__ __launch_bounds__(256) void fused_convT3d_post_kernel(
    const float* __restrict__ input,      // (B, IC, ID, IH, IW)
    const float* __restrict__ weight,     // (IC, OC, KD, KH, KW)
    const float* __restrict__ conv_bias,  // (OC,) or nullptr
    const float* __restrict__ post_bias,  // scalar
    float* __restrict__ out,              // (B, 1, OD, OH, OW)
    const int B, const int IC,
    const int ID, const int IH, const int IW,
    const int OD, const int OH, const int OW,
    const int KD, const int KH, const int KW,
    const int sd, const int sh, const int sw,
    const int pd, const int ph, const int pw)
{
    const long tid = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long total = (long)B * OD * OH * OW;
    if (tid >= total) return;

    const int w = (int)(tid % OW);
    long t = tid / OW;
    const int h = (int)(t % OH); t /= OH;
    const int d = (int)(t % OD);
    const int b = (int)(t / OD);

    const int KVOL = KD * KH * KW;
    const long in_spatial = (long)ID * IH * IW;
    const float* in_base = input + (long)b * IC * in_spatial;

    // Accumulate every output channel for this spatial voxel in registers.
    float acc[OC];
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) acc[oc] = conv_bias ? conv_bias[oc] : 0.0f;

    // Transposed conv: out[d] receives in[id] * weight[kd] whenever
    // d = id*stride - pad + kd  <=>  id = (d + pad - kd) / stride (integer, in range).
    for (int kd = 0; kd < KD; ++kd) {
        const int nd = d + pd - kd;
        if (nd < 0 || (nd % sd) != 0) continue;
        const int id = nd / sd;
        if (id >= ID) continue;
        for (int kh = 0; kh < KH; ++kh) {
            const int nh = h + ph - kh;
            if (nh < 0 || (nh % sh) != 0) continue;
            const int ih = nh / sh;
            if (ih >= IH) continue;
            const long in_off_dh = ((long)id * IH + ih) * IW;
            const int koff_dh = (kd * KH + kh) * KW;
            for (int kw = 0; kw < KW; ++kw) {
                const int nw = w + pw - kw;
                if (nw < 0 || (nw % sw) != 0) continue;
                const int iw = nw / sw;
                if (iw >= IW) continue;

                const long in_off = in_off_dh + iw;
                const int koff = koff_dh + kw;
                for (int ic = 0; ic < IC; ++ic) {
                    const float v = __ldg(in_base + (long)ic * in_spatial + in_off);
                    const float* wptr = weight + (long)ic * OC * KVOL + koff;
                    #pragma unroll
                    for (int oc = 0; oc < OC; ++oc) {
                        acc[oc] = fmaf(v, __ldg(wptr + (long)oc * KVOL), acc[oc]);
                    }
                }
            }
        }
    }

    // LogSumExp over the channel accumulators (numerically stable).
    float max_val = acc[0];
    #pragma unroll
    for (int oc = 1; oc < OC; ++oc) max_val = fmaxf(max_val, acc[oc]);
    float sum = 0.0f;
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) sum += expf(acc[oc] - max_val);
    const float lse = max_val + logf(sum);

    // HardSwish exactly as written in the reference: x * sigmoid(x + 3) / 6
    const float sig = 1.0f / (1.0f + expf(-(lse + 3.0f)));
    float val = lse * sig / 6.0f;

    // subtract scalar bias, clamp to [-1, 1]
    val -= post_bias[0];
    val = fminf(fmaxf(val, -1.0f), 1.0f);

    out[tid] = val;
}

// Fused post-processing: LogSumExp(dim=1) + HardSwish + bias + clamp.
__global__ void fused_post_kernel(
    const float* __restrict__ x,      // (B, C, D, H, W)
    const float* __restrict__ post_bias,
    float* __restrict__ out,          // (B, 1, D, H, W)
    const int C, const long DHW, const long N)
{
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;

    const long b = idx / DHW;
    const long dhw = idx - b * DHW;
    const float* ptr = x + b * (long)C * DHW + dhw;

    float max_val = -INFINITY;
    for (int c = 0; c < C; ++c) max_val = fmaxf(max_val, ptr[(long)c * DHW]);
    float sum = 0.0f;
    for (int c = 0; c < C; ++c) sum += expf(ptr[(long)c * DHW] - max_val);
    const float lse = max_val + logf(sum);

    const float sig = 1.0f / (1.0f + expf(-(lse + 3.0f)));
    float val = lse * sig / 6.0f;
    val -= post_bias[0];
    val = fminf(fmaxf(val, -1.0f), 1.0f);
    out[idx] = val;
}

torch::Tensor fused_conv_transpose3d_post_cuda(
    torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor post_bias,
    int64_t sd, int64_t sh, int64_t sw, int64_t pd, int64_t ph, int64_t pw)
{
    auto x = input.contiguous();
    auto wgt = weight.contiguous();
    auto pb = post_bias.contiguous();

    const int B = x.size(0), IC = x.size(1);
    const int ID = x.size(2), IH = x.size(3), IW = x.size(4);
    const int OC = wgt.size(1);
    const int KD = wgt.size(2), KH = wgt.size(3), KW = wgt.size(4);

    const int OD = (ID - 1) * (int)sd - 2 * (int)pd + KD;
    const int OH = (IH - 1) * (int)sh - 2 * (int)ph + KH;
    const int OW = (IW - 1) * (int)sw - 2 * (int)pw + KW;

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

    const float* cb_ptr = nullptr;
    torch::Tensor cb;
    if (conv_bias.defined() && conv_bias.numel() > 0) {
        cb = conv_bias.contiguous();
        cb_ptr = cb.data_ptr<float>();
    }

    const long N = (long)B * OD * OH * OW;
    const int block = 256;
    const unsigned int grid = (unsigned int)((N + block - 1) / block);
    auto stream = at::cuda::getCurrentCUDAStream();

    #define LAUNCH_CASE(OCV)                                                          \
        fused_convT3d_post_kernel<OCV><<<grid, block, 0, stream>>>(                   \
            x.data_ptr<float>(), wgt.data_ptr<float>(), cb_ptr, pb.data_ptr<float>(), \
            out.data_ptr<float>(), B, IC, ID, IH, IW, OD, OH, OW, KD, KH, KW,         \
            (int)sd, (int)sh, (int)sw, (int)pd, (int)ph, (int)pw)

    switch (OC) {
        case 4:  LAUNCH_CASE(4);  break;
        case 8:  LAUNCH_CASE(8);  break;
        case 16: LAUNCH_CASE(16); break;
        case 32: LAUNCH_CASE(32); break;
        default: TORCH_CHECK(false, "Unsupported out_channels for fused kernel"); break;
    }
    #undef LAUNCH_CASE

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

torch::Tensor fused_post_cuda(torch::Tensor x_in, torch::Tensor post_bias)
{
    auto x = x_in.contiguous();
    auto pb = post_bias.contiguous();

    const int B = x.size(0);
    const int C = x.size(1);
    long DHW = 1;
    for (int i = 2; i < x.dim(); ++i) DHW *= x.size(i);
    const long N = (long)B * DHW;

    auto out = torch::empty({B, 1, x.size(2), x.size(3), x.size(4)}, x.options());

    const int block = 256;
    const unsigned int grid = (unsigned int)((N + block - 1) / block);
    auto stream = at::cuda::getCurrentCUDAStream();

    fused_post_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), pb.data_ptr<float>(), out.data_ptr<float>(), C, DHW, N);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

cpp_source = """
torch::Tensor fused_conv_transpose3d_post_cuda(
    torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor post_bias,
    int64_t sd, int64_t sh, int64_t sw, int64_t pd, int64_t ph, int64_t pw);
torch::Tensor fused_post_cuda(torch::Tensor x_in, torch::Tensor post_bias);
"""

fused_ops = load_inline(
    name="fused_convT3d_post_ops",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_conv_transpose3d_post_cuda", "fused_post_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the 3D transposed convolution, LogSumExp, HardSwish,
    bias subtraction and clamp are all fused into a single custom CUDA kernel,
    eliminating the large intermediate activation tensor entirely.
    """
    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(*bias_shape))
        self.out_channels = out_channels

    def forward(self, x):
        ct = self.conv_transpose
        use_fused = (
            x.is_cuda and x.dim() == 5 and x.dtype == torch.float32
            and ct.weight.dtype == torch.float32
            and self.bias.numel() == 1
            and ct.groups == 1
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
        )
        if use_fused and self.out_channels in (4, 8, 16, 32):
            conv_bias = ct.bias if ct.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
            s = ct.stride
            p = ct.padding
            return fused_ops.fused_conv_transpose3d_post_cuda(
                x, ct.weight, conv_bias, self.bias,
                s[0], s[1], s[2], p[0], p[1], p[2])
        if use_fused:
            y = ct(x)
            return fused_ops.fused_post_cuda(y, self.bias)
        # Generic fallback (e.g. CPU tensors).
        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
```