import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 3D 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 (kernel_depth, kernel_width, kernel_height), 
                             where kernel_width == kernel_height.
        stride (tuple, optional): Stride of the convolution. Defaults to (1, 1, 1).
        padding (tuple, optional): Padding applied to the input. Defaults to (0, 0, 0).
        output_padding (tuple, optional): Additional size added to one side of the output shape. 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, width, height).

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

# Test code
batch_size = 16
in_channels = 32
out_channels = 64
kernel_depth = 3
kernel_width = 5
kernel_height = 5
depth = 64
width = 64
height = 64

def get_inputs():
    x = torch.rand(batch_size, in_channels, depth, width, height)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, (kernel_depth, kernel_width, kernel_height)]  # 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 = '''#include <torch/extension.h>\ntorch::Tensor fwd(torch::Tensor x, torch::Tensor w);\n'''
cu = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void k_co8(const float* __restrict__ x, const float* __restrict__ w, float* __restrict__ y,
                      long total, int N, int Din, int Hin, int Win, int Dout, int Hout, int Wout){
    const int Cin=32, Cout=64, kD=3, kH=5, kW=5;
    long Sper=(long)Dout*Hout*Wout;
    for(long idx=blockIdx.x*(long)blockDim.x+threadIdx.x; idx<total; idx+=(long)blockDim.x*gridDim.x){
        long t=idx;
        long loc=t%Sper; t/=Sper;
        int cg=(int)(t&7); int n=(int)(t>>3); int co0=cg*8;
        int ow=loc%Wout; long q=loc/Wout;
        int oh=q%Hout; int od=q/Hout;
        float a0=0.f,a1=0.f,a2=0.f,a3=0.f,a4=0.f,a5=0.f,a6=0.f,a7=0.f;
        int kd0 = od>63 ? od-63 : 0; int kd1 = od<2 ? od : 2;
        int kh0 = oh>63 ? oh-63 : 0; int kh1 = oh<4 ? oh : 4;
        int kw0 = ow>63 ? ow-63 : 0; int kw1 = ow<4 ? ow : 4;
        for(int ci=0; ci<Cin; ++ci){
            for(int kd=kd0; kd<=kd1; ++kd){
                int id=od-kd;
                for(int kh=kh0; kh<=kh1; ++kh){
                    int ih=oh-kh;
                    long xb=(((long)n*Cin+ci)*Din+id)*Hin*Win + (long)ih*Win;
                    int wb0=((ci*Cout+co0)*kD+kd)*kH*kW + kh*kW;
                    for(int kw=kw0; kw<=kw1; ++kw){
                        float xv=__ldg(x+xb+(ow-kw));
                        int wb=wb0+kw;
                        a0=fmaf(xv,__ldg(w+wb+0*75),a0);
                        a1=fmaf(xv,__ldg(w+wb+1*75),a1);
                        a2=fmaf(xv,__ldg(w+wb+2*75),a2);
                        a3=fmaf(xv,__ldg(w+wb+3*75),a3);
                        a4=fmaf(xv,__ldg(w+wb+4*75),a4);
                        a5=fmaf(xv,__ldg(w+wb+5*75),a5);
                        a6=fmaf(xv,__ldg(w+wb+6*75),a6);
                        a7=fmaf(xv,__ldg(w+wb+7*75),a7);
                    }
                }
            }
        }
        long yo=(long)n*Cout*Sper + co0*Sper + loc;
        y[yo+0*Sper]=a0; y[yo+1*Sper]=a1; y[yo+2*Sper]=a2; y[yo+3*Sper]=a3;
        y[yo+4*Sper]=a4; y[yo+5*Sper]=a5; y[yo+6*Sper]=a6; y[yo+7*Sper]=a7;
    }
}

torch::Tensor fwd(torch::Tensor x, torch::Tensor w){
    int N=x.size(0), Din=x.size(2), Hin=x.size(3), Win=x.size(4);
    int Dout=Din+2, Hout=Hin+4, Wout=Win+4;
    auto y=torch::empty({N,64,Dout,Hout,Wout}, x.options());
    long total=(long)N*8*Dout*Hout*Wout;
    int th=256; int bl=(int)min((long)65535,(total+th-1)/th);
    k_co8<<<bl,th,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),y.data_ptr<float>(),total,N,Din,Hin,Win,Dout,Hout,Wout);
    return y;
}
'''
ext=load_inline(name='ct3d_final_co8_submit', cpp_sources=cpp, cuda_sources=cu, functions=['fwd'], 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,1),padding:tuple=(0,0,0),output_padding:tuple=(0,0,0),groups:int=1,bias:bool=False):
        super().__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):
        return ext.fwd(x.contiguous(), self.conv_transpose3d.weight.contiguous())