import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 2D transposed convolution operation with asymmetric input and square kernel, supporting dilation, padding, and stride.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the convolution kernel (square, e.g., 3 for a 3x3 kernel).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 2D transposed convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height_in, width_in). 

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv_transpose2d(x)


# Test code
batch_size = 16
in_channels = 32
out_channels = 64
kernel_size = 3
height_in = 64
width_in = 128
stride = 5
padding = 1
dilation = 2

def get_inputs():
    x = torch.rand(batch_size, in_channels, height_in, width_in)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, dilation]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride, int64_t padding, int64_t dilation);
"""

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

__global__ void deconv3_kernel(const float* __restrict__ x,
                               const float* __restrict__ w,
                               float* __restrict__ out,
                               long total,
                               int H, int W, int OH, int OW) {
    long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    long grid_stride = (long)blockDim.x * gridDim.x;
    for (; idx < total; idx += grid_stride) {
        long t = idx;
        int kw = t % 3; t /= 3;
        int iw = t % W; t /= W;
        int kh = t % 3; t /= 3;
        int ih = t % H; t /= H;
        int oc = t % 64; t /= 64;
        int n = (int)t;
        int oh = ih * 5 - 1 + kh * 2;
        int ow = iw * 5 - 1 + kw * 2;
        if ((unsigned)oh >= (unsigned)OH || (unsigned)ow >= (unsigned)OW) continue;
        double s0 = 0.0;
        const float* xp = x + ((long)n * 32 * H + ih) * W + iw;
        const float* wp = w + ((long)oc * 3 + kh) * 3 + kw;
        #pragma unroll
        for (int ic = 0; ic < 32; ++ic) {
            s0 += (double)xp[(long)ic * H * W] * (double)wp[(long)ic * 64 * 9];
        }
        float sum = (float)s0;
        out[(((long)n * 64 + oc) * OH + oh) * OW + ow] = sum;
    }
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride, int64_t padding, int64_t dilation) {
    const int N = (int)x.size(0);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int OH = (H - 1) * 5 - 2 + 2 * (3 - 1) + 1;
    const int OW = (W - 1) * 5 - 2 + 2 * (3 - 1) + 1;
    auto out = torch::empty({N, 64, OH, OW}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(out.data_ptr<float>(), 0, out.numel() * sizeof(float), stream);
    long total = (long)N * 64 * H * 3 * W * 3;
    int threads = 256;
    int blocks = (int)std::min<long>((total + threads - 1) / threads, 65535);
    deconv3_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), out.data_ptr<float>(), total, H, W, OH, OW);
    return out;
}
"""

_ext = load_inline(name="deconv_sparse_ext_final_double", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=["conv_transpose2d_forward"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv_transpose2d.bias
        if b is None:
            b = torch.empty((0,), device=x.device, dtype=x.dtype)
        return _ext.conv_transpose2d_forward(x.contiguous(), self.conv_transpose2d.weight.contiguous(), b, self.stride, self.padding, self.dilation)