import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 2D convolution with a square input and an asymmetric kernel.

    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 (int, optional): Stride of the convolution. Defaults to 1.
        padding (int or tuple, optional): Padding applied to the input. Defaults to 0.
        output_padding (int or tuple, optional): Additional size added to one side of the output shape. Defaults to 0.
        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: int = 1, padding: int = 0, output_padding: int = 0, 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, output_padding=output_padding, groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 2D 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 = 64
out_channels = 64
kernel_size = (3, 7)  # larger asymmetric kernel
width = 512
height = 512

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]  # Provide in_channels, out_channels, kernel_size for initialization

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_source='''#include <torch/extension.h>\ntorch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight);'''
cuda_source=r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float tf32_round(float a) {
    unsigned int x = __float_as_uint(a);
    unsigned int r = ((x >> 13) & 1u) + 0x1000u;
    x = (x + r) & 0xffffe000u;
    return __uint_as_float(x);
}

__global__ void round_x_kernel(const float* __restrict__ x, float* __restrict__ xr, int total) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) xr[i] = tf32_round(x[i]);
}

__global__ void pack_weight_kernel(const float* __restrict__ w, float* __restrict__ wp) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= 21 * 64 * 64) return;
    int oc = idx & 63;
    int tmp = idx >> 6;
    int ic = tmp & 63;
    int t = tmp >> 6;
    int kh = t / 7;
    int kw = t - kh * 7;
    wp[idx] = tf32_round(w[((ic * 64 + oc) * 3 + kh) * 7 + kw]);
}

__global__ void conv_transpose_kernel(const float* __restrict__ x,
                                      const float* __restrict__ wp,
                                      float* __restrict__ y,
                                      int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int oc = idx & 63;
    int p = idx >> 6;
    int ow = p % 518;
    p /= 518;
    int oh = p % 514;
    int n = p / 514;

    int kh0 = max(0, oh - 511);
    int kh1 = min(2, oh);
    int kw0 = max(0, ow - 511);
    int kw1 = min(6, ow);

    float acc = 0.0f;
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        if (kh < kh0 || kh > kh1) continue;
        int ih = oh - kh;
        #pragma unroll
        for (int kw = 0; kw < 7; ++kw) {
            if (kw < kw0 || kw > kw1) continue;
            int iw = ow - kw;
            const float* __restrict__ xp = x + ((n * 64) * 512 + ih) * 512 + iw;
            const float* __restrict__ wp2 = wp + ((kh * 7 + kw) * 64) * 64 + oc;
            #pragma unroll
            for (int ic = 0; ic < 64; ++ic) {
                acc = fmaf(xp[ic * 512 * 512], wp2[ic * 64], acc);
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight) {
    auto opts = x.options();
    auto xr = torch::empty_like(x);
    int xtotal = 8 * 64 * 512 * 512;
    round_x_kernel<<<(xtotal + 255) / 256, 256>>>(x.data_ptr<float>(), xr.data_ptr<float>(), xtotal);

    auto wp = torch::empty({21, 64, 64}, opts);
    pack_weight_kernel<<<(21 * 64 * 64 + 255) / 256, 256>>>(weight.data_ptr<float>(), wp.data_ptr<float>());

    auto y = torch::empty_strided({8, 64, 514, 518}, {514 * 518 * 64, 1, 518 * 64, 64}, opts);
    int total = 8 * 514 * 518 * 64;
    conv_transpose_kernel<<<(total + 255) / 256, 256>>>(xr.data_ptr<float>(), wp.data_ptr<float>(), y.data_ptr<float>(), total);
    return y;
}
'''
_ext=load_inline(name='conv_transpose2d_3x7_final_ext',cpp_sources=cpp_source,cuda_sources=cuda_source,functions=['conv_transpose2d_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
    def __init__(self,in_channels:int,out_channels:int,kernel_size:tuple,stride:int=1,padding:int=0,output_padding:int=0,groups:int=1,bias:bool=False):
        super().__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:
        return _ext.conv_transpose2d_forward(x,self.conv_transpose2d.weight)