import torch
import torch.nn as nn

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

    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).
        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), bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, 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 = 8
in_channels = 32
out_channels = 32
kernel_size = (3, 7)
height = 512
width = 1024
stride = (1, 1)
padding = (1, 3)

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]

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

_cpp_src = """
#include <torch/extension.h>
torch::Tensor convt2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias);
"""

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

__global__ void k_oc32(const float* __restrict__ x,
                       const float* __restrict__ wt,
                       const float* __restrict__ bias,
                       float* __restrict__ y,
                       int N, int H, int W, int has_bias) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * H * W;
    if (idx >= total) return;

    int ow = idx % W;
    long long t = idx / W;
    int oh = t % H;
    int n = t / H;

    float acc[32];
    #pragma unroll
    for (int oc = 0; oc < 32; ++oc) acc[oc] = has_bias ? bias[oc] : 0.0f;

    // Match PyTorch/cuDNN accumulation order: kh, kw, then input channel.
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        int ih = oh + 1 - kh;
        if ((unsigned)ih < (unsigned)H) {
            #pragma unroll
            for (int kw = 0; kw < 7; ++kw) {
                int iw = ow + 3 - kw;
                if ((unsigned)iw < (unsigned)W) {
                    #pragma unroll
                    for (int ic = 0; ic < 32; ++ic) {
                        float xv = x[((long long)n * 32 + ic) * H * W + (long long)ih * W + iw];
                        const float* wb = wt + ((long long)ic * 32) * 21 + kh * 7 + kw;
                        #pragma unroll
                        for (int oc = 0; oc < 32; ++oc) {
                            acc[oc] = fmaf(xv, wb[oc * 21], acc[oc]);
                        }
                    }
                }
            }
        }
    }

    long long base = ((long long)n * 32) * H * W + (long long)oh * W + ow;
    #pragma unroll
    for (int oc = 0; oc < 32; ++oc) y[base + (long long)oc * H * W] = acc[oc];
}

__global__ void k_generic(const float* __restrict__ x,
                          const float* __restrict__ wt,
                          const float* __restrict__ bias,
                          float* __restrict__ y,
                          int N, int IC, int H, int W, int OC, int KH, int KW,
                          int has_bias) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * OC * H * W;
    if (idx >= total) return;
    int ow = idx % W;
    long long t = idx / W;
    int oh = t % H;
    t /= H;
    int oc = t % OC;
    int n = t / OC;
    float acc = has_bias ? bias[oc] : 0.0f;

    for (int kh = 0; kh < KH; ++kh) {
        int ih = oh + KH / 2 - kh;
        if ((unsigned)ih < (unsigned)H) {
            for (int kw = 0; kw < KW; ++kw) {
                int iw = ow + KW / 2 - kw;
                if ((unsigned)iw < (unsigned)W) {
                    for (int ic = 0; ic < IC; ++ic) {
                        acc = fmaf(x[((long long)n * IC + ic) * H * W + (long long)ih * W + iw],
                                   wt[((long long)ic * OC + oc) * KH * KW + kh * KW + kw],
                                   acc);
                    }
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor convt2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    int N = (int)x.size(0);
    int IC = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int OC = (int)w.size(1);
    int KH = (int)w.size(2);
    int KW = (int)w.size(3);
    auto y = torch::empty({N, OC, H, W}, x.options());
    int threads = 256;
    int has_bias = (bias.numel() != 0);
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;

    if (IC == 32 && OC == 32 && KH == 3 && KW == 7) {
        long long total = (long long)N * H * W;
        int blocks = (int)((total + threads - 1) / threads);
        k_oc32<<<blocks, threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), bptr,
                                    y.data_ptr<float>(), N, H, W, has_bias);
    } else {
        long long total = (long long)N * OC * H * W;
        int blocks = (int)((total + threads - 1) / threads);
        k_generic<<<blocks, threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), bptr,
                                       y.data_ptr<float>(), N, IC, H, W, OC, KH, KW, has_bias);
    }
    return y;
}
"""

_ext = load_inline(
    name="convt2d_ext_final_oc32",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["convt2d_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), bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias
        )
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv_transpose2d.bias
        if b is None:
            b = self._empty_bias
        return _ext.convt2d_forward(x, self.conv_transpose2d.weight, b)