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

# Enable cuDNN autotuning so the fastest convolution algorithm is selected.
torch.backends.cudnn.benchmark = True
if hasattr(torch.backends.cudnn, "benchmark_limit"):
    torch.backends.cudnn.benchmark_limit = 0
torch.backends.cudnn.allow_tf32 = True

# ---------------------------------------------------------------------------
# Key algorithmic idea:
# A ConvTranspose2d with stride == 1 and output_padding == 0 is mathematically
# identical to a regular Conv2d with:
#   - padding' = kernel_size - 1 - padding
#   - weight'  = conv_transpose weight with channels transposed (in <-> out)
#                and the spatial kernel flipped along both axes.
# This lets us run the (much better optimized) forward-convolution cuDNN
# kernels (including Winograd) instead of the backward-data path used by
# ConvTranspose2d.
#
# The weight transform (flip + channel transpose + contiguous copy) is fused
# into a single custom CUDA kernel below.
# ---------------------------------------------------------------------------

_weight_transform_cpp_source = "torch::Tensor transform_weight(torch::Tensor w);"

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

// Fused kernel: w_out[o, i, ky, kx] = w_in[i, o, kh-1-ky, kw-1-kx]
// Converts a ConvTranspose2d weight (in_c, out_c, kh, kw) into the
// equivalent Conv2d weight (out_c, in_c, kh, kw).
__global__ void transform_weight_kernel(const float* __restrict__ w_in,
                                        float* __restrict__ w_out,
                                        int in_c, int out_c, int kh, int kw) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = in_c * out_c * kh * kw;
    if (idx >= total) return;

    int kx = idx % kw;
    int ky = (idx / kw) % kh;
    int i  = (idx / (kw * kh)) % in_c;
    int o  = idx / (kw * kh * in_c);

    int src = ((i * out_c + o) * kh + (kh - 1 - ky)) * kw + (kw - 1 - kx);
    w_out[idx] = w_in[src];
}

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

    auto wc = w.contiguous();
    const int in_c  = wc.size(0);
    const int out_c = wc.size(1);
    const int kh    = wc.size(2);
    const int kw    = wc.size(3);

    auto w_out = at::empty({out_c, in_c, kh, kw}, wc.options());
    const int total = in_c * out_c * kh * kw;
    if (total == 0) return w_out;

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

    transform_weight_kernel<<<grid, block, 0, stream>>>(
        wc.data_ptr<float>(), w_out.data_ptr<float>(), in_c, out_c, kh, kw);

    return w_out;
}
"""

_weight_transform = load_inline(
    name="conv_transpose_weight_transform",
    cpp_sources=_weight_transform_cpp_source,
    cuda_sources=_weight_transform_cuda_source,
    functions=["transform_weight"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    For the common case of stride == 1 and output_padding == 0 (the
    configuration used by this model), the transposed convolution is exactly
    equivalent to a regular convolution with a flipped, channel-transposed
    kernel and padding' = kernel_size - 1 - padding. The weight conversion is
    performed by a custom fused CUDA kernel and cached; the heavy lifting is
    then done by cuDNN's forward-convolution path (autotuned, Winograd where
    applicable), which is significantly faster than the transposed-conv
    (backward-data) path.

    Any configuration outside this fast path falls back to the original
    nn.ConvTranspose2d implementation, so results remain correct in general.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, output_padding: int = 0,
                 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,
            groups=groups, bias=bias)

        st = self.conv_transpose2d.stride
        pd = self.conv_transpose2d.padding
        op = self.conv_transpose2d.output_padding
        ks = self.conv_transpose2d.kernel_size

        self._fast_path = (
            st[0] == 1 and st[1] == 1 and
            op[0] == 0 and op[1] == 0 and
            self.conv_transpose2d.groups == 1
        )
        self._conv_padding = (ks[0] - 1 - pd[0], ks[1] - 1 - pd[1])

        # Cache for the transformed weight (invalidated if the weight tensor
        # object changes or is modified in place, e.g. via load_state_dict).
        self._w_ref = None
        self._w_version = -1
        self._w_conv = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast_path and x.is_cuda and x.dtype == torch.float32:
            w = self.conv_transpose2d.weight
            if self._w_ref is not w or self._w_version != w._version:
                self._w_conv = _weight_transform.transform_weight(w)
                self._w_ref = w
                self._w_version = w._version
            return F.conv2d(x, self._w_conv, self.conv_transpose2d.bias,
                            stride=(1, 1), padding=self._conv_padding)
        return self.conv_transpose2d(x)
```