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

# Fast path: fully fused ConvTranspose3d (K=3, stride=2, pad=1, Cin=3, Cout=16)
# + LogSumExp(dim=1) + sigmoid-HardSwish + bias + clamp in ONE kernel.
#   * Never materializes the huge (B,16,Do,Ho,Wo) intermediate.
#   * Each thread computes TWO same-parity output voxels spaced by 2 along W,
#     so every shared-memory weight load feeds 2 FMAs (halves shared traffic,
#     which was the bottleneck of the 1-voxel/thread version).
#   * Work ids are permuted so each warp handles a single W-parity
#     (no divergent tap counts inside a warp).
# Safe path: the exact reference cuDNN ConvTranspose3d + one fused post kernel.
# The first forward of each configuration cross-validates the fast path
# against the safe path and permanently falls back if it does not match.

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) {
    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 post_op_fast(float lse, float bias) {
    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) {
    unsigned u = __float_as_uint(f);
    u = (u + 0x1000u) & 0xFFFFE000u;
    return __uint_as_float(u);
}

// ---------------------------------------------------------------------------
// Safe-path post kernel: LogSumExp over C + HardSwish + bias + clamp.
// ---------------------------------------------------------------------------
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 kernel. Shared weights are stored oc-major:
//   sw[koff * 48 + ic * 16 + oc]  with koff = (kd*3 + kh)*3 + kw.
// Work-id decode (keeps warp parity-uniform):
//   W -> j = W % WP; t = W / WP; row = t % nRows; p = t / nRows;
//   row -> oh = row % Ho; od = (row / Ho) % Do; b = row / (Ho*Do).
// Thread handles voxels ow0 = 4j + p and ow1 = 4j + p + 2 (if < Wo).
// ---------------------------------------------------------------------------
template <bool TF32>
__global__ __launch_bounds__(256) void fused_convt3d_lse_post_v2_kernel(
        const float* __restrict__ x,      // (B, 3, Di, Hi, Wi)
        const float* __restrict__ w,      // (3, 16, 3, 3, 3)
        const float* __restrict__ cbias,  // (16)
        const float* __restrict__ bias,   // scalar
        float* __restrict__ out,          // (B, 1, Do, Ho, Wo)
        int Di, int Hi, int Wi,
        int Do, int Ho, int Wo,
        int nRows, int WP, int totalWork) {
    __shared__ float sw[27 * 48];
    __shared__ float scb[16];
    for (int i = threadIdx.x; i < 27 * 48; i += 256) {
        int koff = i / 48;
        int rem = i - koff * 48;
        int ic = rem >> 4;
        int oc = rem & 15;
        int kd = koff / 9;
        int kh = (koff - kd * 9) / 3;
        int kw = koff - kd * 9 - kh * 3;
        float v = w[(ic * 16 + oc) * 27 + koff];  // kd*9+kh*3+kw == koff
        (void)kd; (void)kh; (void)kw;
        sw[i] = TF32 ? round_tf32(v) : v;
    }
    if (threadIdx.x < 16) scb[threadIdx.x] = cbias[threadIdx.x];
    __syncthreads();

    int W = blockIdx.x * 256 + threadIdx.x;
    if (W >= totalWork) return;

    int j = W % WP;
    int t = W / WP;
    int row = t % nRows;
    int p = t / nRows;

    int oh = row % Ho;
    int t2 = row / Ho;
    int od = t2 % Do;
    int b = t2 / Do;

    const int ow0 = 4 * j + p;
    const int ow1 = ow0 + 2;
    const bool v0 = ow0 < Wo;
    const bool v1 = ow1 < Wo;
    if (!v0) return;

    // Valid kernel offsets per dim (s=2, p=1): parity of kd fixed by (od+1)&1.
    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 iw0 = 2 * j;  // see analysis: voxel 4j+p uses iw around 2j

    float acc0[16], acc1[16];
    #pragma unroll
    for (int oc = 0; oc < 16; ++oc) {
        acc0[oc] = scb[oc];
        acc1[oc] = scb[oc];
    }

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

    const float gbias = __ldg(bias);

    for (int a = 0; a < nd; ++a) {
        const int kd = kdA + 2 * a;
        const int id = idA - a;
        for (int bb = 0; bb < nh; ++bb) {
            const int kh = khA + 2 * bb;
            const int ih = ihA - bb;
            const int irow = (id * Hi + ih) * Wi + iw0;
            const int krow = (kd * 3 + kh) * 3;  // + kw
            if (p == 0) {
                const int koff = (krow + 1) * 48;
                #pragma unroll
                for (int ic = 0; ic < 3; ++ic) {
                    const float* xp = xb + ic * plane + irow;
                    float xv0 = __ldg(xp);
                    float xv1 = v1 ? __ldg(xp + 1) : 0.f;
                    if (TF32) { xv0 = round_tf32(xv0); xv1 = round_tf32(xv1); }
                    const float* wp = sw + koff + ic * 16;
                    #pragma unroll
                    for (int oc = 0; oc < 16; ++oc) {
                        float wv = wp[oc];
                        acc0[oc] = fmaf(xv0, wv, acc0[oc]);
                        acc1[oc] = fmaf(xv1, wv, acc1[oc]);
                    }
                }
            } else {
                const int koff0 = (krow + 0) * 48;
                const int koff2 = (krow + 2) * 48;
                #pragma unroll
                for (int ic = 0; ic < 3; ++ic) {
                    const float* xp = xb + ic * plane + irow;
                    float xv0 = __ldg(xp);
                    float xv1 = __ldg(xp + 1);
                    float xv2 = v1 ? __ldg(xp + 2) : 0.f;
                    if (TF32) { xv0 = round_tf32(xv0); xv1 = round_tf32(xv1); xv2 = round_tf32(xv2); }
                    const float* wp0 = sw + koff0 + ic * 16;
                    const float* wp2 = sw + koff2 + ic * 16;
                    #pragma unroll
                    for (int oc = 0; oc < 16; ++oc) {
                        float wl = wp0[oc];
                        float wh = wp2[oc];
                        acc0[oc] = fmaf(xv0, wh, fmaf(xv1, wl, acc0[oc]));
                        acc1[oc] = fmaf(xv1, wh, fmaf(xv2, wl, acc1[oc]));
                    }
                }
            }
        }
    }

    const int orow = ((b * Do + od) * Ho + oh) * Wo;
    {
        float m = acc0[0];
        #pragma unroll
        for (int oc = 1; oc < 16; ++oc) m = fmaxf(m, acc0[oc]);
        float sum = 0.f;
        #pragma unroll
        for (int oc = 0; oc < 16; ++oc) sum += __expf(acc0[oc] - m);
        float lse = m + __logf(sum);
        out[orow + ow0] = post_op_fast(lse, gbias);
    }
    if (v1) {
        float m = acc1[0];
        #pragma unroll
        for (int oc = 1; oc < 16; ++oc) m = fmaxf(m, acc1[oc]);
        float sum = 0.f;
        #pragma unroll
        for (int oc = 0; oc < 16; ++oc) sum += __expf(acc1[oc] - m);
        float lse = m + __logf(sum);
        out[orow + ow1] = post_op_fast(lse, gbias);
    }
}

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

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();
    TORCH_CHECK(cc.numel() == 16, "custom kernel expects 16 conv biases");
    TORCH_CHECK(xc.numel() < (int64_t)INT_MAX, "input too large");

    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;
    TORCH_CHECK(Do > 0 && Ho > 0 && Wo > 0, "bad output size");

    int64_t nRows64 = B * Do * Ho;
    int64_t WP = (Wo + 3) / 4;
    int64_t totalWork64 = nRows64 * 2 * WP;
    TORCH_CHECK(nRows64 < (int64_t)INT_MAX && totalWork64 < (int64_t)INT_MAX,
                "output too large");
    int totalWork = (int)totalWork64;

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

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

    if (tf32) {
        fused_convt3d_lse_post_v2_kernel<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, (int)Do, (int)Ho, (int)Wo,
            (int)nRows64, (int)WP, totalWork);
    } else {
        fused_convt3d_lse_post_v2_kernel<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, (int)Do, (int)Ho, (int)Wo,
            (int)nRows64, (int)WP, totalWork);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_ops = load_inline(
    name="fused_convt3d_lse_post_v6",
    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, using two same-parity output voxels
    per thread so each shared-memory weight load feeds two FMAs. Safe path:
    the exact reference cuDNN ConvTranspose3d plus one fused post kernel. 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)

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