```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

# Strategy: keep the exact same cuDNN ConvTranspose3d call as the reference
# (its numerics must match bitwise), but call it directly from C++ to remove
# Python/dispatcher overhead, and fuse ALL subsequent ops (LogSumExp over
# channels + sigmoid-HardSwish + bias subtraction + clamp) into a single
# memory-bound CUDA kernel that reads the large conv output exactly once.
# Streaming loads/stores (__ldcs/__stcs) are used since the data has no reuse.

fused_cpp_source = (
    "torch::Tensor full_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias, "
    "std::vector<int64_t> stride, std::vector<int64_t> padding);"
)

fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <vector>
#include <climits>

__device__ __forceinline__ float post_op(float lse, float bias) {
    // x * sigmoid(x + 3) / 6  ->  subtract bias  ->  clamp(-1, 1)
    float sig = 1.f / (1.f + expf(-(lse + 3.f)));
    float y = lse * sig * (1.f / 6.f) - bias;
    return fminf(fmaxf(y, -1.f), 1.f);
}

// Channel count known at compile time: all channel values cached in registers,
// exactly one coalesced read of the input, one write of the output.
template <int C>
__global__ __launch_bounds__(256) void fused_lse_kernel(
        const float* __restrict__ x,
        const float* __restrict__ bias,
        float* __restrict__ out,
        int S) {
    int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= S) return;
    int b = blockIdx.y;
    const float* ptr = x + (size_t)b * C * S + s;

    float vals[C];
    float m = -INFINITY;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        float v = __ldcs(ptr + c * S);
        vals[c] = v;
        m = fmaxf(m, v);
    }
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        sum += expf(vals[c] - m);
    }
    float lse = m + logf(sum);
    __stcs(out + (size_t)b * S + s, post_op(lse, __ldg(bias)));
}

// Generic fallback for arbitrary channel counts (two passes over global memory).
__global__ void fused_lse_kernel_generic(const float* __restrict__ x,
                                         const float* __restrict__ bias,
                                         float* __restrict__ out,
                                         int C, int S) {
    int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= S) return;
    int b = blockIdx.y;
    const float* ptr = x + (size_t)b * C * S + s;

    float m = -INFINITY;
    for (int c = 0; c < C; ++c) m = fmaxf(m, __ldcs(ptr + (size_t)c * S));
    float sum = 0.f;
    for (int c = 0; c < C; ++c) sum += expf(__ldcs(ptr + (size_t)c * S) - m);
    float lse = m + logf(sum);
    __stcs(out + (size_t)b * S + s, post_op(lse, __ldg(bias)));
}

torch::Tensor full_forward(torch::Tensor x, torch::Tensor weight,
                           torch::Tensor conv_bias, torch::Tensor bias,
                           std::vector<int64_t> stride,
                           std::vector<int64_t> padding) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 5, "x must be 5D (B, C, D, H, W)");
    TORCH_CHECK(bias.numel() == 1, "bias must be scalar-like");

    auto xc = x.contiguous();

    // Identical ATen/cuDNN call as nn.ConvTranspose3d.forward
    // (output_padding=0, groups=1, dilation=1 are the module defaults).
    auto y = at::conv_transpose3d(xc, weight, conv_bias, stride, padding,
                                  {0, 0, 0}, 1, {1, 1, 1});
    auto yc = y.contiguous();

    int64_t B = yc.size(0);
    int64_t C = yc.size(1);
    int64_t S64 = yc.numel() / (B * C);
    TORCH_CHECK(S64 < (int64_t)INT_MAX, "spatial size too large");
    int S = (int)S64;

    auto out = at::empty({B, 1, yc.size(2), yc.size(3), yc.size(4)}, yc.options());
    auto bc = bias.contiguous();

    const float* xp = yc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    const int block = 256;
    dim3 grid((unsigned)((S + block - 1) / block), (unsigned)B);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    switch (C) {
        case 1:  fused_lse_kernel<1> <<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 2:  fused_lse_kernel<2> <<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 3:  fused_lse_kernel<3> <<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 4:  fused_lse_kernel<4> <<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 8:  fused_lse_kernel<8> <<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 16: fused_lse_kernel<16><<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 32: fused_lse_kernel<32><<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        case 64: fused_lse_kernel<64><<<grid, block, 0, stream>>>(xp, bp, op, S); break;
        default: fused_lse_kernel_generic<<<grid, block, 0, stream>>>(xp, bp, op, (int)C, S); break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_ops = load_inline(
    name="fused_convt3d_lse_post_v3",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["full_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the ConvTranspose3d (same cuDNN call as the reference, so
    numerics match) plus LogSumExp over channels, the sigmoid-based HardSwish,
    bias subtraction, and clamp are executed through a single fused C++/CUDA
    entry point. The large conv output is read from global memory exactly once.
    """
    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 _triple(v):
            if isinstance(v, (tuple, list)):
                return tuple(int(t) for t in v)
            return (int(v), int(v), int(v)

        self._stride = list(_triple(stride))
        self._padding = list(_triple(padding))

    def forward(self, x):
        if self.bias.numel() == 1 and x.is_cuda and x.dtype == torch.float32:
            return fused_ops.full_forward(
                x, self.conv_transpose.weight, self.conv_transpose.bias,
                self.bias, self._stride, self._padding)
        # Safe fallback matching the reference implementation.
        x = self.conv_transpose(x)
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = x * torch.sigmoid(x + 3) / 6
        x = x - self.bias
        return torch.clamp(x, min=-1, max=1)
```