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

_cpp_src = """
torch::Tensor conv_transpose2d_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, int64_t pad_h, int64_t pad_w);
"""

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

#define OC_TILE 8
#define TILE_X 32
#define TILE_Y 16
#define THREADS_X 16
#define THREADS_Y 8

// Rearranges the transposed-conv weight from [C_in][C_out][K][K] to
// [C_in][K][K][C_out] so the 8 output-channel weights are contiguous
// (allowing float4 broadcast loads inside the main kernel).
__global__ void transpose_weight_kernel(const float* __restrict__ w,
                                        float* __restrict__ wt,
                                        int C_in, int C_out, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = C_in * C_out * K * K;
    if (idx < total) {
        int q  = idx % K;
        int p  = (idx / K) % K;
        int oc = (idx / (K * K)) % C_out;
        int ic =  idx / (K * K * C_out);
        wt[((ic * K + p) * K + q) * C_out + oc] = w[idx];
    }
}

// Optimized direct transposed-convolution kernel for stride=1, groups=1.
// out[b, o, i, j] = sum_c sum_{p,q} in[b, c, i + pad_h - p, j + pad_w - q] * W[c, o, p, q]
// Each block computes a TILE_Y x TILE_X spatial tile for OC_TILE output channels.
// Each thread computes a 2x2 spatial micro-tile x OC_TILE channels (32 accumulators).
template <int K>
__global__ void conv_transpose2d_s1_kernel(
        const float* __restrict__ input,
        const float* __restrict__ wt,
        const float* __restrict__ bias,
        float* __restrict__ output,
        int B, int C_in, int C_out, int H, int W,
        int H_out, int W_out, int pad_h, int pad_w) {
    constexpr int SM_H = TILE_Y + K - 1;
    constexpr int SM_W = TILE_X + K - 1;
    __shared__ float tile[SM_H * SM_W];

    const int tx  = threadIdx.x;
    const int ty  = threadIdx.y;
    const int tid = ty * THREADS_X + tx;

    const int ox0 = blockIdx.x * TILE_X;
    const int oy0 = blockIdx.y * TILE_Y;
    const int oc0 = (blockIdx.z % (C_out / OC_TILE)) * OC_TILE;
    const int b   =  blockIdx.z / (C_out / OC_TILE);

    float acc[2][2][OC_TILE];
    #pragma unroll
    for (int dy = 0; dy < 2; ++dy)
        #pragma unroll
        for (int dx = 0; dx < 2; ++dx)
            #pragma unroll
            for (int o = 0; o < OC_TILE; ++o) acc[dy][dx][o] = 0.f;

    const int in_y0 = oy0 + pad_h - (K - 1);
    const int in_x0 = ox0 + pad_w - (K - 1);
    const float* in_b = input + (size_t)b * C_in * H * W;

    for (int ic = 0; ic < C_in; ++ic) {
        const float* in_ic = in_b + (size_t)ic * H * W;
        for (int idx = tid; idx < SM_H * SM_W; idx += THREADS_X * THREADS_Y) {
            int sy = idx / SM_W;
            int sx = idx % SM_W;
            int iy = in_y0 + sy;
            int ix = in_x0 + sx;
            float v = 0.f;
            if (iy >= 0 && iy < H && ix >= 0 && ix < W)
                v = in_ic[(size_t)iy * W + ix];
            tile[idx] = v;
        }
        __syncthreads();

        const float* w_ic = wt + ((size_t)ic * K * K) * C_out + oc0;
        const int ry = ty * 2;
        const int rx = tx * 2;

        #pragma unroll
        for (int p = 0; p < K; ++p) {
            #pragma unroll
            for (int q = 0; q < K; ++q) {
                float iv[2][2];
                #pragma unroll
                for (int dy = 0; dy < 2; ++dy)
                    #pragma unroll
                    for (int dx = 0; dx < 2; ++dx)
                        iv[dy][dx] = tile[(ry + dy + (K - 1 - p)) * SM_W + (rx + dx + (K - 1 - q))];

                const float4 wa = *reinterpret_cast<const float4*>(w_ic + (p * K + q) * C_out);
                const float4 wb = *reinterpret_cast<const float4*>(w_ic + (p * K + q) * C_out + 4);
                const float wv[OC_TILE] = {wa.x, wa.y, wa.z, wa.w, wb.x, wb.y, wb.z, wb.w};
                #pragma unroll
                for (int o = 0; o < OC_TILE; ++o)
                    #pragma unroll
                    for (int dy = 0; dy < 2; ++dy)
                        #pragma unroll
                        for (int dx = 0; dx < 2; ++dx)
                            acc[dy][dx][o] = fmaf(iv[dy][dx], wv[o], acc[dy][dx][o]);
            }
        }
        __syncthreads();
    }

    const int gy = oy0 + ty * 2;
    const int gx = ox0 + tx * 2;
    #pragma unroll
    for (int dy = 0; dy < 2; ++dy) {
        const int y = gy + dy;
        if (y >= H_out) continue;
        #pragma unroll
        for (int dx = 0; dx < 2; ++dx) {
            const int x = gx + dx;
            if (x >= W_out) continue;
            #pragma unroll
            for (int o = 0; o < OC_TILE; ++o) {
                float v = acc[dy][dx][o];
                if (bias != nullptr) v += bias[oc0 + o];
                output[(((size_t)b * C_out + (oc0 + o)) * H_out + y) * W_out + x] = v;
            }
        }
    }
}

// Generic fallback kernel for arbitrary (square) kernel sizes, stride=1, groups=1.
__global__ void conv_transpose2d_s1_generic(
        const float* __restrict__ input,
        const float* __restrict__ wt,
        const float* __restrict__ bias,
        float* __restrict__ output,
        int B, int C_in, int C_out, int H, int W, int K,
        int H_out, int W_out, int pad_h, int pad_w) {
    extern __shared__ float tile[];
    const int GT   = 16;
    const int SM_H = GT + K - 1;
    const int SM_W = GT + K - 1;
    const int tx  = threadIdx.x;
    const int ty  = threadIdx.y;
    const int tid = ty * GT + tx;

    const int ox0 = blockIdx.x * GT;
    const int oy0 = blockIdx.y * GT;
    const int oc0 = (blockIdx.z % (C_out / OC_TILE)) * OC_TILE;
    const int b   =  blockIdx.z / (C_out / OC_TILE);

    float acc[OC_TILE];
    #pragma unroll
    for (int o = 0; o < OC_TILE; ++o) acc[o] = 0.f;

    const int in_y0 = oy0 + pad_h - (K - 1);
    const int in_x0 = ox0 + pad_w - (K - 1);
    const float* in_b = input + (size_t)b * C_in * H * W;

    for (int ic = 0; ic < C_in; ++ic) {
        const float* in_ic = in_b + (size_t)ic * H * W;
        for (int idx = tid; idx < SM_H * SM_W; idx += GT * GT) {
            int sy = idx / SM_W;
            int sx = idx % SM_W;
            int iy = in_y0 + sy;
            int ix = in_x0 + sx;
            float v = 0.f;
            if (iy >= 0 && iy < H && ix >= 0 && ix < W)
                v = in_ic[(size_t)iy * W + ix];
            tile[idx] = v;
        }
        __syncthreads();

        const float* w_ic = wt + ((size_t)ic * K * K) * C_out + oc0;
        for (int p = 0; p < K; ++p) {
            for (int q = 0; q < K; ++q) {
                float iv = tile[(ty + (K - 1 - p)) * SM_W + (tx + (K - 1 - q))];
                const float* wp = w_ic + (p * K + q) * C_out;
                #pragma unroll
                for (int o = 0; o < OC_TILE; ++o)
                    acc[o] = fmaf(iv, wp[o], acc[o]);
            }
        }
        __syncthreads();
    }

    const int y = oy0 + ty;
    const int x = ox0 + tx;
    if (y < H_out && x < W_out) {
        #pragma unroll
        for (int o = 0; o < OC_TILE; ++o) {
            float v = acc[o];
            if (bias != nullptr) v += bias[oc0 + o];
            output[(((size_t)b * C_out + (oc0 + o)) * H_out + y) * W_out + x] = v;
        }
    }
}

torch::Tensor conv_transpose2d_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, int64_t pad_h, int64_t pad_w) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    c10::cuda::CUDAGuard device_guard(input.device());

    auto x = input.contiguous();
    auto w = weight.contiguous();

    const int B     = 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 = w.size(1);
    const int K     = w.size(2);
    TORCH_CHECK(w.size(3) == K, "kernel must be square");
    TORCH_CHECK(C_out % OC_TILE == 0, "out_channels must be divisible by 8");

    const int H_out = H + K - 1 - 2 * (int)pad_h;
    const int W_out = W + K - 1 - 2 * (int)pad_w;
    TORCH_CHECK(H_out > 0 && W_out > 0, "invalid output size");

    auto out = at::empty({(int64_t)B, (int64_t)C_out, (int64_t)H_out, (int64_t)W_out}, x.options());
    auto wt  = at::empty({(int64_t)C_in, (int64_t)K, (int64_t)K, (int64_t)C_out}, x.options());

    const float* bias_ptr = (bias.defined() && bias.numel() == C_out) ? bias.data_ptr<float>() : nullptr;
    auto stream = at::cuda::getCurrentCUDAStream();

    {
        const int total   = C_in * C_out * K * K;
        const int threads = 256;
        const int blocks  = (total + threads - 1) / threads;
        transpose_weight_kernel<<<blocks, threads, 0, stream>>>(
            w.data_ptr<float>(), wt.data_ptr<float>(), C_in, C_out, K);
    }

    const int gz = B * (C_out / OC_TILE);
    const float* in_p  = x.data_ptr<float>();
    const float* wt_p  = wt.data_ptr<float>();
    float* out_p       = out.data_ptr<float>();

    if (K >= 1 && K <= 5) {
        dim3 block(THREADS_X, THREADS_Y);
        dim3 grid((W_out + TILE_X - 1) / TILE_X, (H_out + TILE_Y - 1) / TILE_Y, gz);
        switch (K) {
            case 1: conv_transpose2d_s1_kernel<1><<<grid, block, 0, stream>>>(in_p, wt_p, bias_ptr, out_p, B, C_in, C_out, H, W, H_out, W_out, (int)pad_h, (int)pad_w); break;
            case 2: conv_transpose2d_s1_kernel<2><<<grid, block, 0, stream>>>(in_p, wt_p, bias_ptr, out_p, B, C_in, C_out, H, W, H_out, W_out, (int)pad_h, (int)pad_w); break;
            case 3: conv_transpose2d_s1_kernel<3><<<grid, block, 0, stream>>>(in_p, wt_p, bias_ptr, out_p, B, C_in, C_out, H, W, H_out, W_out, (int)pad_h, (int)pad_w); break;
            case 4: conv_transpose2d_s1_kernel<4><<<grid, block, 0, stream>>>(in_p, wt_p, bias_ptr, out_p, B, C_in, C_out, H, W, H_out, W_out, (int)pad_h, (int)pad_w); break;
            case 5: conv_transpose2d_s1_kernel<5><<<grid, block, 0, stream>>>(in_p, wt_p, bias_ptr, out_p, B, C_in, C_out, H, W, H_out, W_out, (int)pad_h, (int)pad_w); break;
        }
    } else {
        const int GT = 16;
        dim3 block(GT, GT);
        dim3 grid((W_out + GT - 1) / GT, (H_out + GT - 1) / GT, gz);
        size_t smem = (size_t)(GT + K - 1) * (GT + K - 1) * sizeof(float);
        conv_transpose2d_s1_generic<<<grid, block, smem, stream>>>(
            in_p, wt_p, bias_ptr, out_p, B, C_in, C_out, H, W, K, H_out, W_out, (int)pad_h, (int)pad_w);
    }
    return out;
}
"""

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


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

    Fast custom path is used when: CUDA float32 input, stride == 1, groups == 1,
    dilation == 1, output_padding == 0, square kernel, out_channels % 8 == 0.
    All other configurations fall back to the standard PyTorch implementation.
    """
    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.groups == 1
                and m.stride[0] == 1 and m.stride[1] == 1
                and m.dilation[0] == 1 and m.dilation[1] == 1
                and m.output_padding[0] == 0 and m.output_padding[1] == 0
                and m.kernel_size[0] == m.kernel_size[1]
                and m.out_channels % 8 == 0):
            bias = m.bias if m.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
            return _ext.conv_transpose2d_forward(
                x, m.weight, bias, m.padding[0], m.padding[1])
        return self.conv_transpose2d(x)
```