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(torch::Tensor x,
                                      torch::Tensor weight,
                                      torch::Tensor bias,
                                      int64_t stride,
                                      int64_t padding,
                                      int64_t output_padding,
                                      int64_t dilation,
                                      int64_t groups);
'''

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")

__global__ __launch_bounds__(256, 2)
void deconv3d_k3s1p0_48_24_oc4_kernel(const float* __restrict__ x,
                                       const float* __restrict__ w,
                                       const float* __restrict__ b,
                                       float* __restrict__ y,
                                       int N, int D, int H, int W,
                                       int has_bias) {
    constexpr int IC = 48;
    constexpr int OC = 24;
    constexpr int K = 3;
    const int Do = D + 2;
    const int Ho = H + 2;
    const int Wo = W + 2;
    const int spatial_total = N * Do * Ho * Wo;
    int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= spatial_total) return;
    const int oc0 = blockIdx.y * 4;

    int t = p;
    const int ow = t % Wo; t /= Wo;
    const int oh = t % Ho; t /= Ho;
    const int od = t % Do; t /= Do;
    const int n = t;

    float a0 = has_bias ? __ldg(b + oc0 + 0) : 0.0f;
    float a1 = has_bias ? __ldg(b + oc0 + 1) : 0.0f;
    float a2 = has_bias ? __ldg(b + oc0 + 2) : 0.0f;
    float a3 = has_bias ? __ldg(b + oc0 + 3) : 0.0f;

    const int kd_min = max(0, od - (D - 1));
    const int kd_max = min(2, od);
    const int kh_min = max(0, oh - (H - 1));
    const int kh_max = min(2, oh);
    const int kw_min = max(0, ow - (W - 1));
    const int kw_max = min(2, ow);

    const int x_n_base = n * IC * D * H * W;

    for (int ic = 0; ic < IC; ++ic) {
        const int x_ic_base = x_n_base + ic * D * H * W;
        const int w_ic_base = ic * OC * K * K * K;
        for (int kd = kd_min; kd <= kd_max; ++kd) {
            const int iz = od - kd;
            const int x_d_base = x_ic_base + iz * H * W;
            const int wk_d = kd * 9;
            for (int kh = kh_min; kh <= kh_max; ++kh) {
                const int iy = oh - kh;
                const int x_h_base = x_d_base + iy * W;
                const int wk_dh = wk_d + kh * 3;
#pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    if (kw < kw_min || kw > kw_max) continue;
                    const int ix = ow - kw;
                    const float xv = __ldg(x + x_h_base + ix);
                    const int wk = wk_dh + kw;
                    const int wb0 = w_ic_base + (oc0 + 0) * 27 + wk;
                    a0 = fmaf(xv, __ldg(w + wb0),      a0);
                    a1 = fmaf(xv, __ldg(w + wb0 + 27), a1);
                    a2 = fmaf(xv, __ldg(w + wb0 + 54), a2);
                    a3 = fmaf(xv, __ldg(w + wb0 + 81), a3);
                }
            }
        }
    }

    const int out_sp = Do * Ho * Wo;
    const int y_base = (n * OC + oc0) * out_sp + od * Ho * Wo + oh * Wo + ow;
    y[y_base] = a0;
    y[y_base + out_sp] = a1;
    y[y_base + 2 * out_sp] = a2;
    y[y_base + 3 * out_sp] = a3;
}

__global__ void deconv3d_generic_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ b,
                                        float* __restrict__ y,
                                        int N, int IC, int D, int H, int W,
                                        int OC, int K, int stride, int padding,
                                        int dilation, int groups,
                                        int Do, int Ho, int Wo,
                                        int has_bias) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t total = (int64_t)N * OC * Do * Ho * Wo;
    if (idx >= total) return;

    int t0 = (int)idx;
    const int ow = t0 % Wo; t0 /= Wo;
    const int oh = t0 % Ho; t0 /= Ho;
    const int od = t0 % Do; t0 /= Do;
    const int oc = t0 % OC; t0 /= OC;
    const int n = t0;

    const int OCg = OC / groups;
    const int ICg = IC / groups;
    const int g = oc / OCg;
    const int ocg = oc - g * OCg;
    float acc = has_bias ? __ldg(b + oc) : 0.0f;

    for (int ii = 0; ii < ICg; ++ii) {
        const int ic = g * ICg + ii;
        for (int kd = 0; kd < K; ++kd) {
            int zd = od + padding - kd * dilation;
            if (zd < 0 || zd % stride) continue;
            int iz = zd / stride;
            if (iz < 0 || iz >= D) continue;
            for (int kh = 0; kh < K; ++kh) {
                int yh = oh + padding - kh * dilation;
                if (yh < 0 || yh % stride) continue;
                int iy = yh / stride;
                if (iy < 0 || iy >= H) continue;
                for (int kw = 0; kw < K; ++kw) {
                    int xw = ow + padding - kw * dilation;
                    if (xw < 0 || xw % stride) continue;
                    int ix = xw / stride;
                    if (ix < 0 || ix >= W) continue;
                    int64_t xo = (((int64_t)n * IC + ic) * D + iz) * H * W + iy * W + ix;
                    int64_t wo = ((((int64_t)ic * OCg + ocg) * K + kd) * K + kh) * K + kw;
                    acc = fmaf(__ldg(x + xo), __ldg(w + wo), acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor conv_transpose3d_forward(torch::Tensor x,
                                      torch::Tensor weight,
                                      torch::Tensor bias,
                                      int64_t stride,
                                      int64_t padding,
                                      int64_t output_padding,
                                      int64_t dilation,
                                      int64_t groups) {
    CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight); CHECK_CONTIGUOUS(bias);
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "only float32 is supported");
    const c10::cuda::CUDAGuard device_guard(x.device());

    const int N = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    const int OCg = (int)weight.size(1);
    const int K = (int)weight.size(2);
    const int OC = OCg * (int)groups;
    const int Do = (D - 1) * (int)stride - 2 * (int)padding + (int)dilation * (K - 1) + (int)output_padding + 1;
    const int Ho = (H - 1) * (int)stride - 2 * (int)padding + (int)dilation * (K - 1) + (int)output_padding + 1;
    const int Wo = (W - 1) * (int)stride - 2 * (int)padding + (int)dilation * (K - 1) + (int)output_padding + 1;

    auto y = torch::empty({N, OC, Do, Ho, Wo}, x.options());
    const int has_bias = (bias.numel() > 0) ? 1 : 0;
    const float* bp = has_bias ? bias.data_ptr<float>() : bias.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (IC == 48 && OC == 24 && K == 3 && stride == 1 && padding == 0 && output_padding == 0 && dilation == 1 && groups == 1) {
        const int spatial_total = N * Do * Ho * Wo;
        dim3 block(256);
        dim3 grid((spatial_total + block.x - 1) / block.x, 6);
        deconv3d_k3s1p0_48_24_oc4_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), N, D, H, W, has_bias);
    } else {
        int64_t total = (int64_t)N * OC * Do * Ho * Wo;
        dim3 block(256);
        dim3 grid((unsigned int)((total + block.x - 1) / block.x));
        deconv3d_generic_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),
            N, IC, D, H, W, OC, K, (int)stride, (int)padding, (int)dilation, (int)groups,
            Do, Ho, Wo, has_bias);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='deconv3d_transpose_ext_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: 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)
        self.stride = stride[0] if isinstance(stride, tuple) else stride
        self.padding = padding[0] if isinstance(padding, tuple) else padding
        self.output_padding = output_padding[0] if isinstance(output_padding, tuple) else output_padding
        self.dilation = dilation[0] if isinstance(dilation, tuple) else dilation
        self.groups = groups
        self.register_buffer('_empty_bias', torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        w = self.conv_transpose3d.weight.contiguous()
        b = self.conv_transpose3d.bias if self.conv_transpose3d.bias is not None else self._empty_bias
        if b.device != x.device:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv_transpose3d_forward(x, w, b.contiguous(), int(self.stride), int(self.padding),
                                             int(self.output_padding), int(self.dilation), int(self.groups))