import torch
import torch.nn as nn

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

    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 3 integers representing the kernel size in the form (depth, height, width).
        stride (tuple, optional): Tuple of 3 integers representing the stride in the form (depth, height, width). Defaults to (1, 1, 1).
        padding (tuple, optional): Tuple of 3 integers representing the padding in the form (depth, height, width). Defaults to (0, 0, 0).
        output_padding (tuple, optional): Tuple of 3 integers representing the output padding in the form (depth, height, width). Defaults to (0, 0, 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: tuple = (1, 1, 1), padding: tuple = (0, 0, 0), output_padding: tuple = (0, 0, 0), groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(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 3D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth_in, height_in, width_in).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, depth_out, height_out, width_out).
        """
        return self.conv_transpose3d(x)

# Test code
batch_size = 16
in_channels = 32
out_channels = 16
kernel_size = (3, 5, 7)  # Asymmetric kernel size
depth_in = 16
height_in = 32
width_in = 64

def get_inputs():
    x = torch.rand(batch_size, in_channels, depth_in, 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 = r"""
#include <torch/extension.h>
torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <mma.h>
using namespace nvcuda;

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

template<int WARPS>
__global__ void deconv3d_wmma_pad_kernel(const float* __restrict__ x,
                                         const float* __restrict__ w,
                                         float* __restrict__ y,
                                         int N) {
    constexpr int Cin=32, Cout=16, Din=16, Hin=32, Win=64;
    constexpr int kD=3, kH=5, kW=7;
    constexpr int Dout=18, Hout=36, Wout=70, S=45360;

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int tile = blockIdx.x * WARPS + warp;
    int pos0 = tile * 16;
    int total_pos = N * S;

    extern __shared__ char shc[];
    float* base = (float*)shc;
    float* As = base + warp * (128 + 128 + 256);
    float* Bs = As + 128;
    float* Cs = Bs + 128;

    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 ci = 0; ci < 32; ++ci) {
        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            #pragma unroll
            for (int kh = 0; kh < 5; ++kh) {
                for (int i = lane; i < 128; i += 32) {
                    int co = i >> 3;
                    int q = i & 7;
                    float v = 0.0f;
                    if (q < 7) v = w[(((ci * Cout + co) * kD + kd) * kH + kh) * kW + q];
                    As[i] = tf32_round_dev(v);
                }
                for (int i = lane; i < 128; i += 32) {
                    int q = i >> 4;
                    int col = i & 15;
                    int p = pos0 + col;
                    float v = 0.0f;
                    if (q < 7 && p < total_pos) {
                        int n = p / S;
                        int s = p - n * S;
                        int ow = s % Wout;
                        s /= Wout;
                        int oh = s % Hout;
                        int od = s / Hout;
                        int id = od - kd;
                        int ih = oh - kh;
                        int iw = ow - q;
                        if ((unsigned)id < 16u && (unsigned)ih < 32u && (unsigned)iw < 64u) {
                            v = x[(((long long)n * Cin + ci) * Din * Hin + id * Hin + ih) * Win + iw];
                        }
                    }
                    Bs[i] = tf32_round_dev(v);
                }
                __syncwarp();
                wmma::load_matrix_sync(af, As, 8);
                wmma::load_matrix_sync(bf, Bs, 16);
                wmma::mma_sync(cf, af, bf, cf);
                __syncwarp();
            }
        }
    }

    wmma::store_matrix_sync(Cs, cf, 16, wmma::mem_row_major);
    __syncwarp();
    for (int i = lane; i < 256; i += 32) {
        int co = i >> 4;
        int col = i & 15;
        int p = pos0 + col;
        if (p < total_pos) {
            int n = p / S;
            int s = p - n * S;
            y[((long long)n * Cout + co) * S + s] = Cs[i];
        }
    }
}

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight) {
    int N = x.size(0);
    auto y = torch::empty({N, 16, 18, 36, 70}, x.options());
    int tiles = (N * 45360 + 15) / 16;
    const int WARPS = 8;
    int blocks = (tiles + WARPS - 1) / WARPS;
    size_t smem = WARPS * (128 + 128 + 256) * sizeof(float);
    deconv3d_wmma_pad_kernel<WARPS><<<blocks, WARPS * 32, smem>>>(x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>(), N);
    return y;
}
"""

_ext = load_inline(
    name="deconv3d_wmma_pad_final",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["conv_transpose3d_forward"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: tuple = (1, 1, 1), padding: tuple = (0, 0, 0),
                 output_padding: tuple = (0, 0, 0), groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(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_transpose3d_forward(x.contiguous(), self.conv_transpose3d.weight.contiguous())