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

# The given ConvTranspose2d uses stride=(1,1), padding=(0,0), output_padding=(0,0),
# dilation=(1,1), groups=1. For this configuration the transposed convolution is
# mathematically identical to a regular convolution 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,wo-kw] * 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]
# We implement this "full-padding convolution" as a custom tiled CUDA kernel:
#   * each thread block produces a 32x16 spatial tile for 8 output channels
#   * the input tile (with halo / zero padding) is staged in shared memory
#   * weights are streamed through the read-only (L1) cache, broadcast to threads
#   * every thread accumulates 2 pixels x 8 output channels = 16 FMA chains

_cpp_src = "torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight);"

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

#define TILE_W 32
#define TILE_H 16
#define NTHREADS 256
#define OC_BLK 8

// W2[o,i,kh,kw] = Wt[i,o,KH-1-kh,KW-1-kw]  (Wt is ConvTranspose2d's weight [I,O,KH,KW])
__global__ void transform_weight_kernel(
    const float* __restrict__ wt,
    float* __restrict__ w2,
    int C_in, int C_out, int KH, int KW, int total_pad)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_pad) return;
    int kw = idx % KW;
    int kh = (idx / KW) % KH;
    int ci = (idx / (KW * KH)) % C_in;
    int co = idx / (KW * KH * C_in);
    float v = 0.f;
    if (co < C_out) {
        v = wt[(((size_t)ci * C_out + co) * KH + (KH - 1 - kh)) * KW + (KW - 1 - kw)];
    }
    w2[idx] = v;
}

// Direct convolution, stride 1, dilation 1, padding (KH-1, KW-1)  ("full" padding).
__global__ void conv_full_pad_kernel(
    const float* __restrict__ x,   // [N, C_in, H, W]
    const float* __restrict__ w2,  // [C_out_pad, C_in, KH, KW]
    float* __restrict__ y,         // [N, C_out, HO, WO]
    int C_in, int H, int W,
    int C_out, int KH, int KW,
    int HO, int WO, int OG)
{
    extern __shared__ float s_tile[];

    const int PAD_H = KH - 1;
    const int PAD_W = KW - 1;
    const int SH = TILE_H + KH - 1;
    const int SW = TILE_W + KW - 1;

    const int tx = threadIdx.x & (TILE_W - 1);
    const int ty = threadIdx.x >> 5;  // 0..7

    const int wo0 = blockIdx.x * TILE_W;
    const int ho0 = blockIdx.y * TILE_H;
    const int og  = blockIdx.z % OG;
    const int n   = blockIdx.z / OG;
    const int o0  = og * OC_BLK;

    const int ho = ho0 + ty;
    const int wo = wo0 + tx;

    float acc0[OC_BLK], acc1[OC_BLK];
    #pragma unroll
    for (int j = 0; j < OC_BLK; ++j) { acc0[j] = 0.f; acc1[j] = 0.f; }

    const size_t KK   = (size_t)KH * KW;
    const size_t CiKK = (size_t)C_in * KK;
    const float* wg   = w2 + (size_t)o0 * CiKK;

    for (int i = 0; i < C_in; ++i) {
        // stage the (padded) input tile for this input channel in shared memory
        const float* xb = x + ((size_t)n * C_in + i) * (size_t)H * W;
        for (int idx = threadIdx.x; idx < SH * SW; idx += NTHREADS) {
            int sy = idx / SW;
            int sx = idx - sy * SW;
            int gy = ho0 + sy - PAD_H;
            int gx = wo0 + sx - PAD_W;
            float v = 0.f;
            if (gy >= 0 && gy < H && gx >= 0 && gx < W) v = xb[(size_t)gy * W + gx];
            s_tile[idx] = v;
        }
        __syncthreads();

        const float* wi = wg + (size_t)i * KK;
        for (int kh = 0; kh < KH; ++kh) {
            const float* r0 = s_tile + (ty + kh) * SW + tx;
            const float* r1 = s_tile + (ty + (TILE_H / 2) + kh) * SW + tx;
            const float* wk = wi + (size_t)kh * KW;
            for (int kw = 0; kw < KW; ++kw) {
                float xv0 = r0[kw];
                float xv1 = r1[kw];
                float wv[OC_BLK];
                #pragma unroll
                for (int j = 0; j < OC_BLK; ++j) wv[j] = __ldg(wk + (size_t)j * CiKK + kw);
                #pragma unroll
                for (int j = 0; j < OC_BLK; ++j) {
                    acc0[j] = fmaf(xv0, wv[j], acc0[j]);
                    acc1[j] = fmaf(xv1, wv[j], acc1[j]);
                }
            }
        }
        __syncthreads();
    }

    if (wo < WO) {
        const int oc = min(OC_BLK, C_out - o0);
        if (ho < HO) {
            float* yb = y + ((size_t)n * C_out + o0) * (size_t)HO * WO + (size_t)ho * WO + wo;
            const size_t stride_ch = (size_t)HO * WO;
            for (int j = 0; j < oc; ++j) yb[(size_t)j * stride_ch] = acc0[j];
        }
        int ho1 = ho + (TILE_H / 2);
        if (ho1 < HO) {
            float* yb = y + ((size_t)n * C_out + o0) * (size_t)HO * WO + (size_t)ho1 * WO + wo;
            const size_t stride_ch = (size_t)HO * WO;
            for (int j = 0; j < oc; ++j) yb[(size_t)j * stride_ch] = acc1[j];
        }
    }
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "x and weight must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();

    const int N     = (int)x_c.size(0);
    const int C_in  = (int)x_c.size(1);
    const int H     = (int)x_c.size(2);
    const int W     = (int)x_c.size(3);
    const int C_out = (int)w_c.size(1);   // weight: [C_in, C_out, KH, KW]
    const int KH    = (int)w_c.size(2);
    const int KW    = (int)w_c.size(3);

    const int HO = H + KH - 1;
    const int WO = W + KW - 1;

    const int OG = (C_out + OC_BLK - 1) / OC_BLK;
    const int C_out_pad = OG * OC_BLK;

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

    // 1) transform weights: swap channel axes + flip spatial axes
    auto w2 = torch::empty({C_out_pad, C_in, KH, KW}, x_c.options());
    const int total_pad = C_out_pad * C_in * KH * KW;
    transform_weight_kernel<<<(total_pad + 255) / 256, 256, 0, stream>>>(
        w_c.data_ptr<float>(), w2.data_ptr<float>(), C_in, C_out, KH, KW, total_pad);

    // 2) direct convolution with "full" padding
    auto y = torch::empty({N, C_out, (long)HO, (long)WO}, x_c.options());

    dim3 blocks((WO + TILE_W - 1) / TILE_W,
                (HO + TILE_H - 1) / TILE_H,
                N * OG);
    TORCH_CHECK(N * OG <= 65535, "grid z-dimension too large");

    const size_t shmem = (size_t)(TILE_H + KH - 1) * (TILE_W + KW - 1) * sizeof(float);
    if (shmem > 48 * 1024) {
        cudaFuncSetAttribute(conv_full_pad_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)shmem);
    }

    conv_full_pad_kernel<<<blocks, NTHREADS, shmem, stream>>>(
        x_c.data_ptr<float>(), w2.data_ptr<float>(), y.data_ptr<float>(),
        C_in, H, W, C_out, KH, KW, HO, WO, OG);

    TORCH_CHECK(cudaGetLastError() == cudaSuccess, "conv kernel launch failed");
    return y;
}
"""

conv_transpose_ext = load_inline(
    name="conv_transpose2d_full_conv",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv_transpose2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    Uses the mathematical equivalence between stride-1 transposed convolution and
    a regular convolution with full padding plus permuted/flipped weights, and
    evaluates it with a custom tiled CUDA kernel.
    """

    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 (and parameter names / 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)

        stride = (stride, stride) if isinstance(stride, int) else tuple(stride)
        padding = (padding, padding) if isinstance(padding, int) else tuple(padding)
        output_padding = (output_padding, output_padding) if isinstance(output_padding, int) else tuple(output_padding)
        dilation = (dilation, dilation) if isinstance(dilation, int) else tuple(dilation)

        self._use_custom = (
            tuple(stride) == (1, 1)
            and tuple(padding) == (0, 0)
            and tuple(output_padding) == (0, 0)
            and tuple(dilation) == (1, 1)
            and groups == 1
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv_transpose2d.weight
        b = self.conv_transpose2d.bias
        if (self._use_custom and x.is_cuda and x.dtype == torch.float32
                and w.is_cuda and w.dtype == torch.float32):
            oc = w.shape[1]
            og = (oc + 7) // 8
            if x.shape[0] * og <= 65535:
                y = conv_transpose_ext.conv_transpose2d_forward(x, w)
                if b is not None:
                    y = y + b.view(1, -1, 1, 1)
                return y
        # Fallback to the reference implementation for any other configuration.
        return self.conv_transpose2d(x)
```