```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.nn.modules.utils import _pair
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA kernels for transposed 2D convolution.
#
# A ConvTranspose2d with stride=1, dilation=1, output_padding=0 and groups=1 is
# mathematically identical to a regular Conv2d with the kernel flipped
# spatially, in/out channels swapped and padding P = (K - 1) - p.  We use this
# "full convolution" form to implement a fast, tiled, shared-memory CUDA kernel
# with register blocking (each thread computes CO output channels x 4 output
# pixels).  The tap-remapped weights are precomputed once so the kernel reads
# them with perfectly coalesced, broadcast loads.
# ---------------------------------------------------------------------------

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

// --------------------------- generic kernel -------------------------------
// One thread per output element; works for any kernel size. Correct fallback.
__global__ void tconv2d_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wt,   // (KH*KW, C_in, C_out)
    const float* __restrict__ bias,
    float* __restrict__ out,
    int C_in, int H_in, int W_in,
    int C_out, int H_out, int W_out,
    int KH, int KW, int pad_h, int pad_w)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long total = (long)gridDim.y * C_out * H_out * W_out;
    if (idx >= total) return;
    int n = blockIdx.y;

    int wo = idx % W_out;
    long t1 = idx / W_out;
    int ho = t1 % H_out;
    int co = t1 / H_out;

    float acc = 0.f;
    for (int ci = 0; ci < C_in; ++ci) {
        const float* xb = x + ((long)n * C_in + ci) * H_in * W_in;
        const float* wb = wt + (long)ci * C_out + co;
        for (int kh = 0; kh < KH; ++kh) {
            int gx = ho + kh - pad_h;
            if (gx < 0 || gx >= H_in) continue;
            for (int kw = 0; kw < KW; ++kw) {
                int gy = wo + kw - pad_w;
                if (gy < 0 || gy >= W_in) continue;
                acc += xb[gx * W_in + gy] * wb[(long)(kh * KW + kw) * C_in * C_out];
            }
        }
    }
    if (bias) acc += bias[co];
    out[((long)(n * C_out + co) * H_out + ho) * W_out + wo] = acc;
}

// ------------------------- fast tiled kernel ------------------------------
// Block: 16x16 threads, each thread computes CO channels x 4 pixels.
// Input tile staged in shared memory, weights broadcast from L1.
template <int KH, int KW, int CO>
__global__ void tconv2d_tiled_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wt,   // (KH*KW, C_in, C_out)
    const float* __restrict__ bias,
    float* __restrict__ out,
    int C_in, int H_in, int W_in,
    int C_out, int H_out, int W_out,
    int pad_h, int pad_w)
{
    constexpr int TILE_H = 16;
    constexpr int TILE_W = 64;
    constexpr int TX_H = TILE_H + KH - 1;
    constexpr int TX_W = TILE_W + KW - 1;

    __shared__ float tile[TX_H][TX_W];

    const int tx = threadIdx.x;  // 0..15
    const int ty = threadIdx.y;  // 0..15
    const int tid = ty * 16 + tx;

    const int w0 = blockIdx.x * TILE_W;
    const int h0 = blockIdx.y * TILE_H;
    const int nz = blockIdx.z;
    const int co0 = (nz % (C_out / CO)) * CO;
    const int n  = nz / (C_out / CO);

    const int wo = w0 + tx * 4;
    const int ho = h0 + ty;

    float acc[CO][4];
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) acc[c][j] = 0.f;
    }

    constexpr int TILE_ELEMS = TX_H * TX_W;
    constexpr int LOADS = (TILE_ELEMS + 255) / 256;

    for (int ci = 0; ci < C_in; ++ci) {
        const float* xb = x + ((long)n * C_in + ci) * H_in * W_in;
        #pragma unroll
        for (int i = 0; i < LOADS; ++i) {
            int idx = tid + i * 256;
            if (idx < TILE_ELEMS) {
                int r = idx / TX_W;
                int c = idx % TX_W;
                int gx = h0 + r - pad_h;
                int gy = w0 + c - pad_w;
                float v = 0.f;
                if (gx >= 0 && gx < H_in && gy >= 0 && gy < W_in)
                    v = xb[gx * W_in + gy];
                tile[r][c] = v;
            }
        }
        __syncthreads();

        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            float xs[KW + 3];
            #pragma unroll
            for (int t = 0; t < KW + 3; ++t) xs[t] = tile[ty + kh][tx * 4 + t];
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                const float* wp = wt + ((long)(kh * KW + kw) * C_in + ci) * C_out + co0;
                #pragma unroll
                for (int c = 0; c < CO; ++c) {
                    float wv = wp[c];
                    acc[c][0] = fmaf(xs[kw + 0], wv, acc[c][0]);
                    acc[c][1] = fmaf(xs[kw + 1], wv, acc[c][1]);
                    acc[c][2] = fmaf(xs[kw + 2], wv, acc[c][2]);
                    acc[c][3] = fmaf(xs[kw + 3], wv, acc[c][3]);
                }
            }
        }
        __syncthreads();
    }

    if (ho < H_out) {
        #pragma unroll
        for (int c = 0; c < CO; ++c) {
            int co = co0 + c;
            float b = bias ? bias[co] : 0.f;
            long off = ((long)(n * C_out + co) * H_out + ho) * W_out + wo;
            if (wo + 3 < W_out) {
                float4 v = make_float4(acc[c][0] + b, acc[c][1] + b,
                                       acc[c][2] + b, acc[c][3] + b);
                *reinterpret_cast<float4*>(out + off) = v;
            } else if (wo < W_out) {
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    if (wo + j < W_out) out[off + j] = acc[c][j] + b;
            }
        }
    }
}

torch::Tensor conv_transpose2d_forward(
    torch::Tensor x,
    torch::Tensor wt,
    c10::optional<torch::Tensor> bias_opt,
    int64_t H_out, int64_t W_out,
    int64_t pad_h, int64_t pad_w,
    int64_t KH, int64_t KW)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    TORCH_CHECK(x.is_contiguous() && wt.is_contiguous(), "inputs must be contiguous");

    const int N = x.size(0);
    const int C_in = x.size(1);
    const int H_in = x.size(2);
    const int W_in = x.size(3);
    const int C_out = wt.size(2);

    auto out = torch::empty({N, C_out, H_out, W_out}, x.options());
    const float* bias_ptr = nullptr;
    torch::Tensor bias;
    if (bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0) {
        bias = bias_opt->contiguous();
        bias_ptr = bias.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    bool fast = false;
    dim3 grid, block(16, 16);
    auto launch_cfg = [&](int CO) {
        grid = dim3((W_out + 63) / 64, (H_out + 15) / 16, N * (C_out / CO));
    };

    #define LAUNCH_TILED(KHV, KWV, COV)                                            \
        launch_cfg(COV);                                                           \
        tconv2d_tiled_kernel<KHV, KWV, COV><<<grid, block, 0, stream>>>(           \
            x.data_ptr<float>(), wt.data_ptr<float>(), bias_ptr,                   \
            out.data_ptr<float>(), C_in, H_in, W_in, C_out,                        \
            (int)H_out, (int)W_out, (int)pad_h, (int)pad_w);                       \
        fast = true;

    if (KH == 3 && KW == 5 && (C_out % 4) == 0)      { LAUNCH_TILED(3, 5, 4) }
    else if (KH == 3 && KW == 3 && (C_out % 4) == 0) { LAUNCH_TILED(3, 3, 4) }
    else if (KH == 5 && KW == 5 && (C_out % 4) == 0) { LAUNCH_TILED(5, 5, 4) }
    else if (KH == 1 && KW == 1 && (C_out % 4) == 0) { LAUNCH_TILED(1, 1, 4) }
    else if (KH == 1 && KW == 3 && (C_out % 4) == 0) { LAUNCH_TILED(1, 3, 4) }
    else if (KH == 3 && KW == 1 && (C_out % 4) == 0) { LAUNCH_TILED(3, 1, 4) }
    else if (KH == 5 && KW == 3 && (C_out % 4) == 0) { LAUNCH_TILED(5, 3, 4) }
    else if (KH == 3 && KW == 5 && (C_out % 2) == 0) { LAUNCH_TILED(3, 5, 2) }
    else if (KH == 3 && KW == 3 && (C_out % 2) == 0) { LAUNCH_TILED(3, 3, 2) }

    if (!fast) {
        long total = (long)N * C_out * H_out * W_out;
        long per_n = (long)C_out * H_out * W_out;
        int blocks_n = (int)((per_n + 255) / 256);
        dim3 g(blocks_n, N);
        tconv2d_generic_kernel<<<g, 256, 0, stream>>>(
            x.data_ptr<float>(), wt.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), C_in, H_in, W_in, C_out,
            (int)H_out, (int)W_out, (int)KH, (int)KW, (int)pad_h, (int)pad_w);
    }
    #undef LAUNCH_TILED
    return out;
}
"""

cpp_source = r"""
torch::Tensor conv_transpose2d_forward(
    torch::Tensor x,
    torch::Tensor wt,
    c10::optional<torch::Tensor> bias_opt,
    int64_t H_out, int64_t W_out,
    int64_t pad_h, int64_t pad_w,
    int64_t KH, int64_t KW);
"""

_tconv2d_cuda = load_inline(
    name="conv_transpose2d_tiled_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv_transpose2d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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

    Strategy
    --------
    * For the common case (stride=1, dilation=1, output_padding=0, groups=1,
      float32 CUDA input) the transposed convolution is rewritten as an
      equivalent full convolution and evaluated with a custom tiled CUDA
      kernel (shared-memory input tiling + register blocking).
    * The same equivalence also allows a cuDNN ``conv2d`` evaluation, and a
      one-time runtime autotune (executed on the first forward for a given
      input shape) picks the fastest among the custom kernel, the cuDNN conv2d
      path, and the native ``conv_transpose2d``.
    * Any unsupported configuration falls back to the native PyTorch op, so
      the module is a drop-in replacement.
    """

    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)

        kh, kw = _pair(kernel_size)
        self._kh, self._kw = kh, kw
        self._stride = _pair(stride)
        self._padding = _pair(padding)
        self._output_padding = _pair(output_padding)
        self._dilation = _pair(dilation)
        self._groups = groups
        self._in_channels = in_channels
        self._out_channels = out_channels
        # equivalent conv2d padding
        self._eff_pad = (self._dilation[0] * (kh - 1) - self._padding[0],
                         self._dilation[1] * (kw - 1) - self._padding[1])
        self._fast_ok = (self._stride == (1, 1) and self._dilation == (1, 1) and
                         self._output_padding == (0, 0) and groups == 1 and
                         self._eff_pad[0] >= 0 and self._eff_pad[1] >= 0)

        self._derived = None          # (wt_tap, w_flip, key)
        self._plans = {}              # input-shape key -> path name
        self.register_load_state_dict_post_hook(ModelNew._on_load)

    # -- cache management ----------------------------------------------------
    def _on_load(self, incompatible_keys):
        self._derived = None
        self._plans = {}

    def _weight_key(self):
        w = self.conv_transpose2d.weight
        return (w.data_ptr(), w._version, w.device, w.dtype)

    def _get_derived(self):
        key = self._weight_key()
        if self._derived is not None and self._derived[2] == key:
            return self._derived[0], self._derived[1]
        with torch.no_grad():
            w = self.conv_transpose2d.weight.detach()
            wf = w.flip([2, 3])
            # tap-major weights for the custom kernel: (KH*KW, C_in, C_out)
            wt = wf.permute(2, 3, 0, 1).contiguous().view(self._kh * self._kw,
                                                          self._in_channels,
                                                          self._out_channels)
            # flipped weights for the cuDNN conv2d path: (C_out, C_in, KH, KW)
            w_flip = wf.permute(1, 0, 2, 3).contiguous()
        self._derived = (wt, w_flip, key)
        return wt, w_flip

    # -- candidate paths ------------------------------------------------------
    def _run_custom(self, x, wt, bias, out_hw):
        return _tconv2d_cuda.conv_transpose2d_forward(
            x, wt, bias, out_hw[0], out_hw[1],
            self._eff_pad[0], self._eff_pad[1], self._kh, self._kw)

    def _run_cudnn(self, x, w_flip, bias):
        return F.conv2d(x, w_flip, bias, (1, 1), self._eff_pad, (1, 1), 1)

    def _run_orig(self, x):
        m = self.conv_transpose2d
        return F.conv_transpose2d(x, m.weight, m.bias, self._stride,
                                  self._padding, self._output_padding,
                                  self._groups, self._dilation)

    @staticmethod
    def _time_fn(fn, iters=5):
        for _ in range(2):
            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)

    def _autotune(self, x, wt, w_flip, bias, out_hw):
        cands = []
        ref = self._run_orig(x)

        def check(y):
            return y.shape == ref.shape and torch.allclose(y, ref, atol=1e-2, rtol=1e-2)

        try:
            y = self._run_custom(x, wt, bias, out_hw)
            torch.cuda.synchronize()
            if check(y):
                cands.append(("custom", lambda: self._run_custom(x, wt, bias, out_hw)))
        except Exception:
            pass
        try:
            y = self._run_cudnn(x, w_flip, bias)
            torch.cuda.synchronize()
            if check(y):
                cands.append(("cudnn", lambda: self._run_cudnn(x, w_flip, bias)))
        except Exception:
            pass
        cands.append(("orig", lambda: self._run_orig(x)))

        best, best_t = "orig", float("inf")
        for name, fn in cands:
            try:
                t = self._time_fn(fn)
            except Exception:
                continue
            if t < best_t:
                best, best_t = name, t
        return best

    # -- forward --------------------------------------------------------------
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self._fast_ok:
            return self._run_orig(x)

        m = self.conv_transpose2d
        if not (x.is_cuda and x.dtype == torch.float32):
            return self._run_orig(x)

        x = x.contiguous()
        wt, w_flip = self._get_derived()
        bias = m.bias.detach() if m.bias is not None else None

        H_in, W_in = x.shape[2], x.shape[3]
        H_out = (H_in - 1) * self._stride[0] - 2 * self._padding[0] + \
            self._dilation[0] * (self._kh - 1) + self._output_padding[0] + 1
        W_out = (W_in - 1) * self._stride[1] - 2 * self._padding[1] + \
            self._dilation[1] * (self._kw - 1) + self._output_padding[1] + 1

        key = (x.shape[0], H_in, W_in, x.device.index)
        path = self._plans.get(key)
        if path is None:
            path = self._autotune(x, wt, w_flip, bias, (H_out, W_out))
            self._plans[key] = path

        if path == "custom":
            return self._run_custom(x, wt, bias, (H_out, W_out))
        if path == "cudnn":
            return self._run_cudnn(x, w_flip, bias)
        return self._run_orig(x)
```