import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 3D convolution with square input and 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, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        output_padding (int, optional): Additional size added to one side of the output shape. 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, 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, 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 = 48
kernel_size = 3
depth = 64
height = 64
width = 64

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,
                                       int stride, int padding, int output_padding, int groups, bool has_bias);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.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")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

constexpr int VEC = 8;

__global__ __launch_bounds__(128, 2)
void deconv3d_48_k3_s1p0_kernel(const float* __restrict__ x,
                                const float* __restrict__ w,
                                const float* __restrict__ bias,
                                float* __restrict__ y,
                                int B, int D, int H, int Wd,
                                int Do, int Ho, int Wo,
                                int has_bias) {
    const int outVol = Do * Ho * Wo;
    const int inHW = H * Wd;
    const int inVol = D * inHW;
    const int64_t total = (int64_t)B * outVol * 6; // 48 / 8

    int64_t job = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (job >= total) return;

    const int spatial = (int)(job % outVol);
    int64_t t = job / outVol;
    const int n = (int)(t % B);
    const int oc_base = (int)(t / B) * VEC;

    int rem = spatial;
    const int ow = rem % Wo;
    rem /= Wo;
    const int oh = rem % Ho;
    const int od = rem / Ho;

    float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f;
    float acc4 = 0.0f, acc5 = 0.0f, acc6 = 0.0f, acc7 = 0.0f;

    // Transposed convolution, stride=1, padding=0, kernel=3:
    // y[n,oc,od,oh,ow] = sum_ic,k x[n,ic,od-kd,oh-kh,ow-kw] * W[ic,oc,kd,kh,kw]
    #pragma unroll
    for (int kd = 0; kd < 3; ++kd) {
        const int id = od - kd;
        if ((unsigned)id >= (unsigned)D) continue;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            const int ih = oh - kh;
            if ((unsigned)ih >= (unsigned)H) continue;
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                const int iw = ow - kw;
                if ((unsigned)iw >= (unsigned)Wd) continue;
                const int kidx = (kd * 3 + kh) * 3 + kw;
                const int x_sp = id * inHW + ih * Wd + iw;
                #pragma unroll 4
                for (int ic = 0; ic < 48; ++ic) {
                    const float xv = __ldg(x + ((n * 48 + ic) * inVol + x_sp));
                    const int wb = ((ic * 48 + oc_base) * 27 + kidx);
                    acc0 = fmaf(xv, __ldg(w + wb + 0 * 27), acc0);
                    acc1 = fmaf(xv, __ldg(w + wb + 1 * 27), acc1);
                    acc2 = fmaf(xv, __ldg(w + wb + 2 * 27), acc2);
                    acc3 = fmaf(xv, __ldg(w + wb + 3 * 27), acc3);
                    acc4 = fmaf(xv, __ldg(w + wb + 4 * 27), acc4);
                    acc5 = fmaf(xv, __ldg(w + wb + 5 * 27), acc5);
                    acc6 = fmaf(xv, __ldg(w + wb + 6 * 27), acc6);
                    acc7 = fmaf(xv, __ldg(w + wb + 7 * 27), acc7);
                }
            }
        }
    }

    if (has_bias) {
        acc0 += __ldg(bias + oc_base + 0);
        acc1 += __ldg(bias + oc_base + 1);
        acc2 += __ldg(bias + oc_base + 2);
        acc3 += __ldg(bias + oc_base + 3);
        acc4 += __ldg(bias + oc_base + 4);
        acc5 += __ldg(bias + oc_base + 5);
        acc6 += __ldg(bias + oc_base + 6);
        acc7 += __ldg(bias + oc_base + 7);
    }

    float* yb = y + n * 48 * outVol + spatial;
    yb[(oc_base + 0) * outVol] = acc0;
    yb[(oc_base + 1) * outVol] = acc1;
    yb[(oc_base + 2) * outVol] = acc2;
    yb[(oc_base + 3) * outVol] = acc3;
    yb[(oc_base + 4) * outVol] = acc4;
    yb[(oc_base + 5) * outVol] = acc5;
    yb[(oc_base + 6) * outVol] = acc6;
    yb[(oc_base + 7) * outVol] = acc7;
}

__global__
void deconv3d_generic_kernel(const float* __restrict__ x,
                             const float* __restrict__ w,
                             const float* __restrict__ bias,
                             float* __restrict__ y,
                             int B, int IC, int D, int H, int Wd,
                             int OC, int kD, int kH, int kW,
                             int Do, int Ho, int Wo,
                             int stride, int padding, int groups, int has_bias) {
    const int64_t outVol = (int64_t)Do * Ho * Wo;
    const int64_t total = (int64_t)B * OC * outVol;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int spatial = (int)(idx % outVol);
    int64_t tmp = idx / outVol;
    int oc = (int)(tmp % OC);
    int n = (int)(tmp / OC);

    int rem = spatial;
    int ow = rem % Wo; rem /= Wo;
    int oh = rem % Ho; int od = rem / Ho;

    const int ICpg = IC / groups;
    const int OCpg = OC / groups;
    const int g = oc / OCpg;
    const int woc = oc - g * OCpg;
    const int ic0 = g * ICpg;
    const int ic1 = ic0 + ICpg;
    const int inHW = H * Wd;
    const int inVol = D * inHW;
    float acc = 0.0f;

    for (int ic = ic0; ic < ic1; ++ic) {
        for (int kd = 0; kd < kD; ++kd) {
            int td = od + padding - kd;
            if (td < 0 || td % stride) continue;
            int id = td / stride;
            if ((unsigned)id >= (unsigned)D) continue;
            for (int kh = 0; kh < kH; ++kh) {
                int th = oh + padding - kh;
                if (th < 0 || th % stride) continue;
                int ih = th / stride;
                if ((unsigned)ih >= (unsigned)H) continue;
                for (int kw = 0; kw < kW; ++kw) {
                    int tw = ow + padding - kw;
                    if (tw < 0 || tw % stride) continue;
                    int iw = tw / stride;
                    if ((unsigned)iw >= (unsigned)Wd) continue;
                    float xv = x[((n * IC + ic) * D + id) * inHW + ih * Wd + iw];
                    float wv = w[(((ic * OCpg + woc) * kD + kd) * kH + kh) * kW + kw];
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }
    if (has_bias) acc += bias[oc];
    y[idx] = acc;
}

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                       int stride, int padding, int output_padding, int groups, bool has_bias) {
    CHECK_CUDA(x);
    CHECK_CUDA(weight);
    CHECK_CONTIGUOUS(x);
    CHECK_CONTIGUOUS(weight);
    CHECK_FLOAT(x);
    CHECK_FLOAT(weight);
    if (has_bias) { CHECK_CUDA(bias); CHECK_CONTIGUOUS(bias); CHECK_FLOAT(bias); }

    const int B  = (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 Wd = (int)x.size(4);
    const int kD = (int)weight.size(2);
    const int kH = (int)weight.size(3);
    const int kW = (int)weight.size(4);
    const int OC = (int)weight.size(1) * groups;

    const int Do = (D - 1) * stride - 2 * padding + kD + output_padding;
    const int Ho = (H - 1) * stride - 2 * padding + kH + output_padding;
    const int Wo = (Wd - 1) * stride - 2 * padding + kW + output_padding;

    auto y = torch::empty({B, OC, Do, Ho, Wo}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (IC == 48 && OC == 48 && kD == 3 && kH == 3 && kW == 3 && stride == 1 && padding == 0 && output_padding == 0 && groups == 1) {
        const int64_t jobs = (int64_t)B * Do * Ho * Wo * 6;
        const int threads = 128;
        const dim3 blocks((unsigned int)((jobs + threads - 1) / threads));
        deconv3d_48_k3_s1p0_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), B, D, H, Wd, Do, Ho, Wo, has_bias ? 1 : 0);
    } else {
        const int64_t total = (int64_t)B * OC * Do * Ho * Wo;
        const int threads = 256;
        const dim3 blocks((unsigned int)((total + threads - 1) / threads));
        deconv3d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), B, IC, D, H, Wd, OC, kD, kH, kW, Do, Ho, Wo, stride, padding, groups, has_bias ? 1 : 0);
    }
    return y;
}
"""

_ext = load_inline(
    name="conv_transpose3d_custom_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv_transpose3d_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    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(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels,
            kernel_size=(kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, output_padding=output_padding,
            groups=groups, bias=bias
        )
        self.stride = stride if isinstance(stride, int) else stride[0]
        self.padding = padding if isinstance(padding, int) else padding[0]
        self.output_padding = output_padding if isinstance(output_padding, int) else output_padding[0]
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # The benchmark uses contiguous CUDA fp32 tensors.  Make parameters contiguous if a caller has
        # altered their layout; this does not change numerical behavior.
        xc = x.contiguous()
        w = self.conv_transpose3d.weight.contiguous()
        b = self.conv_transpose3d.bias
        has_b = b is not None
        if has_b:
            bc = b.contiguous()
        else:
            bc = torch.empty(0, device=xc.device, dtype=xc.dtype)
        return _ext.conv_transpose3d_forward(xc, w, bc, int(self.stride), int(self.padding), int(self.output_padding), int(self.groups), bool(has_b))