import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Tuple of integers representing the kernel size (height, width).
        stride (tuple, optional): Tuple of integers representing the stride of the convolution. Defaults to (1, 1).
        padding (tuple, optional): Tuple of integers representing the padding applied to the input. Defaults to (0, 0).
        output_padding (tuple, optional): Tuple of integers representing the additional size added to one side of the output shape. Defaults to (0, 0).
        dilation (tuple, optional): Tuple of integers representing the spacing between kernel elements. 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), output_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, output_padding=output_padding, dilation=dilation, 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_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 = 64
in_channels = 64
out_channels = 128
kernel_size = (3, 5)
height_in = 128
width_in = 256

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]  # 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_src='torch::Tensor convt2d_forward(torch::Tensor x, torch::Tensor weight);'
cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <mma.h>
using namespace nvcuda;

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

template<int WARPS>
__global__ __launch_bounds__(WARPS*32)
void wmma_convt_kernel(const float* __restrict__ x,const float* __restrict__ w,float* __restrict__ y,
                       int B,int IC,int OC,int H,int W,int OH,int OW,int Mtiles){
    extern __shared__ float sh[];
    int tid=threadIdx.x; int lane=tid & 31; int warp=tid >> 5;
    float* Ash = sh + warp*128;
    float* Bsh = sh + WARPS*128 + warp*128;
    float* Csh = sh + WARPS*256 + warp*256;
    int tile_m = blockIdx.x * WARPS + warp;
    int tile_n = blockIdx.y;
    if(tile_m >= Mtiles) return;
    const int S = OH*OW;
    const int M = B*S;
    const int oc_base = tile_n*16;
    const int Ktot = IC*15;

    wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator,16,16,8,float> c_frag;
    wmma::fill_fragment(c_frag, 0.0f);

    int m_base = tile_m*16;
    for(int k0=0;k0<Ktot;k0+=8){
        for(int e=lane;e<128;e+=32){
            int r=e>>3; int kk=e&7; int kidx=k0+kk;
            float val=0.0f;
            int m=m_base+r;
            if(m < M && kidx < Ktot){
                int ic = kidx / 15; int rem = kidx - ic*15; int kh = rem / 5; int kw = rem - kh*5;
                int n = m / S; int sp = m - n*S; int oh = sp / OW; int ow = sp - oh*OW;
                int ih = oh - kh; int iw = ow - kw;
                if((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W){
                    val = tf32r(x[((long long)n*IC + ic)*H*W + ih*W + iw]);
                }
            }
            Ash[e]=val;
        }
        for(int e=lane;e<128;e+=32){
            int kk=e>>4; int c=e&15; int kidx=k0+kk; int oc=oc_base+c;
            float val=0.0f;
            if(kidx < Ktot && oc < OC){
                int ic=kidx/15; int rem=kidx-ic*15; int kh=rem/5; int kw=rem-kh*5;
                val = tf32r(w[((long long)ic*OC + oc)*15 + kh*5 + kw]);
            }
            Bsh[e]=val;
        }
        __syncwarp();
        wmma::load_matrix_sync(a_frag, Ash, 8);
        wmma::load_matrix_sync(b_frag, Bsh, 16);
        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
        __syncwarp();
    }
    wmma::store_matrix_sync(Csh, c_frag, 16, wmma::mem_row_major);
    __syncwarp();
    for(int e=lane;e<256;e+=32){
        int r=e>>4; int c=e&15; int m=m_base+r; int oc=oc_base+c;
        if(m < M && oc < OC){
            int n=m/S; int sp=m-n*S;
            y[((long long)n*OC + oc)*S + sp] = Csh[e];
        }
    }
}

torch::Tensor convt2d_forward(torch::Tensor x,torch::Tensor weight){
    int B=x.size(0),IC=x.size(1),H=x.size(2),W=x.size(3),OC=weight.size(1),OH=H+2,OW=W+4;
    auto y=torch::empty({B,OC,OH,OW},x.options());
    const int WARPS=4; int S=OH*OW; int M=B*S; int Mtiles=(M+15)/16; dim3 grid((Mtiles+WARPS-1)/WARPS, (OC+15)/16);
    size_t shmem = WARPS*(128+128+256)*sizeof(float);
    wmma_convt_kernel<WARPS><<<grid,WARPS*32,shmem,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),weight.data_ptr<float>(),y.data_ptr<float>(),B,IC,OC,H,W,OH,OW,Mtiles);
    return y;
}
'''
_ext=load_inline(name='convt2d_ext_wmma_final',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['convt2d_forward'],verbose=False,extra_cuda_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),output_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,output_padding=output_padding,dilation=dilation,groups=groups,bias=bias)
 def forward(self,x:torch.Tensor)->torch.Tensor:
  return _ext.convt2d_forward(x.contiguous(),self.conv_transpose2d.weight.contiguous())