```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 converts a ConvTranspose2d weight of shape
# (in_channels, out_channels, kH, kW) into the mathematically equivalent
# Conv2d weight of shape (out_channels, in_channels, kH, kW) by swapping the
# channel axes and flipping both spatial axes.
#
# For stride == 1, a transposed convolution is identical to a regular
# convolution that uses this transformed kernel with
#     padding = dilation * (k - 1) - padding.
# The forward convolution path of cuDNN is generally much better optimized
# than the backward-data path used to implement ConvTranspose2d, so this
# transformation plus channels_last + autotuning yields a solid speedup.
# ----------------------------------------------------------------------------
weight_flip_cuda_source = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

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

    // out has shape (B, A, KH, KW), linear layout
    const int x = (int)(idx % KW);
    const int y = (int)((idx / KW) % KH);
    const int a = (int)((idx / ((long long)KW * KH)) % A);
    const int b = (int)(idx / ((long long)KW * KH * A));

    // w has shape (A, B, KH, KW); flip spatial dims
    const long long src =
        (((long long)a * B + b) * KH + (KH - 1 - y)) * KW + (KW - 1 - x);
    out[idx] = w[src];
}

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

    auto wc = w.contiguous();
    const int A  = (int)wc.size(0);
    const int B  = (int)wc.size(1);
    const int KH = (int)wc.size(2);
    const int KW = (int)wc.size(3);

    auto out = torch::empty({B, A, KH, KW}, wc.options());
    const long long total = (long long)A * B * KH * KW;
    if (total == 0) {
        return out;
    }

    const int threads = 256;
    const long long blocks = (total + threads - 1) / threads;
    flip_transpose_weight_kernel<<<(unsigned int)blocks, threads>>>(
        wc.data_ptr<float>(), out.data_ptr<float>(), A, B, KH, KW);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

weight_flip_cpp_source = (
    "torch::Tensor flip_transpose_weight_cuda(torch::Tensor w);"
)

weight_flip = load_inline(
    name="conv_transpose2d_weight_flip_v1",
    cpp_sources=weight_flip_cpp_source,
    cuda_sources=weight_flip_cuda_source,
    functions=["flip_transpose_weight_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    For the common stride-1 configuration, the transposed convolution is
    evaluated as an equivalent regular convolution (weight transformed by the
    custom CUDA kernel above) executed in channels_last memory format with
    cuDNN autotuning enabled. All other configurations fall back to the
    reference nn.ConvTranspose2d implementation.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: tuple = (1, 1), padding: tuple = (0, 0),
                 output_padding: tuple = (0, 0), dilation: tuple = (1, 1),
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding,
            dilation=dilation, groups=groups, bias=bias)
        self._weight_flip = weight_flip

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        m = self.conv_transpose2d
        kh, kw = m.kernel_size
        dh, dw = m.dilation
        ph, pw = m.padding
        # Equivalent Conv2d padding: dilation * (k - 1) - padding
        pad_h = dh * (kh - 1) - ph
        pad_w = dw * (kw - 1) - pw

        if (m.stride == (1, 1) and m.output_padding == (0, 0)
                and m.groups == 1 and pad_h >= 0 and pad_w >= 0
                and x.is_cuda and x.dim() == 4
                and x.dtype == torch.float32):
            # Transform the transposed-conv weight into a conv weight using
            # the custom CUDA kernel (cheap: weight is tiny compared to the
            # activations), then run the heavily optimized cuDNN forward conv.
            w = self._weight_flip.flip_transpose_weight_cuda(m.weight)
            x_cl = x.contiguous(memory_format=torch.channels_last)
            w_cl = w.contiguous(memory_format=torch.channels_last)

            prev_benchmark = torch.backends.cudnn.benchmark
            torch.backends.cudnn.benchmark = True
            try:
                out = F.conv2d(x_cl, w_cl, m.bias, (1, 1),
                               (pad_h, pad_w), (dh, dw), 1)
            finally:
                torch.backends.cudnn.benchmark = prev_benchmark
            return out

        # Fallback for configurations the fast path does not cover.
        return m(x)
```