import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 2D transposed convolution operation with asymmetric input, asymmetric kernel, 
    grouped, padded, and dilated.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Size of the convolution kernel (height, width).
        stride (tuple, optional): Stride of the convolution (height, width). Defaults to (1, 1).
        padding (tuple, optional): Padding applied to the input (height, width). Defaults to (0, 0).
        dilation (tuple, optional): Spacing between kernel elements (height, width). Defaults to (1, 1).
        groups (int, optional): Number of blocked connections from input channels to output channels. 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: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), dilation: tuple = (1, 1), groups: 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, groups=groups, 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, width).

        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, 5)
height = 128
width = 256
stride = (2, 3)
padding = (1, 2)
dilation = (2, 1)
groups = 4

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

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

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, c10::optional<torch::Tensor> bias,
                                      int64_t stride_h, int64_t stride_w,
                                      int64_t pad_h, int64_t pad_w,
                                      int64_t dil_h, int64_t dil_w,
                                      int64_t groups);
"""

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

__global__ __launch_bounds__(256, 2) void deconv2d_w3_kernel(const float* __restrict__ x,
                                const float* __restrict__ w,
                                float* __restrict__ y,
                                int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int q = idx;
    int j = q & 255; q >>= 8;
    int orow = q & 127; q >>= 7;
    int co = q & 63; q >>= 6;
    int n = q;
    int g = co >> 4;
    int cog = co & 15;
    int ci0 = g << 3;
    int x_base0 = n * 1048576 + ci0 * 32768;
    int w_base0 = (ci0 * 16 + cog) * 15;
    float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f;
    bool has12 = (j < 255);

    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        int ih = orow + 1 - kh;
        if ((unsigned)ih < 128u) {
            int xb = x_base0 + ih * 256 + j;
            int wb = w_base0 + kh * 5;
            #pragma unroll
            for (int c = 0; c < 8; ++c) {
                float x0 = __ldg(x + xb + c * 32768);
                a0 = fmaf(x0, __ldg(w + wb + c * 240 + 2), a0);
                if (has12) {
                    float x1 = __ldg(x + xb + c * 32768 + 1);
                    a1 = fmaf(x1, __ldg(w + wb + c * 240 + 0), a1);
                    a1 = fmaf(x0, __ldg(w + wb + c * 240 + 3), a1);
                    a2 = fmaf(x1, __ldg(w + wb + c * 240 + 1), a2);
                    a2 = fmaf(x0, __ldg(w + wb + c * 240 + 4), a2);
                }
            }
        }
    }
    int out0 = (((n * 64 + co) * 257 + (orow * 2 + 1)) * 766 + j * 3);
    y[out0] = a0;
    if (has12) {
        y[out0 + 1] = a1;
        y[out0 + 2] = a2;
    }
}

__global__ void deconv2d_generic_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b, float* __restrict__ y,
    int N,int Cin,int H,int W,int Cout,int kH,int kW,int sh,int sw,int ph,int pw,int dh,int dw,int groups,int Hout,int Wout,int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total) return;
    int ow = idx % Wout; int t = idx / Wout; int oh = t % Hout; t /= Hout; int co = t % Cout; int n = t / Cout;
    int cin_g = Cin/groups, cout_g=Cout/groups, g=co/cout_g, cog=co-g*cout_g; float acc = b ? b[co] : 0.0f;
    for(int kh=0; kh<kH; ++kh){ int ihn=oh+ph-kh*dh; if(ihn%sh) continue; int ih=ihn/sh; if((unsigned)ih>=(unsigned)H) continue;
      for(int kw=0; kw<kW; ++kw){ int iwn=ow+pw-kw*dw; if(iwn%sw) continue; int iw=iwn/sw; if((unsigned)iw>=(unsigned)W) continue;
        for(int c=0;c<cin_g;++c) acc=fmaf(x[((n*Cin+g*cin_g+c)*H+ih)*W+iw], w[(((g*cin_g+c)*cout_g+cog)*kH+kh)*kW+kw], acc);
      }} y[idx]=acc;
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias,
                                      int64_t stride_h, int64_t stride_w, int64_t pad_h, int64_t pad_w,
                                      int64_t dil_h, int64_t dil_w, int64_t groups) {
    x = x.contiguous(); weight = weight.contiguous();
    int N=x.size(0), Cin=x.size(1), H=x.size(2), W=x.size(3), kH=weight.size(2), kW=weight.size(3), Cout=weight.size(1)*groups;
    int Hout=(H-1)*(int)stride_h-2*(int)pad_h+(int)dil_h*(kH-1)+1;
    int Wout=(W-1)*(int)stride_w-2*(int)pad_w+(int)dil_w*(kW-1)+1;
    auto y=torch::empty({N,Cout,Hout,Wout}, x.options());
    if (N==16 && Cin==32 && H==128 && W==256 && Cout==64 && kH==3 && kW==5 && stride_h==2 && stride_w==3 && pad_h==1 && pad_w==2 && dil_h==2 && dil_w==1 && groups==4 && !(bias.has_value() && bias.value().defined())) {
        cudaMemsetAsync(y.data_ptr<float>(), 0, (size_t)N*Cout*Hout*Wout*sizeof(float), at::cuda::getCurrentCUDAStream());
        int total = N * Cout * H * 256;
        int threads = 256, blocks = (total + threads - 1) / threads;
        deconv2d_w3_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>(), total);
    } else {
        const float* bp=nullptr; torch::Tensor bc; if(bias.has_value() && bias.value().defined()){bc=bias.value().contiguous(); bp=bc.data_ptr<float>();}
        int total=N*Cout*Hout*Wout; int threads=256, blocks=(total+threads-1)/threads;
        deconv2d_generic_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bp,y.data_ptr<float>(),N,Cin,H,W,Cout,kH,kW,(int)stride_h,(int)stride_w,(int)pad_h,(int)pad_w,(int)dil_h,(int)dil_w,(int)groups,Hout,Wout,total);
    }
    return y;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: tuple = (1, 1), 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, dilation=dilation, groups=groups, bias=bias)
        self.stride = stride; self.padding = padding; self.dilation = dilation; self.groups = groups
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.conv_transpose2d_forward(x, self.conv_transpose2d.weight, self.conv_transpose2d.bias, self.stride[0], self.stride[1], self.padding[0], self.padding[1], self.dilation[0], self.dilation[1], self.groups)