```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA kernel: direct (gather-based) 3D transposed convolution.
#
# For every output element out[n, o, d, h, w] it accumulates over all input
# channels and kernel taps:
#     out[n,o,d,h,w] = sum_{ci,kd,kh,kw}
#         x[n, ci, (d+pd-kd)/sd, (h+ph-kh)/sh, (w+pw-kw)/sw]
#         * w[ci, o, kd, kh, kw]
# (only when the input indices are integral and in range).
#
# The tiled kernel assigns one thread block per (batch, out_channel, d, h/w
# tile), stages the weights of the block's output channel in shared memory,
# and is specialized for unit stride to avoid integer divisions.
# ---------------------------------------------------------------------------

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>

#define TX 16
#define TY 16

template <bool UNIT_STRIDE, bool W_SMEM>
__global__ void conv_transpose3d_tiled_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int Ci, const int Co,
    const int Di, const int Hi, const int Wi,
    const int Do, const int Ho, const int Wo,
    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 int tilesH)
{
    extern __shared__ float s_w[];

    const int o = blockIdx.z % Co;
    const int n = blockIdx.z / Co;
    const int K3 = Kd * Kh * Kw;

    if (W_SMEM) {
        const int total = Ci * K3;
        const int tid = threadIdx.y * blockDim.x + threadIdx.x;
        const int nthr = blockDim.x * blockDim.y;
        const long long wbase = (long long)o * K3;
        for (int idx = tid; idx < total; idx += nthr) {
            const int ci = idx / K3;
            const int t  = idx - ci * K3;
            s_w[idx] = w[(long long)ci * Co * K3 + wbase + t];
        }
        __syncthreads();
    }

    const int d  = blockIdx.y / tilesH;
    const int oh = (blockIdx.y % tilesH) * TY + threadIdx.y;
    const int ow = blockIdx.x * TX + threadIdx.x;

    if (oh >= Ho || ow >= Wo) return;

    float acc = 0.f;
    for (int ci = 0; ci < Ci; ++ci) {
        const long long xci = ((long long)n * Ci + ci) * Di;
        for (int kd = 0; kd < Kd; ++kd) {
            int di;
            if (UNIT_STRIDE) { di = d + pd - kd; }
            else { const int t = d + pd - kd; if (t % sd) continue; di = t / sd; }
            if ((unsigned)di >= (unsigned)Di) continue;
            const long long xd = (xci + di) * Hi;
            for (int kh = 0; kh < Kh; ++kh) {
                int hi;
                if (UNIT_STRIDE) { hi = oh + ph - kh; }
                else { const int t = oh + ph - kh; if (t % sh) continue; hi = t / sh; }
                if ((unsigned)hi >= (unsigned)Hi) continue;
                const long long xh = (xd + hi) * Wi;
                const int wkh = (kd * Kh + kh) * Kw;
                for (int kw = 0; kw < Kw; ++kw) {
                    int wi;
                    if (UNIT_STRIDE) { wi = ow + pw - kw; }
                    else { const int t = ow + pw - kw; if (t % sw) continue; wi = t / sw; }
                    if ((unsigned)wi >= (unsigned)Wi) continue;
                    const float xv = x[xh + wi];
                    const int t3 = wkh + kw;
                    const float wv = W_SMEM ? s_w[ci * K3 + t3]
                                            : w[((long long)ci * Co + o) * K3 + t3];
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }
    if (bias != nullptr) acc += bias[o];
    out[(((long long)n * Co + o) * Do + d) * Ho * Wo + (long long)oh * Wo + ow] = acc;
}

__global__ void conv_transpose3d_naive_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int Ci, const int Co,
    const int Di, const int Hi, const int Wi,
    const int Do, const int Ho, const int Wo,
    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 long total)
{
    const long long stride = (long long)gridDim.x * blockDim.x;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < total; idx += stride) {
        long long t = idx;
        const int ow = t % Wo; t /= Wo;
        const int oh = t % Ho; t /= Ho;
        const int d  = t % Do; t /= Do;
        const int o  = t % Co; t /= Co;
        const int n  = (int)t;
        const int K3 = Kd * Kh * Kw;
        float acc = 0.f;
        for (int ci = 0; ci < Ci; ++ci) {
            for (int kd = 0; kd < Kd; ++kd) {
                const int td = d + pd - kd;
                if (td % sd) continue;
                const int di = td / sd;
                if ((unsigned)di >= (unsigned)Di) continue;
                for (int kh = 0; kh < Kh; ++kh) {
                    const int th = oh + ph - kh;
                    if (th % sh) continue;
                    const int hi = th / sh;
                    if ((unsigned)hi >= (unsigned)Hi) continue;
                    for (int kw = 0; kw < Kw; ++kw) {
                        const int tw = ow + pw - kw;
                        if (tw % sw) continue;
                        const int wi = tw / sw;
                        if ((unsigned)wi >= (unsigned)Wi) continue;
                        acc = fmaf(x[(((long long)n * Ci + ci) * Di + di) * Hi * Wi
                                     + (long long)hi * Wi + wi],
                                   w[((long long)ci * Co + o) * K3
                                     + (kd * Kh + kh) * Kw + kw], acc);
                    }
                }
            }
        }
        if (bias != nullptr) acc += bias[o];
        out[idx] = acc;
    }
}

torch::Tensor conv_transpose3d_cuda(
    torch::Tensor x, torch::Tensor w, torch::Tensor bias,
    std::vector<int64_t> stride, std::vector<int64_t> padding,
    std::vector<int64_t> output_padding)
{
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "x and w must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(x.dim() == 5 && w.dim() == 5, "x and w must be 5-D");

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int N  = xc.size(0), Ci = xc.size(1);
    const int Di = xc.size(2), Hi = xc.size(3), Wi = xc.size(4);
    const int Co = wc.size(1);
    const int Kd = wc.size(2), Kh = wc.size(3), Kw = wc.size(4);

    const int sd = (int)stride[0], sh = (int)stride[1], sw = (int)stride[2];
    const int pd = (int)padding[0], ph = (int)padding[1], pw = (int)padding[2];
    const int od = (int)output_padding[0], oh_ = (int)output_padding[1], ow_ = (int)output_padding[2];

    const int Do = (Di - 1) * sd - 2 * pd + Kd + od;
    const int Ho = (Hi - 1) * sh - 2 * ph + Kh + oh_;
    const int Wo = (Wi - 1) * sw - 2 * pw + Kw + ow_;

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

    at::Tensor bc;
    const float* bias_ptr = nullptr;
    if (bias.defined() && bias.numel() > 0) {
        bc = bias.contiguous();
        bias_ptr = bc.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    const bool unit = (sd == 1 && sh == 1 && sw == 1);
    const int K3 = Kd * Kh * Kw;
    const size_t smem = (size_t)Ci * K3 * sizeof(float);
    const bool wsmem = smem <= (size_t)(46 * 1024);

    const int tilesW = (Wo + TX - 1) / TX;
    const int tilesH = (Ho + TY - 1) / TY;
    const long long gridY = (long long)Do * tilesH;
    const long long gridZ = (long long)N * Co;
    const bool tiled_ok = (gridY <= 65535LL) && (gridZ <= 65535LL);

    if (tiled_ok) {
        dim3 block(TX, TY);
        dim3 grid((unsigned)tilesW, (unsigned)gridY, (unsigned)gridZ);
        if (unit && wsmem) {
            conv_transpose3d_tiled_kernel<true, true><<<grid, block, smem, stream>>>(
                xp, wp, bias_ptr, op, Ci, Co, Di, Hi, Wi, Do, Ho, Wo,
                Kd, Kh, Kw, sd, sh, sw, pd, ph, pw, tilesH);
        } else if (unit) {
            conv_transpose3d_tiled_kernel<true, false><<<grid, block, 0, stream>>>(
                xp, wp, bias_ptr, op, Ci, Co, Di, Hi, Wi, Do, Ho, Wo,
                Kd, Kh, Kw, sd, sh, sw, pd, ph, pw, tilesH);
        } else if (wsmem) {
            conv_transpose3d_tiled_kernel<false, true><<<grid, block, smem, stream>>>(
                xp, wp, bias_ptr, op, Ci, Co, Di, Hi, Wi, Do, Ho, Wo,
                Kd, Kh, Kw, sd, sh, sw, pd, ph, pw, tilesH);
        } else {
            conv_transpose3d_tiled_kernel<false, false><<<grid, block, 0, stream>>>(
                xp, wp, bias_ptr, op, Ci, Co, Di, Hi, Wi, Do, Ho, Wo,
                Kd, Kh, Kw, sd, sh, sw, pd, ph, pw, tilesH);
        }
    } else {
        const long long total = (long long)N * Co * Do * Ho * Wo;
        const int threads = 256;
        long long blocks = (total + threads - 1) / threads;
        if (blocks > 2147483647LL) blocks = 2147483647LL;
        conv_transpose3d_naive_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            xp, wp, bias_ptr, op, Ci, Co, Di, Hi, Wi, Do, Ho, Wo,
            Kd, Kh, Kw, sd, sh, sw, pd, ph, pw, total);
    }
    return out;
}
"""

_cpp_source = (
    "torch::Tensor conv_transpose3d_cuda(torch::Tensor x, torch::Tensor w, "
    "torch::Tensor bias, std::vector<int64_t> stride, "
    "std::vector<int64_t> padding, std::vector<int64_t> output_padding);"
)

try:
    _conv_transpose3d_ext = load_inline(
        name="conv_transpose3d_tiled_ext",
        cpp_sources=_cpp_source,
        cuda_sources=_cuda_source,
        functions=["conv_transpose3d_cuda"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    _conv_transpose3d_ext = None


class ModelNew(nn.Module):
    """
    Optimized ConvTranspose3d.

    Speedups:
      1. A custom tiled CUDA kernel implementing the transposed 3D
         convolution directly (weights staged in shared memory).
      2. For unit stride / zero output padding, the transposed convolution is
         algebraically rewritten as a regular convolution with full padding
         (flipped, channel-transposed kernel), which usually maps to faster
         cuDNN algorithms.
      3. A one-time runtime autotuner benchmarks all valid implementations
         (custom kernel, cuDNN conv3d, cuDNN conv3d channels-last, cuDNN
         transposed-conv, cuDNN transposed-conv channels-last) with
         cudnn.benchmark enabled, keeps only those matching the reference
         output within tolerance, and permanently selects the fastest.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: tuple = (1, 1, 1), padding: tuple = (0, 0, 0),
                 output_padding: tuple = (0, 0, 0), groups: int = 1,
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding, groups=groups,
            bias=bias)
        self._best = None
        self._cache_key = None
        self._w_conv = None
        self._w_conv_cl = None
        self._w_cl = None
        self._pad_conv = (0, 0, 0)

    # ------------------------------------------------------------------ util

    def _equiv_ok(self) -> bool:
        m = self.conv_transpose3d
        return (m.groups == 1
                and tuple(m.stride) == (1, 1, 1)
                and tuple(m.dilation) == (1, 1, 1)
                and tuple(m.output_padding) == (0, 0, 0))

    def _refresh_weight_cache(self, force: bool = False):
        m = self.conv_transpose3d
        w = m.weight
        key = (w.data_ptr(), w._version, str(w.device), w.dtype)
        if force or self._cache_key != key:
            wt = w.detach()
            # ConvTranspose weight: (Ci, Co, Kd, Kh, Kw). The equivalent
            # regular-conv weight is (Co, Ci, Kd, Kh, Kw) with spatial flip
            # and padding (k - 1 - p).
            self._w_conv = wt.transpose(0, 1).flip([2, 3, 4]).contiguous()
            self._w_conv_cl = self._w_conv.to(
                memory_format=torch.channels_last_3d)
            self._w_cl = wt.to(memory_format=torch.channels_last_3d)
            self._pad_conv = tuple(k - 1 - p for k, p in
                                   zip(m.kernel_size, m.padding))
            self._cache_key = key

    @staticmethod
    def _bench_call(fn, *args):
        prev = torch.backends.cudnn.benchmark
        torch.backends.cudnn.benchmark = True
        try:
            return fn(*args)
        finally:
            torch.backends.cudnn.benchmark = prev

    @staticmethod
    def _timeit(fn, warmup: int = 2, iters: int = 5) -> float:
        for _ in range(warmup):
            fn()
        torch.cuda.synchronize()
        s = torch.cuda.Event(enable_timing=True)
        e = torch.cuda.Event(enable_timing=True)
        s.record()
        for _ in range(iters):
            fn()
        e.record()
        torch.cuda.synchronize()
        return s.elapsed_time(e) / iters

    # -------------------------------------------------------------- autotune

    @torch.no_grad()
    def _autotune(self, x: torch.Tensor):
        m = self.conv_transpose3d
        prev = torch.backends.cudnn.benchmark
        torch.backends.cudnn.benchmark = True
        try:
            ref = m(x)
            cands = [('orig', lambda: m(x))]

            def orig_cl():
                return F.conv_transpose3d(
                    x.contiguous(memory_format=torch.channels_last_3d),
                    m.weight.to(memory_format=torch.channels_last_3d),
                    m.bias, m.stride, m.padding, m.output_padding,
                    m.groups, m.dilation)
            cands.append(('orig_cl', orig_cl))

            if self._equiv_ok():
                self._refresh_weight_cache(force=True)
                pad = self._pad_conv
                w_conv = self._w_conv
                w_conv_cl = self._w_conv_cl
                cands.append(('conv3d', lambda: F.conv3d(
                    x, w_conv, m.bias, (1, 1, 1), pad, (1, 1, 1), m.groups)))

                def conv3d_cl():
                    return F.conv3d(
                        x.contiguous(memory_format=torch.channels_last_3d),
                        w_conv_cl, m.bias, (1, 1, 1), pad, (1, 1, 1), m.groups)
                cands.append(('conv3d_cl', conv3d_cl))

            if (_conv_transpose3d_ext is not None and m.groups == 1
                    and tuple(m.dilation) == (1, 1, 1)
                    and x.dtype == torch.float32):
                bias0 = (m.bias if m.bias is not None else
                         torch.empty(0, device=x.device, dtype=x.dtype))
                w = m.weight

                def custom():
                    return _conv_transpose3d_ext.conv_transpose3d_cuda(
                        x, w, bias0, list(m.stride), list(m.padding),
                        list(m.output_padding))
                cands.append(('custom', custom))

            best_name, best_t = 'orig', float('inf')
            for name, fn in cands:
                try:
                    out = fn()
                    torch.cuda.synchronize()
                    ok = (out.shape == ref.shape and
                          torch.allclose(out, ref, atol=2e-3, rtol=1e-2))
                    del out
                    if not ok:
                        continue
                    t = self._timeit(fn)
                except Exception:
                    continue
                if t < best_t:
                    best_name, best_t = name, t
            self._best = best_name
        finally:
            torch.backends.cudnn.benchmark = prev

    # -------------------------------------------------------------- dispatch

    def _run_best(self, x: torch.Tensor) -> torch.Tensor:
        m = self.conv_transpose3d
        name = self._best
        if name == 'conv3d':
            self._refresh_weight_cache()
            return self._bench_call(
                lambda: F.conv3d(x, self._w_conv, m.bias, (1, 1, 1),
                                 self._pad_conv, (1, 1, 1), m.groups))
        if name == 'conv3d_cl':
            self._refresh_weight_cache()
            return self._bench_call(
                lambda: F.conv3d(
                    x.contiguous(memory_format=torch.channels_last_3d),
                    self._w_conv_cl, m.bias, (1, 1, 1), self._pad_conv,
                    (1, 1, 1), m.groups))
        if name == 'orig_cl':
            self._refresh_weight_cache()
            return self._bench_call(
                lambda: F.conv_transpose3d(
                    x.contiguous(memory_format=torch.channels_last_3d),
                    self._w_cl, m.bias, m.stride, m.padding,
                    m.output_padding, m.groups, m.dilation))
        if name == 'custom' and _conv_transpose3d_ext is not None:
            bias0 = (m.bias if m.bias is not None else
                     torch.empty(0, device=x.device, dtype=x.dtype))
            return _conv_transpose3d_ext.conv_transpose3d_cuda(
                x, m.weight, bias0, list(m.stride), list(m.padding),
                list(m.output_padding))
        return self._bench_call(m, x)

    # --------------------------------------------------------------- forward

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        m = self.conv_transpose3d
        if not x.is_cuda:
            return m(x)

        if torch.is_grad_enabled():
            # Autograd-safe path: differentiable rewrite as regular conv.
            if self._equiv_ok():
                w = m.weight
                w_conv = w.transpose(0, 1).flip([2, 3, 4]).contiguous()
                pad = tuple(k - 1 - p for k, p in
                            zip(m.kernel_size, m.padding))
                return self._bench_call(
                    lambda: F.conv3d(x, w_conv, m.bias, (1, 1, 1), pad,
                                     (1, 1, 1), m.groups))
            return self._bench_call(m, x)

        if self._best is None:
            try:
                self._autotune(x)
            except Exception:
                self._best = 'orig'
        try:
            return self._run_best(x)
        except Exception:
            return m(x)
```