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

# Two execution paths are compiled:
#  1) safe_forward: identical cuDNN ConvTranspose3d call as the reference plus a
#     single fused kernel doing LogSumExp(dim=1) + sigmoid-HardSwish + bias + clamp
#     (the large conv output is read from global memory exactly once).
#  2) custom_forward: a fully fused kernel that computes the transposed
#     convolution itself (gather formulation, FP32 FMA accumulation, optional
#     TF32-rounded multiplicands to mirror cuDNN tensor-core behavior) and
#     immediately reduces over channels with LogSumExp + post ops, writing only
#     the tiny (B,1,Do,Ho,Wo) output. This avoids materializing the huge
#     (B,16,Do,Ho,Wo) intermediate entirely (~1GB of traffic saved).
# At runtime, the first forward for a given configuration validates the custom
# path against the safe path; if they do not match tightly (e.g. unexpected
# accumulation semantics), the safe path is used permanently.

fused_cpp_source = (
    "torch::Tensor safe_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias, "
    "std::vector<int64_t> stride, std::vector<int64_t> padding);\n"
    "torch::Tensor custom_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias, bool tf32);"
)

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

__device__ __forceinline__ float round_tf32(float f) {
    // Round to TF32 (10 explicit mantissa bits), ties away from zero
    // (matches `cvt.rna.tf32.f32` used when feeding tensor cores).
    unsigned u = __float_as_uint(f);
    u = (u + 0x1000u) & 0xFFFFE000u;
    return __uint_as_float(u);
}

// ---------------------------------------------------------------------------
// Fused post-processing kernels: LogSumExp over channel dim + HardSwish
// (sigmoid formulation) + bias subtraction + clamp. One global read of the
// conv output, one write of the final 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 + (size_t)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)));
}

__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)));
}

// ---------------------------------------------------------------------------
// Fully fused ConvTranspose3d (K=3, stride=2, padding=1, groups=1, dilation=1)
// + LogSumExp over output channels + HardSwish + bias + clamp.
// One thread per output voxel; all COUT channel accumulators live in
// registers; weights staged in shared memory. Never materializes the
// (B, COUT, Do, Ho, Wo) intermediate.
// ---------------------------------------------------------------------------
template <int CIN, int COUT, bool TF32>
__global__ __launch_bounds__(256) void fused_convt3d_lse_post_kernel(
        const float* __restrict__ x,      // (B, CIN, Di, Hi, Wi)
        const float* __restrict__ w,      // (CIN, COUT, 3, 3, 3)
        const float* __restrict__ cbias,  // (COUT)
        const float* __restrict__ bias,   // scalar
        float* __restrict__ out,          // (B, 1, Do, Ho, Wo)
        int Di, int Hi, int Wi,
        int total) {
    __shared__ float sw[CIN * COUT * 27];
    __shared__ float scb[COUT];
    for (int i = threadIdx.x; i < CIN * COUT * 27; i += 256) {
        float v = w[i];
        sw[i] = TF32 ? round_tf32(v) : v;
    }
    if (threadIdx.x < COUT) scb[threadIdx.x] = cbias[threadIdx.x];
    __syncthreads();

    int t = blockIdx.x * 256 + threadIdx.x;
    if (t >= total) return;

    const int HiWi = Hi * Wi;
    const int Ho = (Hi - 1) * 2 - 2 + 3;  // (Hi-1)*s - 2p + K
    const int Wo = (Wi - 1) * 2 - 2 + 3;
    const int Do = (Di - 1) * 2 - 2 + 3;

    int ow = t % Wo;
    int q = t / Wo;
    int oh = q % Ho;
    q = q / Ho;
    int od = q % Do;
    int b = q / Do;

    // Valid kernel offsets per dim: kd must satisfy (od + p - kd) % s == 0.
    // With s=2, p=1: parity of kd fixed by (od+1)&1; if that parity is 0,
    // both {0, 2} are valid, otherwise only {1}. Resulting input indices are
    // always within bounds given the output-size formula.
    const int kdA = (od + 1) & 1;
    const int idA = (od + 1 - kdA) >> 1;
    const int nd = (kdA == 0) ? 2 : 1;
    const int khA = (oh + 1) & 1;
    const int ihA = (oh + 1 - khA) >> 1;
    const int nh = (khA == 0) ? 2 : 1;
    const int kwA = (ow + 1) & 1;
    const int iwA = (ow + 1 - kwA) >> 1;
    const int nw = (kwA == 0) ? 2 : 1;

    float vals[COUT];
    #pragma unroll
    for (int oc = 0; oc < COUT; ++oc) vals[oc] = scb[oc];

    const int plane = Di * HiWi;
    const float* xb = x + (size_t)b * CIN * plane;

    const int kbase = kdA * 9 + khA * 3 + kwA;
    const int ibase = (idA * Hi + ihA) * Wi + iwA;

    for (int a = 0; a < nd; ++a) {
        for (int bb = 0; bb < nh; ++bb) {
            for (int c = 0; c < nw; ++c) {
                // offsets (kdA+2a, khA+2bb, kwA+2c), inputs (idA-a, ihA-bb, iwA-c)
                const int koff = kbase + 18 * a + 6 * bb + 2 * c;
                const int ioff = ibase - a * HiWi - bb * Wi - c;
                #pragma unroll
                for (int ic = 0; ic < CIN; ++ic) {
                    float xv = __ldg(xb + ic * plane + ioff);
                    if (TF32) xv = round_tf32(xv);
                    const float* wp = sw + (ic * COUT) * 27 + koff;
                    #pragma unroll
                    for (int oc = 0; oc < COUT; ++oc) {
                        vals[oc] = fmaf(xv, wp[oc * 27], vals[oc]);
                    }
                }
            }
        }
    }

    float m = vals[0];
    #pragma unroll
    for (int oc = 1; oc < COUT; ++oc) m = fmaxf(m, vals[oc]);
    float sum = 0.f;
    #pragma unroll
    for (int oc = 0; oc < COUT; ++oc) sum += expf(vals[oc] - m);
    float lse = m + logf(sum);
    out[t] = post_op(lse, __ldg(bias));
}

// ---------------------------------------------------------------------------
// Safe path: exact same ATen/cuDNN conv_transpose3d as the reference module
// (output_padding=0, groups=1, dilation=1) followed by the fused post kernel.
// ---------------------------------------------------------------------------
torch::Tensor safe_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();
    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;
}

// ---------------------------------------------------------------------------
// Fast path: fully fused transposed conv + all post ops.
// ---------------------------------------------------------------------------
torch::Tensor custom_forward(torch::Tensor x, torch::Tensor weight,
                             torch::Tensor conv_bias, torch::Tensor bias,
                             bool tf32) {
    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(x.size(1) == 3, "custom kernel expects 3 input channels");
    TORCH_CHECK(weight.dim() == 5 && weight.size(0) == 3 && weight.size(1) == 16,
                "custom kernel expects weight shape (3, 16, 3, 3, 3)");
    TORCH_CHECK(bias.numel() == 1, "bias must be scalar-like");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto cc = conv_bias.contiguous();
    auto bc = bias.contiguous();

    int64_t B = xc.size(0);
    int64_t Di = xc.size(2), Hi = xc.size(3), Wi = xc.size(4);
    int64_t Do = (Di - 1) * 2 - 2 + 3;
    int64_t Ho = (Hi - 1) * 2 - 2 + 3;
    int64_t Wo = (Wi - 1) * 2 - 2 + 3;
    int64_t total64 = B * Do * Ho * Wo;
    TORCH_CHECK(total64 < (int64_t)INT_MAX, "output too large");
    int total = (int)total64;

    auto out = at::empty({B, 1, Do, Ho, Wo}, xc.options());

    const int block = 256;
    int grid = (total + block - 1) / block;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (tf32) {
        fused_convt3d_lse_post_kernel<3, 16, true><<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), cc.data_ptr<float>(),
            bc.data_ptr<float>(), out.data_ptr<float>(),
            (int)Di, (int)Hi, (int)Wi, total);
    } else {
        fused_convt3d_lse_post_kernel<3, 16, false><<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), cc.data_ptr<float>(),
            bc.data_ptr<float>(), out.data_ptr<float>(),
            (int)Di, (int)Hi, (int)Wi, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

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


class ModelNew(nn.Module):
    """
    Optimized model. Fast path: a single fully-fused CUDA kernel computing
    ConvTranspose3d + LogSumExp + HardSwish + bias + clamp without ever
    materializing the large intermediate. Safe path: the exact same cuDNN
    ConvTranspose3d call as the reference plus one fused kernel for all post
    ops. On the first forward for a given configuration the fast path is
    validated against the safe path; if it does not match tightly, the safe
    path is used permanently.
    """
    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))
        k = _triple(kernel_size)
        self._custom_ok = (
            in_channels == 3 and out_channels == 16 and
            k == (3, 3, 3) and tuple(self._stride) == (2, 2, 2) and
            tuple(self._padding) == (1, 1, 1)
        )
        self._choice = {}

    def _safe(self, x):
        return fused_ops.safe_forward(
            x, self.conv_transpose.weight, self.conv_transpose.bias,
            self.bias, self._stride, self._padding)

    def forward(self, x):
        if not (x.is_cuda and x.dtype == torch.float32) or self.bias.numel() != 1:
            y = self.conv_transpose(x)
            y = torch.logsumexp(y, dim=1, keepdim=True)
            y = y * torch.sigmoid(y + 3) / 6
            y = y - self.bias
            return torch.clamp(y, min=-1, max=1)

        if not self._custom_ok:
            return self._safe(x)

        w = self.conv_transpose.weight
        cb = self.conv_transpose.bias
        key = (tuple(x.shape), w._version, cb._version, self.bias._version,
               w.data_ptr(), cb.data_ptr())
        choice = self._choice.get(key, -1)

        if choice == 1:
            return fused_ops.custom_forward(x, w, cb, self.bias, False)
        if choice == 2:
            return fused_ops.custom_forward(x, w, cb, self.bias, True)
        if choice == 0:
            return self._safe(x)

        # First call for this configuration: validate the fast path.
        safe = self._safe(x)
        fast = fused_ops.custom_forward(x, w, cb, self.bias, False)
        if torch.allclose(fast, safe, rtol=5e-6, atol=5e-6):
            self._choice[key] = 1
            return fast
        fast_t = fused_ops.custom_forward(x, w, cb, self.bias, True)
        if torch.allclose(fast_t, safe, rtol=5e-6, atol=5e-6):
            self._choice[key] = 2
            return fast_t
        self._choice[key] = 0
        return safe
```