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_src = r'''
#include <torch/extension.h>

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,
    int64_t sD, int64_t sH, int64_t sW,
    int64_t pD, int64_t pH, int64_t pW,
    int64_t opD, int64_t opH, int64_t opW,
    int64_t groups, bool has_bias);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define KDV 3
#define KHV 5
#define KWV 5
#define KVOL 75

__global__ __launch_bounds__(128, 2)
void ct3d_default_interior8_kernel(const float* __restrict__ x,
                                   const float* __restrict__ w,
                                   const float* __restrict__ b,
                                   float* __restrict__ y,
                                   int N, int Cin, int D, int H, int W,
                                   int Cout, int Dout, int Hout, int Wout,
                                   int Dint, int Hint, int Wint,
                                   int has_bias) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * (Cout >> 3) * Dint * Hint * Wint;
    if (tid >= total) return;

    int sp = (int)(tid % ((long long)Dint * Hint * Wint));
    int cb = (int)((tid / ((long long)Dint * Hint * Wint)) % (Cout >> 3));
    int n  = (int)(tid / (((long long)(Cout >> 3)) * Dint * Hint * Wint));

    int ow = sp % Wint + (KWV - 1);
    int t = sp / Wint;
    int oh = t % Hint + (KHV - 1);
    int od = t / Hint + (KDV - 1);
    int co = cb << 3;

    float a0 = has_bias ? __ldg(b + co + 0) : 0.0f;
    float a1 = has_bias ? __ldg(b + co + 1) : 0.0f;
    float a2 = has_bias ? __ldg(b + co + 2) : 0.0f;
    float a3 = has_bias ? __ldg(b + co + 3) : 0.0f;
    float a4 = has_bias ? __ldg(b + co + 4) : 0.0f;
    float a5 = has_bias ? __ldg(b + co + 5) : 0.0f;
    float a6 = has_bias ? __ldg(b + co + 6) : 0.0f;
    float a7 = has_bias ? __ldg(b + co + 7) : 0.0f;

    const long long xN = (long long)n * Cin * D * H * W;
    #pragma unroll 1
    for (int ci = 0; ci < Cin; ++ci) {
        const long long xC = xN + (long long)ci * D * H * W;
        const long long wC = (long long)ci * Cout * KVOL + (long long)co * KVOL;
        #pragma unroll
        for (int kd = 0; kd < KDV; ++kd) {
            int id = od - kd;
            #pragma unroll
            for (int kh = 0; kh < KHV; ++kh) {
                int ih = oh - kh;
                const float* xp = x + xC + ((long long)id * H + ih) * W + ow;
                const float* wp = w + wC + (kd * KHV + kh) * KWV;
                #pragma unroll
                for (int kw = 0; kw < KWV; ++kw) {
                    float xv = __ldg(xp - kw);
                    int kk = kw;
                    a0 = fmaf(xv, __ldg(wp + 0 * KVOL + kk), a0);
                    a1 = fmaf(xv, __ldg(wp + 1 * KVOL + kk), a1);
                    a2 = fmaf(xv, __ldg(wp + 2 * KVOL + kk), a2);
                    a3 = fmaf(xv, __ldg(wp + 3 * KVOL + kk), a3);
                    a4 = fmaf(xv, __ldg(wp + 4 * KVOL + kk), a4);
                    a5 = fmaf(xv, __ldg(wp + 5 * KVOL + kk), a5);
                    a6 = fmaf(xv, __ldg(wp + 6 * KVOL + kk), a6);
                    a7 = fmaf(xv, __ldg(wp + 7 * KVOL + kk), a7);
                }
            }
        }
    }

    long long ybase = ((long long)n * Cout * Dout * Hout * Wout) + ((long long)co * Dout + od) * Hout * Wout + (long long)oh * Wout + ow;
    long long cstride = (long long)Dout * Hout * Wout;
    y[ybase + 0 * cstride] = a0;
    y[ybase + 1 * cstride] = a1;
    y[ybase + 2 * cstride] = a2;
    y[ybase + 3 * cstride] = a3;
    y[ybase + 4 * cstride] = a4;
    y[ybase + 5 * cstride] = a5;
    y[ybase + 6 * cstride] = a6;
    y[ybase + 7 * cstride] = a7;
}

__global__ __launch_bounds__(128, 2)
void ct3d_default_edge8_kernel(const float* __restrict__ x,
                               const float* __restrict__ w,
                               const float* __restrict__ b,
                               float* __restrict__ y,
                               int N, int Cin, int D, int H, int W,
                               int Cout, int Dout, int Hout, int Wout,
                               int Dint, int Hint, int Wint,
                               int has_bias) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long spatial = (long long)Dout * Hout * Wout;
    long long total = (long long)N * (Cout >> 3) * spatial;
    if (tid >= total) return;

    int sp = (int)(tid % spatial);
    int cb = (int)((tid / spatial) % (Cout >> 3));
    int n  = (int)(tid / (((long long)(Cout >> 3)) * spatial));
    int ow = sp % Wout;
    int t = sp / Wout;
    int oh = t % Hout;
    int od = t / Hout;

    if (od >= KDV-1 && od < D && oh >= KHV-1 && oh < H && ow >= KWV-1 && ow < W) return;

    int co = cb << 3;
    float a0 = has_bias ? __ldg(b + co + 0) : 0.0f;
    float a1 = has_bias ? __ldg(b + co + 1) : 0.0f;
    float a2 = has_bias ? __ldg(b + co + 2) : 0.0f;
    float a3 = has_bias ? __ldg(b + co + 3) : 0.0f;
    float a4 = has_bias ? __ldg(b + co + 4) : 0.0f;
    float a5 = has_bias ? __ldg(b + co + 5) : 0.0f;
    float a6 = has_bias ? __ldg(b + co + 6) : 0.0f;
    float a7 = has_bias ? __ldg(b + co + 7) : 0.0f;

    const long long xN = (long long)n * Cin * D * H * W;
    for (int ci = 0; ci < Cin; ++ci) {
        const long long xC = xN + (long long)ci * D * H * W;
        const long long wC = (long long)ci * Cout * KVOL + (long long)co * KVOL;
        #pragma unroll
        for (int kd = 0; kd < KDV; ++kd) {
            int id = od - kd;
            if ((unsigned)id >= (unsigned)D) continue;
            #pragma unroll
            for (int kh = 0; kh < KHV; ++kh) {
                int ih = oh - kh;
                if ((unsigned)ih >= (unsigned)H) continue;
                const float* xp = x + xC + ((long long)id * H + ih) * W;
                const float* wp = w + wC + (kd * KHV + kh) * KWV;
                #pragma unroll
                for (int kw = 0; kw < KWV; ++kw) {
                    int iw = ow - kw;
                    if ((unsigned)iw >= (unsigned)W) continue;
                    float xv = __ldg(xp + iw);
                    a0 = fmaf(xv, __ldg(wp + 0 * KVOL + kw), a0);
                    a1 = fmaf(xv, __ldg(wp + 1 * KVOL + kw), a1);
                    a2 = fmaf(xv, __ldg(wp + 2 * KVOL + kw), a2);
                    a3 = fmaf(xv, __ldg(wp + 3 * KVOL + kw), a3);
                    a4 = fmaf(xv, __ldg(wp + 4 * KVOL + kw), a4);
                    a5 = fmaf(xv, __ldg(wp + 5 * KVOL + kw), a5);
                    a6 = fmaf(xv, __ldg(wp + 6 * KVOL + kw), a6);
                    a7 = fmaf(xv, __ldg(wp + 7 * KVOL + kw), a7);
                }
            }
        }
    }
    long long ybase = ((long long)n * Cout * Dout * Hout * Wout) + ((long long)co * Dout + od) * Hout * Wout + (long long)oh * Wout + ow;
    long long cstride = (long long)Dout * Hout * Wout;
    y[ybase + 0 * cstride] = a0;
    y[ybase + 1 * cstride] = a1;
    y[ybase + 2 * cstride] = a2;
    y[ybase + 3 * cstride] = a3;
    y[ybase + 4 * cstride] = a4;
    y[ybase + 5 * cstride] = a5;
    y[ybase + 6 * cstride] = a6;
    y[ybase + 7 * cstride] = a7;
}

__global__
void ct3d_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,
                         const float* __restrict__ b, float* __restrict__ y,
                         int N, int Cin, int D, int H, int W,
                         int Cout, int KD, int KH, int KW,
                         int Dout, int Hout, int Wout,
                         int sD, int sH, int sW, int pD, int pH, int pW,
                         int groups, int has_bias) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * Cout * Dout * Hout * Wout;
    if (tid >= total) return;
    int ow = tid % Wout;
    long long q = tid / Wout;
    int oh = q % Hout; q /= Hout;
    int od = q % Dout; q /= Dout;
    int co = q % Cout;
    int n = q / Cout;
    int cpg_in = Cin / groups;
    int cpg_out = Cout / groups;
    int g = co / cpg_out;
    int co_g = co - g * cpg_out;
    float acc = has_bias ? __ldg(b + co) : 0.0f;
    for (int ci0 = 0; ci0 < cpg_in; ++ci0) {
        int ci = g * cpg_in + ci0;
        for (int kd = 0; kd < KD; ++kd) {
            int td = od + pD - kd;
            if (td % sD) continue;
            int id = td / sD;
            if ((unsigned)id >= (unsigned)D) continue;
            for (int kh = 0; kh < KH; ++kh) {
                int th = oh + pH - kh;
                if (th % sH) continue;
                int ih = th / sH;
                if ((unsigned)ih >= (unsigned)H) continue;
                for (int kw = 0; kw < KW; ++kw) {
                    int tw = ow + pW - kw;
                    if (tw % sW) continue;
                    int iw = tw / sW;
                    if ((unsigned)iw >= (unsigned)W) continue;
                    float xv = __ldg(x + ((long long)n*Cin*D*H*W + (long long)ci*D*H*W + ((long long)id*H+ih)*W + iw));
                    float wv = __ldg(w + ((((long long)ci*cpg_out + co_g)*KD + kd)*KH + kh)*KW + kw);
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }
    y[tid] = acc;
}

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,
    int64_t sD64, int64_t sH64, int64_t sW64,
    int64_t pD64, int64_t pH64, int64_t pW64,
    int64_t opD64, int64_t opH64, int64_t opW64,
    int64_t groups64, bool has_bias_bool) {

    int N = (int)x.size(0), Cin = (int)x.size(1), D = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);
    int Cout_per_g = (int)w.size(1);
    int KD = (int)w.size(2), KH = (int)w.size(3), KW = (int)w.size(4);
    int groups = (int)groups64;
    int Cout = Cout_per_g * groups;
    int sD=(int)sD64, sH=(int)sH64, sW=(int)sW64, pD=(int)pD64, pH=(int)pH64, pW=(int)pW64;
    int Dout = (D - 1) * sD - 2 * pD + KD + (int)opD64;
    int Hout = (H - 1) * sH - 2 * pH + KH + (int)opH64;
    int Wout = (W - 1) * sW - 2 * pW + KW + (int)opW64;
    auto y = torch::empty({N, Cout, Dout, Hout, Wout}, x.options());
    const float* bp = has_bias_bool ? b.data_ptr<float>() : nullptr;
    int has_bias = has_bias_bool ? 1 : 0;
    int threads = 128;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat &&
        groups == 1 && sD == 1 && sH == 1 && sW == 1 && pD == 0 && pH == 0 && pW == 0 &&
        KD == 3 && KH == 5 && KW == 5 && (Cout % 8) == 0) {
        int Dint = D - KD + 1;
        int Hint = H - KH + 1;
        int Wint = W - KW + 1;
        if (Dint > 0 && Hint > 0 && Wint > 0) {
            long long total_int = (long long)N * (Cout >> 3) * Dint * Hint * Wint;
            ct3d_default_interior8_kernel<<<(unsigned)((total_int + threads - 1) / threads), threads, 0, stream>>>(
                x.data_ptr<float>(), w.data_ptr<float>(), bp, y.data_ptr<float>(),
                N, Cin, D, H, W, Cout, Dout, Hout, Wout, Dint, Hint, Wint, has_bias);
        }
        long long total_edge = (long long)N * (Cout >> 3) * Dout * Hout * Wout;
        ct3d_default_edge8_kernel<<<(unsigned)((total_edge + threads - 1) / threads), threads, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bp, y.data_ptr<float>(),
            N, Cin, D, H, W, Cout, Dout, Hout, Wout, Dint, Hint, Wint, has_bias);
    } else {
        long long total = (long long)N * Cout * Dout * Hout * Wout;
        ct3d_generic_kernel<<<(unsigned)((total + threads - 1) / threads), threads, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bp, y.data_ptr<float>(),
            N, Cin, D, H, W, Cout, KD, KH, KW, Dout, Hout, Wout,
            sD, sH, sW, pD, pH, pW, groups, has_bias);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='ct3d_ext_fast_default_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv_transpose3d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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)
        self.stride = stride
        self.padding = padding
        self.output_padding = output_padding
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bias = self.conv_transpose3d.bias
        has_bias = bias is not None
        if bias is None:
            bias = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv_transpose3d_forward(
            x, self.conv_transpose3d.weight, bias,
            int(self.stride[0]), int(self.stride[1]), int(self.stride[2]),
            int(self.padding[0]), int(self.padding[1]), int(self.padding[2]),
            int(self.output_padding[0]), int(self.output_padding[1]), int(self.output_padding[2]),
            int(self.groups), bool(has_bias)
        )