import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 3D transposed convolution operation with asymmetric input and square kernel.
    The input is padded before the convolution.

    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, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 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: int, stride: int = 1, padding: int = 0, output_padding: int = 0, 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, kernel_size), stride=stride, padding=padding, groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 3D transposed 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 = 4
in_channels = 32
out_channels = 32
kernel_size = 3
depth = 32
height = 64
width = 128
stride = 2
padding = 1
groups = 4

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, stride, padding, groups]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp = r'''
#include <torch/extension.h>
torch::Tensor convt3d_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias_opt,
                              int64_t stride, int64_t padding, int64_t groups);
'''

_cuda = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float tf32r(float a) {
    unsigned u = __float_as_uint(a);
    unsigned r = 0x00000fffu + ((u >> 13) & 1u);
    u += r;
    u &= 0xffffe000u;
    return __uint_as_float(u);
}

__global__ void fast_g1_k3s2p1(const float* __restrict__ x, const float* __restrict__ w,
                               const float* __restrict__ b, float* __restrict__ y,
                               long spatial, int D, int H, int W, int OD, int OH, int OW) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    long s = blockIdx.x * (long)(blockDim.x >> 5) + warp;
    if (s >= spatial) return;
    int ow = s % OW;
    long t = s / OW;
    int oh = t % OH; t /= OH;
    int od = t % OD;
    int n = t / OD;

    float acc = b ? b[lane] : 0.0f;
    int kdA[2], khA[2], kwA[2], nd, nh, nw;
    if ((od & 1) == 0) { kdA[0] = 1; nd = 1; } else { kdA[0] = 0; kdA[1] = 2; nd = 2; }
    if ((oh & 1) == 0) { khA[0] = 1; nh = 1; } else { khA[0] = 0; khA[1] = 2; nh = 2; }
    if ((ow & 1) == 0) { kwA[0] = 1; nw = 1; } else { kwA[0] = 0; kwA[1] = 2; nw = 2; }

    unsigned mask = 0xffffffffu;
    int DHW = D * H * W;
    #pragma unroll
    for (int ic = 0; ic < 32; ++ic) {
        const float* xb = x + ((n * 32 + ic) * DHW);
        const float* wb = w + (ic * 32 + lane) * 27;
        for (int kdi = 0; kdi < nd; ++kdi) {
            int kd = kdA[kdi]; int id = (od + 1 - kd) >> 1;
            if ((unsigned)id >= (unsigned)D) continue;
            for (int khi = 0; khi < nh; ++khi) {
                int kh = khA[khi]; int ih = (oh + 1 - kh) >> 1;
                if ((unsigned)ih >= (unsigned)H) continue;
                for (int kwi = 0; kwi < nw; ++kwi) {
                    int kw = kwA[kwi]; int iw = (ow + 1 - kw) >> 1;
                    if ((unsigned)iw >= (unsigned)W) continue;
                    float xv = 0.0f;
                    if (lane == 0) xv = xb[(id * H + ih) * W + iw];
                    xv = __shfl_sync(mask, xv, 0);
                    acc = fmaf(tf32r(xv), tf32r(wb[(kd * 3 + kh) * 3 + kw]), acc);
                }
            }
        }
    }
    y[((n * 32 + lane) * OD + od) * OH * OW + oh * OW + ow] = acc;
}

__global__ void generic_k(const float* __restrict__ x, const float* __restrict__ w,
                          const float* __restrict__ b, float* __restrict__ y, long total,
                          int N, int Cin, int D, int H, int W, int Cout, int K,
                          int stride, int pad, int groups, int OD, int OH, int OW) {
    long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW; long t = idx / OW;
    int oh = t % OH; t /= OH;
    int od = t % OD; t /= OD;
    int co = t % Cout;
    int n = t / Cout;
    int opg = Cout / groups, ipg = Cin / groups, g = co / opg, cog = co - g * opg, ic0 = g * ipg;
    float acc = b ? b[co] : 0.0f;
    for (int icg = 0; icg < ipg; ++icg) {
        int ic = ic0 + icg;
        for (int kd = 0; kd < K; ++kd) {
            int idp = od + pad - kd; if (idp < 0 || idp % stride) continue;
            int id = idp / stride; if (id >= D) continue;
            for (int kh = 0; kh < K; ++kh) {
                int ihp = oh + pad - kh; if (ihp < 0 || ihp % stride) continue;
                int ih = ihp / stride; if (ih >= H) continue;
                for (int kw = 0; kw < K; ++kw) {
                    int iwp = ow + pad - kw; if (iwp < 0 || iwp % stride) continue;
                    int iw = iwp / stride; if (iw >= W) continue;
                    float xv = x[(((n * Cin + ic) * D + id) * H + ih) * W + iw];
                    float wv = w[(((ic * opg + cog) * K + kd) * K + kh) * K + kw];
                    acc = fmaf(tf32r(xv), tf32r(wv), acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor convt3d_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias_opt,
                              int64_t stride, int64_t padding, int64_t groups) {
    x = x.contiguous(); weight = weight.contiguous();
    int N = x.size(0), Cin = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4);
    int K = weight.size(2), Cout = weight.size(1) * groups;
    int OD = (D - 1) * stride - 2 * padding + K;
    int OH = (H - 1) * stride - 2 * padding + K;
    int OW = (W - 1) * stride - 2 * padding + K;
    auto y = torch::empty({N, Cout, OD, OH, OW}, x.options());
    const float* bp = nullptr; torch::Tensor bias;
    if (bias_opt.has_value() && bias_opt.value().defined()) { bias = bias_opt.value().contiguous(); bp = bias.data_ptr<float>(); }
    if (Cin == 32 && Cout == 32 && groups == 1 && K == 3 && stride == 2 && padding == 1) {
        long sp = (long)N * OD * OH * OW;
        int th = 128, bl = (int)((sp + 3) / 4);
        fast_g1_k3s2p1<<<bl, th>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), sp, D, H, W, OD, OH, OW);
    } else {
        long total = (long)N * Cout * OD * OH * OW;
        int th = 256, bl = (int)((total + th - 1) / th);
        generic_k<<<bl, th>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), total, N, Cin, D, H, W, Cout, K, stride, padding, groups, OD, OH, OW);
    }
    return y;
}
'''

_ext = load_inline(name='convt3d_final_ext_a', cpp_sources=_cpp, cuda_sources=_cuda,
                   functions=['convt3d_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], 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, groups: int = 1, bias: bool = False):
        super().__init__()
        self.stride = stride
        self.padding = padding
        self.groups = groups
        self.conv_transpose3d = nn.ConvTranspose3d(in_channels, out_channels,
            (kernel_size, kernel_size, kernel_size), stride=stride, padding=padding, groups=groups, bias=bias)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.convt3d_forward(x, self.conv_transpose3d.weight, self.conv_transpose3d.bias,
                                    self.stride, self.padding, self.groups)