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>
#include <vector>

torch::Tensor conv_transpose3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                            std::vector<int64_t> stride, std::vector<int64_t> padding,
                                            std::vector<int64_t> output_padding, int64_t groups, bool has_bias);

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                       std::vector<int64_t> stride, std::vector<int64_t> padding,
                                       std::vector<int64_t> output_padding, int64_t groups, bool has_bias) {
    return conv_transpose3d_forward_cuda(x, weight, bias, stride, padding, output_padding, groups, has_bias);
}
'''

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

__global__ void pack_w_kernel(const float* __restrict__ w, float* __restrict__ wp,
                              int Ci, int CoG, int KD, int KH, int KW, int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int t = idx;
    int kw = t % KW; t /= KW;
    int kh = t % KH; t /= KH;
    int kd = t % KD; t /= KD;
    int co = t % CoG; t /= CoG;
    int ci = t;
    int out = (((ci * KD + kd) * KH + kh) * KW + kw) * CoG + co;
    wp[out] = w[idx];
}

// Optimized path for the benchmark case: groups=1 and out_channels=16.
// One half-warp computes all 16 output channels for one (n,d,h,w) output position.
__global__ void deconv3d_s1_co16_kernel(const float* __restrict__ x,
                                        const float* __restrict__ wp,
                                        const float* __restrict__ bias,
                                        float* __restrict__ y,
                                        int N, int Ci, int D, int H, int W,
                                        int Do, int Ho, int Wo,
                                        int KD, int KH, int KW,
                                        int pD, int pH, int pW,
                                        bool has_bias, long long positions) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int lane = threadIdx.x & 31;
    int half = lane >> 4;
    int co = lane & 15;
    long long pos = ((long long)tid) >> 4;
    if (pos >= positions) return;

    int ow = (int)(pos % Wo);
    long long q = pos / Wo;
    int oh = (int)(q % Ho); q /= Ho;
    int od = (int)(q % Do); q /= Do;
    int n = (int)q;

    float acc = has_bias ? __ldg(bias + co) : 0.0f;

    int odpg = od + pD;
    int ohpg = oh + pH;
    int owpg = ow + pW;

    int kd0 = odpg - (D - 1); if (kd0 < 0) kd0 = 0;
    int kh0 = ohpg - (H - 1); if (kh0 < 0) kh0 = 0;
    int kw0 = owpg - (W - 1); if (kw0 < 0) kw0 = 0;
    int kd1 = odpg; if (kd1 > KD - 1) kd1 = KD - 1;
    int kh1 = ohpg; if (kh1 > KH - 1) kh1 = KH - 1;
    int kw1 = owpg; if (kw1 > KW - 1) kw1 = KW - 1;

    unsigned mask = half ? 0xFFFF0000u : 0x0000FFFFu;
    int leader = half << 4;

    long long nbase = ((long long)n * Ci * D * H * W);
    for (int ci = 0; ci < Ci; ++ci) {
        long long cbase = nbase + (long long)ci * D * H * W;
        for (int kd = kd0; kd <= kd1; ++kd) {
            int id = odpg - kd;
            long long dbase = cbase + (long long)id * H * W;
            for (int kh = kh0; kh <= kh1; ++kh) {
                int ih = ohpg - kh;
                long long hbase = dbase + (long long)ih * W;
                int wpo_base = (((ci * KD + kd) * KH + kh) * KW + kw0) * 16 + co;
                for (int kw = kw0; kw <= kw1; ++kw) {
                    float xv = 0.0f;
                    if (co == 0) {
                        int iw = owpg - kw;
                        xv = __ldg(x + hbase + iw);
                    }
                    xv = __shfl_sync(mask, xv, leader, 32);
                    float wv = __ldg(wp + wpo_base);
                    acc = fmaf(xv, wv, acc);
                    wpo_base += 16;
                }
            }
        }
    }
    long long yidx = (((((long long)n * 16 + co) * Do + od) * Ho + oh) * Wo + ow);
    y[yidx] = acc;
}

__global__ void deconv3d_generic_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ bias,
                                        float* __restrict__ y,
                                        int N, int Ci, int Co, int D, int H, int W,
                                        int Do, int Ho, int Wo,
                                        int KD, int KH, int KW,
                                        int sD, int sH, int sW,
                                        int pD, int pH, int pW,
                                        int groups, bool has_bias,
                                        long long total) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = (int)(idx % Wo); long long t = idx / Wo;
    int oh = (int)(t % Ho); t /= Ho;
    int od = (int)(t % Do); t /= Do;
    int co = (int)(t % Co); t /= Co;
    int n = (int)t;
    int CoG = Co / groups;
    int CiG = Ci / groups;
    int g = co / CoG;
    int cog = co - g * CoG;
    float acc = has_bias ? __ldg(bias + co) : 0.0f;
    for (int cig = 0; cig < CiG; ++cig) {
        int ci = g * CiG + cig;
        for (int kd = 0; kd < KD; ++kd) {
            int idn = od + pD - kd;
            if (idn < 0 || idn % sD) continue;
            int id = idn / sD;
            if (id < 0 || id >= D) continue;
            for (int kh = 0; kh < KH; ++kh) {
                int ihn = oh + pH - kh;
                if (ihn < 0 || ihn % sH) continue;
                int ih = ihn / sH;
                if (ih < 0 || ih >= H) continue;
                for (int kw = 0; kw < KW; ++kw) {
                    int iwn = ow + pW - kw;
                    if (iwn < 0 || iwn % sW) continue;
                    int iw = iwn / sW;
                    if (iw < 0 || iw >= W) continue;
                    long long xidx = (((((long long)n * Ci + ci) * D + id) * H + ih) * W + iw);
                    long long widx = (((((long long)ci * CoG + cog) * KD + kd) * KH + kh) * KW + kw);
                    acc = fmaf(__ldg(x + xidx), __ldg(w + widx), acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor conv_transpose3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                            std::vector<int64_t> stride, std::vector<int64_t> padding,
                                            std::vector<int64_t> output_padding, int64_t groups64, bool has_bias) {
    int N = (int)x.size(0);
    int Ci = (int)x.size(1);
    int D = (int)x.size(2);
    int H = (int)x.size(3);
    int W = (int)x.size(4);
    int CoG = (int)weight.size(1);
    int KD = (int)weight.size(2);
    int KH = (int)weight.size(3);
    int KW = (int)weight.size(4);
    int groups = (int)groups64;
    int Co = CoG * groups;
    int sD = (int)stride[0], sH = (int)stride[1], sW = (int)stride[2];
    int pD = (int)padding[0], pH = (int)padding[1], pW = (int)padding[2];
    int opD = (int)output_padding[0], opH = (int)output_padding[1], opW = (int)output_padding[2];
    int Do = (D - 1) * sD - 2 * pD + KD + opD;
    int Ho = (H - 1) * sH - 2 * pH + KH + opH;
    int Wo = (W - 1) * sW - 2 * pW + KW + opW;
    auto y = torch::empty({N, Co, Do, Ho, Wo}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat && groups == 1 && Co == 16 && sD == 1 && sH == 1 && sW == 1) {
        auto wp = torch::empty({Ci, KD, KH, KW, 16}, weight.options());
        int totalw = Ci * 16 * KD * KH * KW;
        pack_w_kernel<<<(totalw + 255) / 256, 256, 0, stream>>>(weight.data_ptr<float>(), wp.data_ptr<float>(), Ci, 16, KD, KH, KW, totalw);
        long long positions = (long long)N * Do * Ho * Wo;
        int threads = 128;
        int blocks = (int)(((positions * 16) + threads - 1) / threads);
        const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
        deconv3d_s1_co16_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), wp.data_ptr<float>(), bptr, y.data_ptr<float>(),
            N, Ci, D, H, W, Do, Ho, Wo, KD, KH, KW, pD, pH, pW, has_bias, positions);
    } else {
        long long total = (long long)N * Co * Do * Ho * Wo;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
        deconv3d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(),
            N, Ci, Co, D, H, W, Do, Ho, Wo, KD, KH, KW, sD, sH, sW, pD, pH, pW, groups, has_bias, total);
    }
    return y;
}
'''

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # The benchmark tensors/parameters are contiguous; these preserve correctness for the common case at no-copy cost.
        if not x.is_contiguous():
            x = x.contiguous()
        w = self.conv_transpose3d.weight
        if not w.is_contiguous():
            w = w.contiguous()
        b = self.conv_transpose3d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
            has_b = False
        else:
            has_b = True
        return _ext.conv_transpose3d_forward(x, w, b, list(self.stride), list(self.padding), list(self.output_padding), int(self.groups), bool(has_b))