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, 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, bool has_bias);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>

__global__ void deconv2d_3x5_s2x3_opt_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    int N, int H, int W, int Ho, int Wo, bool has_bias)
{
    // One thread computes all 16 output channels belonging to one group for one (n,oh,ow).
    // Specialized for Cin=32, Cout=64, groups=4, k=(3,5), stride=(2,3), pad=(1,2), dil=(2,1).
    const int G = 4;
    const int CinPG = 8;
    const int CoutPG = 16;
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * G * Ho * Wo;
    if (tid >= total) return;

    int ow = tid % Wo;
    tid /= Wo;
    int oh = tid % Ho;
    tid /= Ho;
    int g = tid % G;
    int n = tid / G;

    const int Cout = 64;
    const int Cin = 32;
    const long long out_base = ((long long)n * Cout + g * CoutPG) * Ho * Wo + (long long)oh * Wo + ow;

    // With pad_h=1,dil_h=2,stride_h=2, only odd output rows receive convolution contributions.
    if ((oh & 1) == 0) {
        #pragma unroll
        for (int oc = 0; oc < 16; ++oc) {
            y[out_base + (long long)oc * Ho * Wo] = has_bias ? b[g * CoutPG + oc] : 0.0f;
        }
        return;
    }

    float acc0  = has_bias ? b[g*CoutPG +  0] : 0.0f;
    float acc1  = has_bias ? b[g*CoutPG +  1] : 0.0f;
    float acc2  = has_bias ? b[g*CoutPG +  2] : 0.0f;
    float acc3  = has_bias ? b[g*CoutPG +  3] : 0.0f;
    float acc4  = has_bias ? b[g*CoutPG +  4] : 0.0f;
    float acc5  = has_bias ? b[g*CoutPG +  5] : 0.0f;
    float acc6  = has_bias ? b[g*CoutPG +  6] : 0.0f;
    float acc7  = has_bias ? b[g*CoutPG +  7] : 0.0f;
    float acc8  = has_bias ? b[g*CoutPG +  8] : 0.0f;
    float acc9  = has_bias ? b[g*CoutPG +  9] : 0.0f;
    float acc10 = has_bias ? b[g*CoutPG + 10] : 0.0f;
    float acc11 = has_bias ? b[g*CoutPG + 11] : 0.0f;
    float acc12 = has_bias ? b[g*CoutPG + 12] : 0.0f;
    float acc13 = has_bias ? b[g*CoutPG + 13] : 0.0f;
    float acc14 = has_bias ? b[g*CoutPG + 14] : 0.0f;
    float acc15 = has_bias ? b[g*CoutPG + 15] : 0.0f;

    const int ih_base = (oh + 1) >> 1;
    const int r = (ow + 2) % 3;

    #pragma unroll
    for (int icl = 0; icl < 8; ++icl) {
        const int ic = g * CinPG + icl;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            const int ih = ih_base - kh;
            if ((unsigned)ih < (unsigned)H) {
                #pragma unroll
                for (int kk = 0; kk < 2; ++kk) {
                    int kw = r + kk * 3;
                    if (kw < 5) {
                        int iw = (ow + 2 - kw) / 3;
                        if ((unsigned)iw < (unsigned)W) {
                            float xv = x[((long long)n * Cin + ic) * H * W + (long long)ih * W + iw];
                            const float* wp = w + ((ic * CoutPG) * 3 + kh) * 5 + kw;
                            acc0  = fmaf(xv, wp[ 0 * 15], acc0);
                            acc1  = fmaf(xv, wp[ 1 * 15], acc1);
                            acc2  = fmaf(xv, wp[ 2 * 15], acc2);
                            acc3  = fmaf(xv, wp[ 3 * 15], acc3);
                            acc4  = fmaf(xv, wp[ 4 * 15], acc4);
                            acc5  = fmaf(xv, wp[ 5 * 15], acc5);
                            acc6  = fmaf(xv, wp[ 6 * 15], acc6);
                            acc7  = fmaf(xv, wp[ 7 * 15], acc7);
                            acc8  = fmaf(xv, wp[ 8 * 15], acc8);
                            acc9  = fmaf(xv, wp[ 9 * 15], acc9);
                            acc10 = fmaf(xv, wp[10 * 15], acc10);
                            acc11 = fmaf(xv, wp[11 * 15], acc11);
                            acc12 = fmaf(xv, wp[12 * 15], acc12);
                            acc13 = fmaf(xv, wp[13 * 15], acc13);
                            acc14 = fmaf(xv, wp[14 * 15], acc14);
                            acc15 = fmaf(xv, wp[15 * 15], acc15);
                        }
                    }
                }
            }
        }
    }

    const long long ch_stride = (long long)Ho * Wo;
    y[out_base +  0 * ch_stride] = acc0;
    y[out_base +  1 * ch_stride] = acc1;
    y[out_base +  2 * ch_stride] = acc2;
    y[out_base +  3 * ch_stride] = acc3;
    y[out_base +  4 * ch_stride] = acc4;
    y[out_base +  5 * ch_stride] = acc5;
    y[out_base +  6 * ch_stride] = acc6;
    y[out_base +  7 * ch_stride] = acc7;
    y[out_base +  8 * ch_stride] = acc8;
    y[out_base +  9 * ch_stride] = acc9;
    y[out_base + 10 * ch_stride] = acc10;
    y[out_base + 11 * ch_stride] = acc11;
    y[out_base + 12 * ch_stride] = acc12;
    y[out_base + 13 * ch_stride] = acc13;
    y[out_base + 14 * ch_stride] = acc14;
    y[out_base + 15 * ch_stride] = acc15;
}

__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 Ho, int Wo,
    int kH, int kW, int sh, int sw, int ph, int pw, int dh, int dw, int groups, bool has_bias)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * Cout * Ho * Wo;
    if (idx >= total) return;
    int ow = idx % Wo; idx /= Wo;
    int oh = idx % Ho; idx /= Ho;
    int oc = idx % Cout;
    int n = idx / Cout;
    int CoutPG = Cout / groups;
    int CinPG = Cin / groups;
    int g = oc / CoutPG;
    int ocl = oc - g * CoutPG;
    float acc = has_bias ? b[oc] : 0.0f;
    for (int icl = 0; icl < CinPG; ++icl) {
        int ic = g * CinPG + icl;
        for (int kh = 0; kh < kH; ++kh) {
            int th = oh + ph - kh * dh;
            if (th % sh) continue;
            int ih = th / sh;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < kW; ++kw) {
                int tw = ow + pw - kw * dw;
                if (tw % sw) continue;
                int iw = tw / sw;
                if ((unsigned)iw < (unsigned)W) {
                    float xv = x[((long long)n * Cin + ic) * H * W + (long long)ih * W + iw];
                    float wv = w[((long long)ic * CoutPG + ocl) * kH * kW + kh * kW + kw];
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }
    y[((long long)n * Cout + oc) * Ho * Wo + (long long)oh * Wo + ow] = acc;
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight, 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, bool has_bias) {
    const int N = (int)x.size(0);
    const int Cin = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int kH = (int)weight.size(2);
    const int kW = (int)weight.size(3);
    const int CoutPG = (int)weight.size(1);
    const int Cout = CoutPG * (int)groups;
    const int Ho = (H - 1) * (int)stride_h - 2 * (int)pad_h + (int)dil_h * (kH - 1) + 1;
    const int Wo = (W - 1) * (int)stride_w - 2 * (int)pad_w + (int)dil_w * (kW - 1) + 1;
    auto y = torch::empty({N, Cout, Ho, Wo}, x.options());

    const int threads = 128;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;

    if (Cin == 32 && Cout == 64 && (int)groups == 4 && kH == 3 && kW == 5 &&
        stride_h == 2 && stride_w == 3 && pad_h == 1 && pad_w == 2 && dil_h == 2 && dil_w == 1) {
        long long total = (long long)N * (int)groups * Ho * Wo;
        int blocks = (int)((total + threads - 1) / threads);
        deconv2d_3x5_s2x3_opt_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), N, H, W, Ho, Wo, has_bias);
    } else {
        long long total = (long long)N * Cout * Ho * Wo;
        int blocks = (int)((total + threads - 1) / threads);
        deconv2d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),
            N, Cin, H, W, Cout, Ho, Wo, kH, kW, (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w, (int)dil_h, (int)dil_w, (int)groups, has_bias);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

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

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:
        b = self.conv_transpose2d.bias
        has_bias = b is not None
        if b is None:
            b = x.new_empty((0,))
        return _ext.conv_transpose2d_forward(
            x, self.conv_transpose2d.weight, b,
            int(self.stride[0]), int(self.stride[1]),
            int(self.padding[0]), int(self.padding[1]),
            int(self.dilation[0]), int(self.dilation[1]),
            int(self.groups), bool(has_bias)
        )