import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        output_padding (int, 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: int, 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  # double channels for heavier compute
out_channels = 64
kernel_size = 3
# larger square input
height = 1024
width = 1024

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_src = r'''
#include <torch/extension.h>
torch::Tensor deconv2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,int stride,int padding,int output_padding,int groups);
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <mma.h>
using namespace nvcuda;
#define WARPS 4
__device__ __forceinline__ float to_tf32(float x) { unsigned y; asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(y) : "f"(x)); return __uint_as_float(y); }
__global__ void deconv_wmma_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b,
                                   float* __restrict__ y, int N, int IC, int H, int W, int OC, int OH, int OW, int has_bias) {
    __shared__ float As[WARPS][16*8];
    __shared__ float Bs[8*16];
    __shared__ float Cs[WARPS][16*16];
    int tid=threadIdx.x, warp=tid>>5, lane=tid&31;
    int tiles_w = (OW + 15) >> 4;
    int tile = blockIdx.x * WARPS + warp;
    int out_row = tile / tiles_w;
    int out_col0 = (tile - out_row * tiles_w) << 4;
    int oc0 = blockIdx.y << 4;
    int n_batch=blockIdx.z;
    bool valid_tile = out_row < OH;
    wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> af;
    wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::row_major> bf;
    wmma::fragment<wmma::accumulator,16,16,8,float> cf;
    wmma::fill_fragment(cf,0.0f);
    for(int k0=0;k0<IC*9;k0+=8){
        for(int idx=lane; idx<16*8; idx+=32){
            int mi=idx>>3; int kk=idx&7; int gk=k0+kk; int ow=out_col0+mi; float v=0.0f;
            if(valid_tile && ow<OW && gk<IC*9){
                int ic=gk/9; int rem=gk-ic*9; int r=rem/3; int s=rem-r*3; int ih=out_row-r; int iw=ow-s;
                if((unsigned)ih<(unsigned)H && (unsigned)iw<(unsigned)W) v=x[((n_batch*IC+ic)*H+ih)*W+iw];
            }
            As[warp][idx]=to_tf32(v);
        }
        for(int idx=tid; idx<8*16; idx+=blockDim.x){
            int kk=idx>>4; int cn=idx&15; int gk=k0+kk; int oc=oc0+cn; float v=0.0f;
            if(gk<IC*9 && oc<OC){int ic=gk/9; int rem=gk-ic*9; v=w[((ic*OC+oc)*3+(rem/3))*3+(rem%3)];}
            Bs[idx]=to_tf32(v);
        }
        __syncthreads();
        wmma::load_matrix_sync(af,&As[warp][0],8);
        wmma::load_matrix_sync(bf,Bs,16);
        wmma::mma_sync(cf,af,bf,cf);
        __syncthreads();
    }
    wmma::store_matrix_sync(&Cs[warp][0],cf,16,wmma::mem_row_major);
    __syncthreads();
    for(int idx=tid; idx<WARPS*16*16; idx+=blockDim.x){
        int wp=idx>>8; int rem=idx&255; int mi=rem>>4; int cn=rem&15;
        int t=blockIdx.x*WARPS+wp; int row=t/tiles_w; int col=((t-row*tiles_w)<<4)+mi; int oc=oc0+cn;
        if(row<OH && col<OW && oc<OC){float val=Cs[wp][rem]; if(has_bias) val+=b[oc]; y[((n_batch*OC+oc)*OH+row)*OW+col]=val;}
    }
}
torch::Tensor deconv2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,int stride,int padding,int output_padding,int groups){
 int64_t N=x.size(0),IC=x.size(1),H=x.size(2),W=x.size(3),K=weight.size(2),OC=weight.size(1)*groups; int64_t OH=(H-1)*stride-2*padding+K+output_padding, OW=(W-1)*stride-2*padding+K+output_padding; auto y=torch::empty({N,OC,OH,OW},x.options()); dim3 block(32*WARPS); int tiles_w=((int)OW+15)/16; int tile_count=tiles_w*(int)OH; dim3 grid((tile_count+WARPS-1)/WARPS,(OC+15)/16,N); int has_bias=bias.numel()!=0; deconv_wmma_kernel<<<grid,block>>>(x.data_ptr<float>(),weight.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,y.data_ptr<float>(),(int)N,(int)IC,(int)H,(int)W,(int)OC,(int)OH,(int)OW,has_bias); return y;}
'''
_ext=load_inline(name='deconv2d_ext_final',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['deconv2d_forward'],extra_cuda_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
 def __init__(self,in_channels:int,out_channels:int,kernel_size:int,stride:int=1,padding:int=0,output_padding:int=0,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,groups=groups,bias=bias); self.stride=stride; self.padding=padding; self.output_padding=output_padding; self.groups=groups
 def forward(self,x:torch.Tensor)->torch.Tensor:
  b=self.conv_transpose2d.bias
  if b is None: b=torch.empty(0,device=x.device,dtype=x.dtype)
  return _ext.deconv2d_forward(x.contiguous(),self.conv_transpose2d.weight.contiguous(),b,self.stride,self.padding,self.output_padding,self.groups)