```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 conv engine (incl. TF32 tensor cores).
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.backends.cudnn.benchmark_limit = 0  # exhaustive engine search
except Exception:
    pass

# ---------------------------------------------------------------------------
# stride-1 / pad-0 / dil-1 ConvTranspose2d is mathematically identical to a
# regular Conv2d with "full" padding (kernel - 1) and a weight tensor whose
# channel axes are swapped and spatial axes flipped:
#     W2[o,i,kh,kw] = Wt[i,o,KH-1-kh,KW-1-kw]
# The forward-conv path uses cuDNN's highly optimized implicit-GEMM
# tensor-core kernels, much faster than the backward-data kernels used by
# ConvTranspose2d. The custom CUDA kernel below performs the weight
# permutation (tiny, cached) and we additionally autotune the memory layout.
# ---------------------------------------------------------------------------

_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]
__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_wt_xform_v3",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["flip_transpose_weight"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class _WeightXformFn(torch.autograd.Function):
    """Differentiable wrapper around the custom weight-permutation kernel."""

    @staticmethod
    def forward(ctx, w):
        return weight_xform.flip_transpose_weight(w)

    @staticmethod
    def backward(ctx, grad_out):
        # The permutation is an involution: the same kernel maps grad back.
        return weight_xform.flip_transpose_weight(grad_out.contiguous())


class ModelNew(nn.Module):
    """
    Optimized ConvTranspose2d for stride=1, padding=0, output_padding=0,
    dilation=1, groups=1. Replaces it with the mathematically equivalent
    full-padding forward convolution evaluated by cuDNN tensor-core kernels,
    using a custom CUDA kernel for the weight flip/transpose (cached), plus a
    one-time autotune over NCHW / channels-last execution layouts.
    """

    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__()
        # Same submodule (parameter names/init) as the reference model.
        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)
        self._full_pad = (k[0] - 1, k[1] - 1)
        self._wcache = None   # [data_ptr, version, {layout: w2}]
        self._layout = {}     # input shape -> 0 (NCHW), 1 (NHWC->NCHW), 2 (NHWC)

    # ------------------------------------------------------------------
    def _get_w2(self, w, channels_last):
        if torch.is_grad_enabled() and w.requires_grad:
            w2 = _WeightXformFn.apply(w)
            if channels_last:
                w2 = w2.contiguous(memory_format=torch.channels_last)
            return w2
        ver = getattr(w, "_version", 0)
        c = self._wcache
        if c is None or c[0] != w.data_ptr() or c[1] != ver:
            with torch.inference_mode(False), torch.no_grad():
                w2 = weight_xform.flip_transpose_weight(w.detach())
            c = [w.data_ptr(), ver, {0: w2}]
            self._wcache = c
        d = c[2]
        if channels_last:
            w2 = d.get(1)
            if w2 is None:
                with torch.inference_mode(False), torch.no_grad():
                    w2 = d[0].contiguous(memory_format=torch.channels_last)
                d[1] = w2
            return w2
        return d[0]

    # ------------------------------------------------------------------
    @staticmethod
    def _time_op(fn, warmup=3, iters=6):
        for _ in range(warmup):
            fn()
        torch.cuda.synchronize()
        s = torch.cuda.Event(enable_timing=True)
        e = torch.cuda.Event(enable_timing=True)
        s.record()
        for _ in range(iters):
            fn()
        e.record()
        torch.cuda.synchronize()
        return s.elapsed_time(e)

    @torch.no_grad()
    def _autotune(self, x, w):
        bias = self.conv_transpose2d.bias
        pad = self._full_pad
        with torch.inference_mode(False):
            w2 = weight_xform.flip_transpose_weight(w.detach())
        self._wcache = [w.data_ptr(), getattr(w, "_version", 0), {0: w2}]

        t0 = self._time_op(lambda: F.conv2d(x, w2, bias, (1, 1), pad))
        t1 = t2 = float("inf")
        try:
            x_c = x.contiguous(memory_format=torch.channels_last)
            with torch.inference_mode(False):
                w2c = w2.contiguous(memory_format=torch.channels_last)
            self._wcache[2][1] = w2c
            t1 = self._time_op(lambda: F.conv2d(
                x_c, w2c, bias, (1, 1), pad).contiguous(
                    memory_format=torch.contiguous_format))
            t2 = self._time_op(lambda: F.conv2d(x_c, w2c, bias, (1, 1), pad))
        except Exception:
            pass

        best, best_t = 0, t0
        if t1 < best_t * 0.98:
            best, best_t = 1, t1
        if t2 < best_t * 0.98:
            best = 2
        return best

    # ------------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv_transpose2d.weight
        if (self._use_fast and x.dim() == 4 and x.is_cuda and w.is_cuda
                and x.dtype == torch.float32 and w.dtype == torch.float32):
            layout = self._layout.get(x.shape)
            if layout is None:
                layout = self._autotune(x, w)
                self._layout[x.shape] = layout
            bias = self.conv_transpose2d.bias
            if layout == 0:
                return F.conv2d(x, self._get_w2(w, False), bias,
                                (1, 1), self._full_pad)
            x_c = x.contiguous(memory_format=torch.channels_last)
            out = F.conv2d(x_c, self._get_w2(w, True), bias,
                           (1, 1), self._full_pad)
            if layout == 1:
                return out.contiguous(memory_format=torch.contiguous_format)
            return out
        return self.conv_transpose2d(x)
```