```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 for stride-1 / dilation-1 / groups-1 ConvTranspose2d.
#
# For stride 1, transposed conv is a gather (correlation) over the input:
#   out[n, co, oy, ox] = sum_{ci, ky, kx} x[n, ci, oy + padH - ky, ox + padW - kx]
#                                        * w[ci, co, ky, kx]
# The kernel is output-stationary: each thread block owns a 13x26 spatial tile
# of the output for 8 output channels of one batch element. The required input
# halo tile (including the kernel - 1 border) is staged in shared memory in
# chunks of 8 input channels, and every loaded input value is reused across all
# 8 output channels kept in registers.

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

#define TILE_H 13
#define TILE_W 26
#define HALF_W (TILE_W / 2)   // 13
#define CO_BLOCK 8
#define CI_CHUNK 8
#define NTHREADS (TILE_H * HALF_W)  // 169

__global__ void __launch_bounds__(NTHREADS)
tconv2d_s1_kernel(const float* __restrict__ x,
                  const float* __restrict__ w,
                  const float* __restrict__ bias,
                  float* __restrict__ y,
                  int CI, int CO,
                  int HI, int WI, int HO, int WO,
                  int KH, int KW, int padH, int padW,
                  int tilesX) {
    extern __shared__ float smem[];
    const int haloH = TILE_H + KH - 1;
    const int haloW = TILE_W + KW - 1;
    const int inFloats = CI_CHUNK * haloH * haloW;
    const int inFloatsPad = (inFloats + 3) & ~3;  // keep s_w 16B aligned
    float* s_in = smem;
    float* s_w  = smem + inFloatsPad;

    const int tid = threadIdx.x;
    const int nthreads = blockDim.x;

    const int tile = blockIdx.x;
    const int ty0 = tile / tilesX;
    const int tx0 = tile - ty0 * tilesX;
    const int oy0 = ty0 * TILE_H;
    const int ox0 = tx0 * TILE_W;
    const int co0 = blockIdx.y * CO_BLOCK;
    const int n   = blockIdx.z;

    const int ty = tid / HALF_W;
    const int tx = tid - ty * HALF_W;
    const int oy  = oy0 + ty;
    const int oxa = ox0 + tx;
    const int oxb = ox0 + tx + HALF_W;
    const bool validY = (oy < HO);
    const bool validA = validY && (oxa < WO);
    const bool validB = validY && (oxb < WO);

    float accA[CO_BLOCK];
    float accB[CO_BLOCK];
    #pragma unroll
    for (int j = 0; j < CO_BLOCK; ++j) { accA[j] = 0.f; accB[j] = 0.f; }

    const long xBatchOff = (long)n * CI * HI * WI;
    const int coLimit = CO - co0;

    const int rOff = oy0 + padH - (KH - 1);
    const int cOff = ox0 + padW - (KW - 1);

    for (int ci0 = 0; ci0 < CI; ci0 += CI_CHUNK) {
        const int cchunk = min(CI_CHUNK, CI - ci0);
        __syncthreads();
        // Stage input halo tile for channels [ci0, ci0 + cchunk)
        {
            const int total = cchunk * haloH * haloW;
            const int plane = haloH * haloW;
            for (int i = tid; i < total; i += nthreads) {
                const int c = i / plane;
                const int rem = i - c * plane;
                const int r = rem / haloW;
                const int cc = rem - r * haloW;
                const int iy = rOff + r;
                const int ix = cOff + cc;
                float v = 0.f;
                if ((unsigned)iy < (unsigned)HI && (unsigned)ix < (unsigned)WI)
                    v = __ldg(x + xBatchOff + ((long)(ci0 + c) * HI + iy) * WI + ix);
                s_in[i] = v;
            }
        }
        // Stage weights as [c][ky][kx][co] (co fastest, float4 friendly)
        {
            const int kkw = KH * KW;
            const int total = cchunk * kkw * CO_BLOCK;
            for (int i = tid; i < total; i += nthreads) {
                const int j = i & (CO_BLOCK - 1);
                const int rem = i >> 3;  // log2(CO_BLOCK)
                const int kx = rem % KW;
                const int rem2 = rem / KW;
                const int ky = rem2 % KH;
                const int c = rem2 / KH;
                float v = 0.f;
                if (j < coLimit)
                    v = __ldg(w + (((long)(ci0 + c) * CO + co0 + j) * KH + ky) * KW + kx);
                s_w[i] = v;
            }
        }
        __syncthreads();

        if (!validY) continue;  // all threads still hit the sync at loop top

        for (int c = 0; c < cchunk; ++c) {
            const float* in_c = s_in + c * haloH * haloW;
            const float* w_c  = s_w + c * KH * KW * CO_BLOCK;
            for (int ky = 0; ky < KH; ++ky) {
                const int r = ty + (KH - 1) - ky;
                const float* rowA = in_c + r * haloW + tx;
                const float* rowB = rowA + HALF_W;
                const float* wrow = w_c + ky * KW * CO_BLOCK;
                for (int kx = 0; kx < KW; ++kx) {
                    const int off = (KW - 1) - kx;
                    const float a = rowA[off];
                    const float b = rowB[off];
                    const float4 wlo = *(const float4*)(wrow + kx * CO_BLOCK);
                    const float4 whi = *(const float4*)(wrow + kx * CO_BLOCK + 4);
                    accA[0] += a * wlo.x; accA[1] += a * wlo.y;
                    accA[2] += a * wlo.z; accA[3] += a * wlo.w;
                    accA[4] += a * whi.x; accA[5] += a * whi.y;
                    accA[6] += a * whi.z; accA[7] += a * whi.w;
                    accB[0] += b * wlo.x; accB[1] += b * wlo.y;
                    accB[2] += b * wlo.z; accB[3] += b * wlo.w;
                    accB[4] += b * whi.x; accB[5] += b * whi.y;
                    accB[6] += b * whi.z; accB[7] += b * whi.w;
                }
            }
        }
    }

    if (!validY) return;
    #pragma unroll
    for (int j = 0; j < CO_BLOCK; ++j) {
        if (j < coLimit) {
            const int co = co0 + j;
            const float bv = (bias != nullptr) ? __ldg(bias + co) : 0.f;
            float* out_ptr = y + (((long)n * CO + co) * HO + oy) * WO;
            if (validA) out_ptr[oxa] = accA[j] + bv;
            if (validB) out_ptr[oxb] = accB[j] + bv;
        }
    }
}

torch::Tensor conv_transpose2d_s1(torch::Tensor x,
                                  torch::Tensor weight,
                                  c10::optional<torch::Tensor> bias_opt,
                                  int64_t padH, int64_t padW) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == at::kFloat, "w must be CUDA float32");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    const int N  = xc.size(0);
    const int CI = xc.size(1);
    const int HI = xc.size(2);
    const int WI = xc.size(3);
    const int CO = wc.size(1);
    const int KH = wc.size(2);
    const int KW = wc.size(3);
    // stride = 1, dilation = 1, output_padding = 0
    const int HO = HI + KH - 1 - 2 * (int)padH;
    const int WO = WI + KW - 1 - 2 * (int)padW;
    TORCH_CHECK(HO > 0 && WO > 0, "invalid output size");

    auto y = at::empty({N, CO, HO, WO}, xc.options());

    const int tilesX = (WO + TILE_W - 1) / TILE_W;
    const int tilesY = (HO + TILE_H - 1) / TILE_H;
    dim3 grid(tilesY * tilesX, (CO + CO_BLOCK - 1) / CO_BLOCK, N);
    dim3 block(NTHREADS);

    const int haloH = TILE_H + KH - 1;
    const int haloW = TILE_W + KW - 1;
    const int inFloats = CI_CHUNK * haloH * haloW;
    const int inFloatsPad = (inFloats + 3) & ~3;
    const size_t shmem = (size_t)(inFloatsPad + CI_CHUNK * KH * KW * CO_BLOCK) * sizeof(float);

    const float* bias_ptr = nullptr;
    torch::Tensor bc;
    if (bias_opt.has_value() && bias_opt->defined()) {
        bc = bias_opt->contiguous();
        bias_ptr = bc.data_ptr<float>();
    }

    auto stream = at::cuda::getCurrentCUDAStream();
    tconv2d_s1_kernel<<<grid, block, shmem, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
        y.data_ptr<float>(),
        CI, CO, HI, WI, HO, WO, KH, KW,
        (int)padH, (int)padW, tilesX);
    return y;
}
"""

tconv2d_cpp_source = (
    "torch::Tensor conv_transpose2d_s1(torch::Tensor x, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias_opt, int64_t padH, int64_t padW);"
)

tconv2d_ext = load_inline(
    name="tconv2d_s1_ext",
    cpp_sources=tconv2d_cpp_source,
    cuda_sources=tconv2d_cuda_source,
    functions=["conv_transpose2d_s1"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose2d accelerated with a custom CUDA kernel.

    The custom kernel handles the common case of stride=1, dilation=1,
    groups=1, output_padding=0 (which covers the target architecture). Any
    other configuration falls back to the native PyTorch implementation so the
    module remains fully general. Parameters are stored in an identically
    named nn.ConvTranspose2d submodule so state_dicts are interchangeable with
    the reference Model.
    """

    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)

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

        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.output_padding = _pair(output_padding)
        self.dilation = _pair(dilation)
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (x.is_cuda and x.dtype == torch.float32
                and self.stride == (1, 1) and self.dilation == (1, 1)
                and self.groups == 1 and self.output_padding == (0, 0)):
            return tconv2d_ext.conv_transpose2d_s1(
                x, self.conv_transpose2d.weight,
                self.conv_transpose2d.bias,
                self.padding[0], self.padding[1])
        return self.conv_transpose2d(x)
```