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_src = 'torch::Tensor convt2d_forward(torch::Tensor x, torch::Tensor w);'

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <mma.h>
using namespace nvcuda;

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

__global__ void convt_wmma_kernel(const float* __restrict__ x,
                                  const float* __restrict__ w,
                                  float* __restrict__ y,
                                  int Mtotal) {
    __shared__ float As[16 * 8];
    __shared__ float Bs[4 * 8 * 16];
    __shared__ float Cs[4 * 16 * 16];

    int warpId = threadIdx.x >> 5;
    int mbase = blockIdx.x * 16;

    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);

    for (int k0 = 0; k0 < 1344; k0 += 8) {
        int tid = threadIdx.x;
        if (tid < 128) {
            int row = tid >> 3;
            int kk = tid & 7;
            int m = mbase + row;
            float val = 0.0f;
            if (m < Mtotal) {
                int ow = m % 518;
                int tmp = m / 518;
                int oh = tmp % 514;
                int n = tmp / 514;
                int k = k0 + kk;
                int ic = k / 21;
                int rem = k - ic * 21;
                int kh = rem / 7;
                int kw = rem - kh * 7;
                int ih = oh - kh;
                int iw = ow - kw;
                if ((unsigned)ih < 512u && (unsigned)iw < 512u) {
                    val = tf32_round(x[((long)n * 64 + ic) * 512 * 512 + (long)ih * 512 + iw]);
                }
            }
            As[row * 8 + kk] = val;
        }

        for (int idx = threadIdx.x; idx < 512; idx += 128) {
            int tile = idx >> 7;
            int e = idx & 127;
            int kk = e >> 4;
            int col = e & 15;
            int k = k0 + kk;
            int ic = k / 21;
            int rem = k - ic * 21;
            int kh = rem / 7;
            int kw = rem - kh * 7;
            int oc = tile * 16 + col;
            Bs[idx] = tf32_round(w[((ic * 64 + oc) * 3 + kh) * 7 + kw]);
        }

        __syncthreads();
        wmma::load_matrix_sync(a_frag, As, 8);
        wmma::load_matrix_sync(b_frag, Bs + warpId * 128, 16);
        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
        __syncthreads();
    }

    wmma::store_matrix_sync(Cs + warpId * 256, c_frag, 16, wmma::mem_row_major);
    __syncthreads();

    for (int idx = threadIdx.x; idx < 1024; idx += 128) {
        int tile = idx >> 8;
        int e = idx & 255;
        int row = e >> 4;
        int col = e & 15;
        int m = mbase + row;
        if (m < Mtotal) {
            int ow = m % 518;
            int tmp = m / 518;
            int oh = tmp % 514;
            int n = tmp / 514;
            int oc = tile * 16 + col;
            y[((long)n * 64 + oc) * 514 * 518 + (long)oh * 518 + ow] = Cs[idx];
        }
    }
}

torch::Tensor convt2d_forward(torch::Tensor x, torch::Tensor w) {
    int N = x.size(0);
    int Mtotal = N * 514 * 518;
    auto y = torch::empty({N, 64, 514, 518}, x.options());
    convt_wmma_kernel<<<(Mtotal + 15) / 16, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), y.data_ptr<float>(), Mtotal);
    return y;
}
'''

_ext = load_inline(
    name='convt2d_wmma_final',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['convt2d_forward'],
    with_cuda=True,
    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(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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.convt2d_forward(x.contiguous(), self.conv_transpose2d.weight.contiguous())