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

# Let cuDNN autotune the fastest convolution algorithm (incl. TF32 tensor cores).
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True

# ---------------------------------------------------------------------------
# Algorithmic change:
# A stride-1, padding-0, dilation-1 ConvTranspose2d is mathematically identical
# to a regular Conv2d with "full" padding (pad = kernel - 1) and a weight
# tensor whose channel axes are swapped and whose spatial axes are flipped:
#
#     y[n,o,ho,wo] = sum_{i,kh,kw} x[n,i,ho-kh+KH-1, wo-kw+KW-1] * Wt[i,o,kh,kw]
#                  = conv2d(x, W2, padding=(KH-1, KW-1))
#     W2[o,i,kh,kw] = Wt[i,o,KH-1-kh,KW-1-kw]
#
# The forward-convolution path lets cuDNN use its highly optimized implicit
# GEMM / tensor-core kernels, which are considerably faster than the
# backward-data kernels used by ConvTranspose2d. The weight permutation itself
# is done by the custom CUDA kernel below (O(0.5 MB), negligible cost).
# ---------------------------------------------------------------------------

_cpp_src = "torch::Tensor flip_transpose_weight(torch::Tensor w);"

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

// out[o, i, kh, kw] = w[i, o, KH-1-kh, KW-1-kw]
// w   : [C_in,  C_out, KH, KW]  (ConvTranspose2d weight layout)
// out : [C_out, C_in,  KH, KW]  (Conv2d weight layout)
__global__ void flip_transpose_kernel(
    const float* __restrict__ w,
    float* __restrict__ out,
    int C_in, int C_out, int KH, int KW, int total)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int kw = idx % KW;
    int t  = idx / KW;
    int kh = t % KH;
    t /= KH;
    int i  = t % C_in;
    int o  = t / C_in;

    out[idx] = w[(((size_t)i * C_out + o) * KH + (KH - 1 - kh)) * KW + (KW - 1 - kw)];
}

torch::Tensor flip_transpose_weight(torch::Tensor w)
{
    TORCH_CHECK(w.is_cuda() && w.dim() == 4, "weight must be a 4D CUDA tensor");
    auto w_c = w.contiguous();

    const int C_in  = (int)w_c.size(0);
    const int C_out = (int)w_c.size(1);
    const int KH    = (int)w_c.size(2);
    const int KW    = (int)w_c.size(3);

    auto out = torch::empty({C_out, C_in, KH, KW}, w_c.options());
    const int total = C_in * C_out * KH * KW;

    auto stream = at::cuda::getCurrentCUDAStream();
    flip_transpose_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        w_c.data_ptr<float>(), out.data_ptr<float>(), C_in, C_out, KH, KW, total);

    TORCH_CHECK(cudaGetLastError() == cudaSuccess, "flip_transpose_kernel failed");
    return out;
}
"""

weight_xform = load_inline(
    name="conv_transpose2d_weight_xform",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["flip_transpose_weight"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized ConvTranspose2d for the stride=1, padding=0, output_padding=0,
    dilation=1, groups=1 configuration.

    Replaces the transposed convolution with the mathematically equivalent
    full-padding forward convolution (evaluated by cuDNN's tensor-core
    kernels), using a custom CUDA kernel to permute/flip the weights.
    """

    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__()
        # Keep the exact same submodule (parameter names and initialization)
        # as the reference model so state dicts and RNG behavior match.
        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)

        def _pair(v):
            return (v, v) if isinstance(v, int) else tuple(v)

        stride = _pair(stride)
        padding = _pair(padding)
        output_padding = _pair(output_padding)
        dilation = _pair(dilation)
        k = _pair(kernel_size)

        self._use_fast = (
            stride == (1, 1)
            and padding == (0, 0)
            and output_padding == (0, 0)
            and dilation == (1, 1)
            and groups == 1
        )
        # "Full" padding for the equivalent forward convolution.
        self._full_pad = (k[0] - 1, k[1] - 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv_transpose2d.weight
        if (self._use_fast and x.is_cuda and x.dtype == torch.float32
                and w.is_cuda and w.dtype == torch.float32):
            # Custom CUDA kernel: transpose channel axes + flip spatial axes.
            w2 = weight_xform.flip_transpose_weight(w)
            # Equivalent forward convolution -> fast cuDNN tensor-core path.
            return F.conv2d(x, w2, self.conv_transpose2d.bias,
                            (1, 1), self._full_pad)
        # Fallback to the reference implementation for other configurations.
        return self.conv_transpose2d(x)
```