import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 3D convolution operation with asymmetric input and a 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 or tuple, 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 each dimension in the output shape. 
                                                  Defaults to 0.
        dilation (int or tuple, optional): Spacing between kernel elements. Defaults to 1.
        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, 
                 dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(in_channels, out_channels, (kernel_size, kernel_size, kernel_size), 
                                                stride=stride, padding=padding, output_padding=output_padding, 
                                                dilation=dilation, 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, height, width).

        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 = 8
in_channels = 48
out_channels = 24
kernel_size = 3
depth = 96
height = 96
width = 96

def get_inputs():
    x = torch.rand(batch_size, in_channels, depth, 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 conv_transpose3d_forward_cuda(torch::Tensor x, torch::Tensor weight);

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight) {
    return conv_transpose3d_forward_cuda(x, weight);
}
'''

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

#define CI 48
#define CO 24
#define DI 96
#define HI 96
#define WI 96
#define DO 98
#define HO 98
#define WO 98
#define XS (DI*HI*WI)

__device__ __forceinline__ float tf32_round(float x) {
    unsigned y;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(y) : "f"(x));
    return __uint_as_float(y);
}

__global__ void round_kernel(const float* __restrict__ a, float* __restrict__ b, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) b[i] = tf32_round(a[i]);
}

__global__ __launch_bounds__(256, 2)
void deconv3d_k3_kernel(const float* __restrict__ x,
                        const float* __restrict__ w,
                        float* __restrict__ y,
                        int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int t = idx;
    const int ow = t % WO; t /= WO;
    const int oh = t % HO; t /= HO;
    const int od = t % DO; t /= DO;
    const int oc = t % CO; t /= CO;
    const int n = t;

    float sum = 0.0f;
    const int nbase = n * CI * XS;

    int kd0 = od - (DI - 1); if (kd0 < 0) kd0 = 0;
    int kd1 = od;            if (kd1 > 2) kd1 = 2;
    int kh0 = oh - (HI - 1); if (kh0 < 0) kh0 = 0;
    int kh1 = oh;            if (kh1 > 2) kh1 = 2;
    int kw0 = ow - (WI - 1); if (kw0 < 0) kw0 = 0;
    int kw1 = ow;            if (kw1 > 2) kw1 = 2;

    for (int ic = 0; ic < CI; ++ic) {
        const float* __restrict__ xp = x + nbase + ic * XS;
        const float* __restrict__ wp = w + (ic * CO + oc) * 27;
        for (int kd = kd0; kd <= kd1; ++kd) {
            int zoff = (od - kd) * (HI * WI);
            for (int kh = kh0; kh <= kh1; ++kh) {
                int yoff = (oh - kh) * WI;
                for (int kw = kw0; kw <= kw1; ++kw) {
                    int wk = (kd * 3 + kh) * 3 + kw;
                    sum = fmaf(xp[zoff + yoff + (ow - kw)], wp[wk], sum);
                }
            }
        }
    }
    y[idx] = sum;
}

torch::Tensor conv_transpose3d_forward_cuda(torch::Tensor x, torch::Tensor weight) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
    auto xr = torch::empty_like(x);
    auto wr = torch::empty_like(weight);
    int xn = (int)x.numel();
    int wn = (int)weight.numel();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    round_kernel<<<(xn + 255) / 256, 256, 0, stream>>>(x.data_ptr<float>(), xr.data_ptr<float>(), xn);
    round_kernel<<<(wn + 255) / 256, 256, 0, stream>>>(weight.data_ptr<float>(), wr.data_ptr<float>(), wn);

    auto y = torch::empty({x.size(0), (long)CO, (long)DO, (long)HO, (long)WO}, x.options());
    int total = (int)(x.size(0) * CO * DO * HO * WO);
    deconv3d_k3_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        xr.data_ptr<float>(), wr.data_ptr<float>(), y.data_ptr<float>(), total);
    return y;
}
'''

_ext = load_inline(
    name='deconv3d_k3_final_ext',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['conv_transpose3d_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
)

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,
                 dilation: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels, (kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, output_padding=output_padding,
            dilation=dilation, groups=groups, bias=bias)

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