```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 that rewrites a ConvTranspose3d weight of shape
#   (in_channels, out_channels, kD, kH, kW)
# into the mathematically equivalent Conv3d weight of shape
#   (out_channels, in_channels, kD, kH, kW)
# with every spatial kernel dimension flipped.
#
# For stride == 1 and dilation == 1, a transposed convolution is exactly a
# regular convolution with padding (k - 1 - p) and this flipped /
# channel-transposed kernel.  This lets the heavy computation run through
# cuDNN's highly-optimized forward-convolution kernels instead of the slower
# transposed-convolution (backward-data) path.
# ---------------------------------------------------------------------------
_weight_transform_cpp_src = "torch::Tensor flip_permute_cuda(torch::Tensor w);"

_weight_transform_cuda_src = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void flip_permute_kernel(const float* __restrict__ w,
                                    float* __restrict__ out,
                                    const int CI, const int CO,
                                    const int KD, const int KH, const int KW) {
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long total = (long)CI * CO * KD * KH * KW;
    if (idx >= total) return;

    // Decompose the flat output index into (co, ci, kd, kh, kw).
    long t = idx;
    const int kw = (int)(t % KW); t /= KW;
    const int kh = (int)(t % KH); t /= KH;
    const int kd = (int)(t % KD); t /= KD;
    const int ci = (int)(t % CI); t /= CI;
    const int co = (int)t;

    // out[co, ci, kd, kh, kw] = w[ci, co, KD-1-kd, KH-1-kh, KW-1-kw]
    const long src = ((((long)ci * CO + co) * KD + (KD - 1 - kd)) * KH
                      + (KH - 1 - kh)) * KW + (KW - 1 - kw);
    out[idx] = w[src];
}

torch::Tensor flip_permute_cuda(torch::Tensor w) {
    TORCH_CHECK(w.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(w.dim() == 5, "weight must be 5-dimensional");
    TORCH_CHECK(w.scalar_type() == at::kFloat, "weight must be float32");

    auto wc = w.contiguous();
    const int CI = (int)wc.size(0);
    const int CO = (int)wc.size(1);
    const int KD = (int)wc.size(2);
    const int KH = (int)wc.size(3);
    const int KW = (int)wc.size(4);

    auto out = torch::empty({CO, CI, KD, KH, KW}, wc.options());

    const long total = (long)CI * CO * KD * KH * KW;
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    flip_permute_kernel<<<blocks, threads>>>(
        wc.data_ptr<float>(), out.data_ptr<float>(), CI, CO, KD, KH, KW);
    return out;
}
"""

_weight_transform = load_inline(
    name="conv_transpose3d_weight_transform",
    cpp_sources=_weight_transform_cpp_src,
    cuda_sources=_weight_transform_cuda_src,
    functions=["flip_permute_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized transposed 3D convolution.

    When stride == 1, dilation == 1, output_padding == 0 and groups == 1 the
    transposed convolution is replaced by an equivalent regular Conv3d with
    padding P = kernel_size - 1 - padding and a weight tensor whose spatial
    axes are flipped and whose input/output channel axes are swapped.  The
    weight rewrite is performed by a custom CUDA kernel, and the convolution
    itself runs on cuDNN's forward path with autotuning enabled, which is
    significantly faster than the transposed-convolution path.  All other
    configurations fall back to the original nn.ConvTranspose3d.
    """

    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)

        s = self.conv_transpose3d.stride
        p = self.conv_transpose3d.padding
        op = self.conv_transpose3d.output_padding
        d = self.conv_transpose3d.dilation
        k = self.conv_transpose3d.kernel_size

        self._fast_path = (
            self.conv_transpose3d.groups == 1
            and all(v == 1 for v in s)
            and all(v == 1 for v in d)
            and all(v == 0 for v in op)
            and all(k[i] - 1 - p[i] >= 0 for i in range(3))
        )
        self._conv_padding = tuple(k[i] - 1 - p[i] for i in range(3))

    def _equivalent_conv_weight(self, w: torch.Tensor) -> torch.Tensor:
        if w.is_cuda and w.dtype == torch.float32:
            return _weight_transform.flip_permute_cuda(w)
        # Portable fallback for CPU tensors or non-float32 dtypes.
        return w.flip(dims=[2, 3, 4]).permute(1, 0, 2, 3, 4).contiguous()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast_path:
            w = self._equivalent_conv_weight(self.conv_transpose3d.weight)
            prev_benchmark = torch.backends.cudnn.benchmark
            torch.backends.cudnn.benchmark = True
            try:
                out = F.conv3d(
                    x, w, self.conv_transpose3d.bias,
                    stride=1, padding=self._conv_padding, groups=1)
            finally:
                torch.backends.cudnn.benchmark = prev_benchmark
            return out
        return self.conv_transpose3d(x)
```