```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:
# Transforms a ConvTranspose2d weight (IC, OC, KH, KW) into the mathematically
# equivalent Conv2d weight (OC, IC, KH, KW) by swapping the in/out channel
# axes and flipping both spatial axes.
#
# For stride == 1, dilation == 1, output_padding == 0, groups == 1:
#     ConvTranspose2d(x, w, padding=P)
#     == Conv2d(x, flip_transpose(w), padding=K-1-P)
# This lets us run the (much better optimized) forward-convolution cuDNN
# engines (e.g. Winograd / implicit-GEMM) instead of the backward-data
# convolution engines used by ConvTranspose2d, yielding a solid speedup.
# ---------------------------------------------------------------------------
_weight_transform_cuda_src = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

__global__ void transpose_conv_weight_to_conv_kernel(
        const float* __restrict__ w,
        float* __restrict__ out,
        int IC, int OC, int KH, int KW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = IC * OC * KH * KW;
    if (idx >= total) return;
    // Decompose idx over the output layout (o, c, p, q)
    int q = idx % KW;
    int t = idx / KW;
    int p = t % KH;
    t /= KH;
    int c = t % IC;
    int o = t / IC;
    // Source element in the ConvTranspose2d weight (IC, OC, KH, KW)
    out[idx] = w[((c * OC + o) * KH + (KH - 1 - p)) * KW + (KW - 1 - q)];
}

torch::Tensor transpose_conv_weight_to_conv(torch::Tensor w) {
    TORCH_CHECK(w.is_cuda() && w.dim() == 4 && w.scalar_type() == at::kFloat,
                "expected a CUDA float32 4D weight tensor");
    auto wc = w.contiguous();
    const int IC = wc.size(0), OC = wc.size(1), KH = wc.size(2), KW = wc.size(3);
    auto out = at::empty({OC, IC, KH, KW}, wc.options());
    const int total = IC * OC * KH * KW;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    transpose_conv_weight_to_conv_kernel<<<blocks, threads>>>(
        wc.data_ptr<float>(), out.data_ptr<float>(), IC, OC, KH, KW);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_weight_transform_cpp_src = (
    "torch::Tensor transpose_conv_weight_to_conv(torch::Tensor w);"
)

_weight_transform_ext = load_inline(
    name="transpose_conv_weight_to_conv_ext",
    cpp_sources=_weight_transform_cpp_src,
    cuda_sources=_weight_transform_cuda_src,
    functions=["transpose_conv_weight_to_conv"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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

    A stride-1 ConvTranspose2d is algebraically identical to a regular Conv2d
    whose kernel is channel-swapped and spatially flipped, applied with
    padding = kernel_size - 1 - padding. We perform that kernel conversion
    once with a custom CUDA kernel (cached afterwards) and then execute the
    forward pass with cuDNN's forward-convolution path, which is significantly
    faster than the backward-data path used by ConvTranspose2d.

    All other configurations fall back to the reference implementation.
    """
    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)
        self._conv2d_weight = None
        self._conv2d_weight_key = None

    def _transform_weight(self, w: torch.Tensor) -> torch.Tensor:
        if w.is_cuda and w.dtype == torch.float32:
            return _weight_transform_ext.transpose_conv_weight_to_conv(w)
        # Generic fallback (CPU / other dtypes).
        return w.transpose(0, 1).flip([2, 3]).contiguous()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        ct = self.conv_transpose2d
        kh, kw = ct.kernel_size
        ph, pw = ct.padding

        fast_path = (
            ct.stride == (1, 1)
            and ct.dilation == (1, 1)
            and ct.groups == 1
            and ct.output_padding == (0, 0)
            and ph <= kh - 1
            and pw <= kw - 1
        )
        if not fast_path:
            return ct(x)

        w = ct.weight
        if torch.is_grad_enabled() and w.requires_grad:
            # Differentiable transform (permute + flip) keeps autograd exact.
            w2 = w.transpose(0, 1).flip([2, 3]).contiguous()
        else:
            # Cache the transformed weight; refresh if the parameter changes.
            key = (w.data_ptr(), w._version, w.device, w.dtype)
            if self._conv2d_weight is None or self._conv2d_weight_key != key:
                self._conv2d_weight = self._transform_weight(w.detach())
                self._conv2d_weight_key = key
            w2 = self._conv2d_weight

        return F.conv2d(x, w2, ct.bias, stride=1,
                        padding=(kh - 1 - ph, kw - 1 - pw))
```