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

# Custom CUDA kernel: direct transposed-2D-convolution (stride=1, padding=0,
# 3x3 kernel, groups=1) computed in "gather" form, i.e. each output pixel
# accumulates its 3x3xC_in contributions directly. This avoids the
# GEMM + col2im (atomic scatter) path entirely. The kernel uses shared-memory
# input tiles and computes a 32x32 spatial tile x 8 output channels per block,
# with 2x2 outputs per thread for high arithmetic intensity.

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

#define TILE 32
#define IN_TILE 34
#define COT 8
#define NTHREADS 256

__global__ void __launch_bounds__(NTHREADS)
conv_transpose2d_k3_s1_p0_kernel(const float* __restrict__ X,
                                 const float* __restrict__ Wt,
                                 const float* __restrict__ Bias,
                                 float* __restrict__ Y,
                                 int C_in, int H, int W,
                                 int C_out, int H_out, int W_out) {
    __shared__ float in_tile[IN_TILE][IN_TILE];
    __shared__ float w_tile[COT][9];

    const int tid = threadIdx.x;
    const int tx = tid & 15;
    const int ty = tid >> 4;

    const int ow0 = blockIdx.x * TILE;
    const int oh0 = blockIdx.y * TILE;
    const int n   = blockIdx.z / (C_out / COT);
    const int co0 = (blockIdx.z % (C_out / COT)) * COT;

    const int r0 = oh0 - 2;
    const int c0 = ow0 - 2;

    float acc[COT][2][2];
    #pragma unroll
    for (int c = 0; c < COT; ++c)
        #pragma unroll
        for (int i = 0; i < 2; ++i)
            #pragma unroll
            for (int j = 0; j < 2; ++j)
                acc[c][i][j] = 0.0f;

    for (int ci = 0; ci < C_in; ++ci) {
        // Load the (TILE+2)x(TILE+2) input halo tile for this input channel.
        const float* xbase = X + ((long)(n * C_in + ci)) * H * W;
        #pragma unroll
        for (int k = 0; k < (IN_TILE * IN_TILE + NTHREADS - 1) / NTHREADS; ++k) {
            int idx = tid + k * NTHREADS;
            if (idx < IN_TILE * IN_TILE) {
                int tr = idx / IN_TILE;
                int tc = idx - tr * IN_TILE;
                int gr = r0 + tr;
                int gc = c0 + tc;
                float v = 0.0f;
                if (gr >= 0 && gr < H && gc >= 0 && gc < W)
                    v = xbase[(long)gr * W + gc];
                in_tile[tr][tc] = v;
            }
        }
        // Load the 8x9 weight slice for this input channel.
        if (tid < COT * 9) {
            int cl  = tid / 9;
            int tap = tid - cl * 9;
            w_tile[cl][tap] =
                Wt[(((long)(ci * C_out + co0 + cl)) * 3 + tap / 3) * 3 + (tap % 3)];
        }
        __syncthreads();

        // ConvTranspose2d (stride=1, pad=0):
        //   Y[n, co, h, w] = sum_{ci,p,q} Wt[ci,co,p,q] * X[n,ci,h-p,w-q]
        // In tile coords, input row = (h - p) - (oh0 - 2) = (h - oh0) + 2 - p.
        #pragma unroll
        for (int p = 0; p < 3; ++p) {
            #pragma unroll
            for (int q = 0; q < 3; ++q) {
                float xv[2][2];
                #pragma unroll
                for (int i = 0; i < 2; ++i)
                    #pragma unroll
                    for (int j = 0; j < 2; ++j)
                        xv[i][j] = in_tile[ty * 2 + i + (2 - p)][tx * 2 + j + (2 - q)];
                #pragma unroll
                for (int c = 0; c < COT; ++c) {
                    float wv = w_tile[c][p * 3 + q];
                    acc[c][0][0] = fmaf(wv, xv[0][0], acc[c][0][0]);
                    acc[c][0][1] = fmaf(wv, xv[0][1], acc[c][0][1]);
                    acc[c][1][0] = fmaf(wv, xv[1][0], acc[c][1][0]);
                    acc[c][1][1] = fmaf(wv, xv[1][1], acc[c][1][1]);
                }
            }
        }
        __syncthreads();
    }

    #pragma unroll
    for (int c = 0; c < COT; ++c) {
        int co = co0 + c;
        float bv = (Bias != nullptr) ? Bias[co] : 0.0f;
        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            int oh = oh0 + ty * 2 + i;
            if (oh >= H_out) continue;
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                int ow = ow0 + tx * 2 + j;
                if (ow >= W_out) continue;
                Y[(((long)(n * C_out + co)) * H_out + oh) * W_out + ow] =
                    acc[c][i][j] + bv;
            }
        }
    }
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight,
                                       torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    const int N     = x.size(0);
    const int C_in  = x.size(1);
    const int H     = x.size(2);
    const int W     = x.size(3);
    const int C_out = weight.size(1);
    const int H_out = H + 2;
    const int W_out = W + 2;

    auto y = torch::empty({N, C_out, H_out, W_out}, x.options());

    dim3 block(NTHREADS);
    dim3 grid((W_out + TILE - 1) / TILE, (H_out + TILE - 1) / TILE,
              N * (C_out / COT));
    const float* bptr =
        (bias.defined() && bias.numel() > 0) ? bias.data_ptr<float>() : nullptr;

    auto stream = at::cuda::getCurrentCUDAStream();
    conv_transpose2d_k3_s1_p0_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bptr,
        y.data_ptr<float>(), C_in, H, W, C_out, H_out, W_out);
    return y;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized transposed 2D convolution. Uses a custom tiled CUDA kernel for
    the common fast path (3x3 kernel, stride=1, padding=0, output_padding=0,
    groups=1, dilation=1, float32 CUDA input) and falls back to the native
    PyTorch implementation otherwise.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, output_padding: int = 0,
                 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, groups=groups,
            bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        m = self.conv_transpose2d
        if (x.is_cuda and x.dtype == torch.float32 and
                m.kernel_size == (3, 3) and m.stride == (1, 1) and
                m.padding == (0, 0) and m.output_padding == (0, 0) and
                m.groups == 1 and m.dilation == (1, 1) and
                x.size(1) == m.in_channels and (m.out_channels % 8 == 0)):
            bias = m.bias
            if bias is None:
                bias = torch.empty(0, device=x.device, dtype=x.dtype)
            return fast_conv_transpose2d.conv_transpose2d_forward(
                x.contiguous(), m.weight.contiguous(), bias)
        return m(x)
```