import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 3D convolution operation with a square input and an asymmetric kernel.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Size of the convolution kernel (kernel_width, kernel_height, kernel_depth).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int or tuple, optional): Padding applied to the input. 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: tuple, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 3D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, width, height, depth).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, width_out, height_out, depth_out).
        """
        return self.conv3d(x)

# Test code
batch_size = 16
in_channels = 3
out_channels = 64
kernel_size = (3, 5, 7)  # Asymmetric kernel
width = 64
height = 64
depth = 64

def get_inputs():
    x = torch.rand(batch_size, in_channels, width, height, depth)
    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 conv3d_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                          int64_t sD, int64_t sH, int64_t sW,
                          int64_t pD, int64_t pH, int64_t pW,
                          int64_t dD, int64_t dH, int64_t dW);

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t sD, int64_t sH, int64_t sW,
                             int64_t pD, int64_t pH, int64_t pW,
                             int64_t dD, int64_t dH, int64_t dW) {
    return conv3d_cuda(x, weight, bias, sD, sH, sW, pD, pH, pW, dD, dH, dW);
}
'''

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

__global__ __launch_bounds__(256)
void conv3d_3_64_3x5x7_d2_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 Do, int Ho, int Wo,
                                  int has_bias) {
    extern __shared__ float sw[];
    const int tx = threadIdx.x;        // 0..63, output W
    const int ty = threadIdx.y;        // 0..3, output H
    const int tid = ty * 64 + tx;

    int gz = blockIdx.z;
    const int od_pair = gz % ((Do + 1) >> 1);
    gz /= ((Do + 1) >> 1);
    const int oc = gz & 63;            // OC == 64
    const int n = gz >> 6;

    for (int i = tid; i < 315; i += 256) {
        sw[i] = w[oc * 315 + i];
    }
    __syncthreads();

    const int ow = blockIdx.x * 64 + tx;
    const int oh = blockIdx.y * 4 + ty;
    const int od0 = od_pair << 1;

    if (ow >= Wo || oh >= Ho || n >= N || od0 >= Do) return;

    float sum0 = has_bias ? b[oc] : 0.0f;
    float sum1 = has_bias ? b[oc] : 0.0f;

    const int HW = H * W;
    const int in_n_base = n * 3 * D * HW;

    #pragma unroll
    for (int c = 0; c < 3; ++c) {
        const int cbase = in_n_base + c * D * HW;
        #pragma unroll
        for (int kh = 0; kh < 5; ++kh) {
            const int row_base = cbase + od0 * HW + (oh + kh) * W + ow;
            #pragma unroll
            for (int kw = 0; kw < 7; ++kw) {
                const float v0 = x[row_base + kw];
                const float v1 = x[row_base + HW + kw];
                const float v2 = x[row_base + 2 * HW + kw];
                const float v3 = x[row_base + 3 * HW + kw];
                const int wbase = c * 105 + kh * 7 + kw;
                const float w0 = sw[wbase];
                const float w1 = sw[wbase + 35];
                const float w2 = sw[wbase + 70];
                sum0 = fmaf(v0, w0, sum0);
                sum0 = fmaf(v1, w1, sum0);
                sum0 = fmaf(v2, w2, sum0);
                sum1 = fmaf(v1, w0, sum1);
                sum1 = fmaf(v2, w1, sum1);
                sum1 = fmaf(v3, w2, sum1);
            }
        }
    }

    const int out_base = ((n * 64 + oc) * Do + od0) * Ho * Wo + oh * Wo + ow;
    y[out_base] = sum0;
    if (od0 + 1 < Do) {
        y[out_base + Ho * Wo] = sum1;
    }
}

__global__ __launch_bounds__(256)
void conv3d_generic_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           long total,
                           int N, int C, int D, int H, int Wd,
                           int OC, int kD, int kH, int kW,
                           int Do, int Ho, int Wo,
                           int sD, int sH, int sW,
                           int pD, int pH, int pW,
                           int dilD, int dilH, int dilW,
                           int groups, int has_bias) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

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

    int ocpg = OC / groups;
    int cpg = C / groups;
    int g = oc / ocpg;
    int c_start = g * cpg;

    float sum = has_bias ? b[oc] : 0.0f;
    for (int ci = 0; ci < cpg; ++ci) {
        int c = c_start + ci;
        for (int kd = 0; kd < kD; ++kd) {
            int id = od * sD - pD + kd * dilD;
            if ((unsigned)id >= (unsigned)D) continue;
            for (int kh = 0; kh < kH; ++kh) {
                int ih = oh * sH - pH + kh * dilH;
                if ((unsigned)ih >= (unsigned)H) continue;
                for (int kw = 0; kw < kW; ++kw) {
                    int iw = ow * sW - pW + kw * dilW;
                    if ((unsigned)iw >= (unsigned)Wd) continue;
                    float xv = x[((n * C + c) * D + id) * H * Wd + ih * Wd + iw];
                    float wv = w[(((oc * cpg + ci) * kD + kd) * kH + kh) * kW + kw];
                    sum = fmaf(xv, wv, sum);
                }
            }
        }
    }
    y[(((n * OC + oc) * Do + od) * Ho + oh) * Wo + ow] = sum;
}

torch::Tensor conv3d_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                          int64_t sD64, int64_t sH64, int64_t sW64,
                          int64_t pD64, int64_t pH64, int64_t pW64,
                          int64_t dD64, int64_t dH64, int64_t dW64) {
    const int N = (int)x.size(0);
    const int C = (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 OC = (int)weight.size(0);
    const int Cw = (int)weight.size(1);
    const int kD = (int)weight.size(2);
    const int kH = (int)weight.size(3);
    const int kW = (int)weight.size(4);
    const int sD = (int)sD64, sH = (int)sH64, sW = (int)sW64;
    const int pD = (int)pD64, pH = (int)pH64, pW = (int)pW64;
    const int dilD = (int)dD64, dilH = (int)dH64, dilW = (int)dW64;
    const int Do = (D + 2 * pD - dilD * (kD - 1) - 1) / sD + 1;
    const int Ho = (H + 2 * pH - dilH * (kH - 1) - 1) / sH + 1;
    const int Wo = (Wd + 2 * pW - dilW * (kW - 1) - 1) / sW + 1;
    auto y = torch::empty({N, OC, Do, Ho, Wo}, x.options());
    if (y.numel() == 0) return y;

    const int has_bias = (bias.numel() != 0);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat &&
        C == 3 && OC == 64 && Cw == 3 && kD == 3 && kH == 5 && kW == 7 &&
        sD == 1 && sH == 1 && sW == 1 && pD == 0 && pH == 0 && pW == 0 &&
        dilD == 1 && dilH == 1 && dilW == 1) {
        dim3 block(64, 4, 1);
        dim3 grid((Wo + 63) / 64, (Ho + 3) / 4, N * OC * ((Do + 1) / 2));
        conv3d_3_64_3x5x7_d2_kernel<<<grid, block, 315 * sizeof(float), stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(),
            N, D, H, Wd, Do, Ho, Wo, has_bias);
    } else {
        int groups = C / Cw;
        long total = (long)N * OC * Do * Ho * Wo;
        dim3 block(256);
        dim3 grid((total + 255) / 256);
        conv3d_generic_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(), total,
            N, C, D, H, Wd, OC, kD, kH, kW, Do, Ho, Wo,
            sD, sH, sW, pD, pH, pW, dilD, dilH, dilW, groups, has_bias);
    }
    return y;
}
'''

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


def _triple(v):
    if isinstance(v, tuple):
        return v
    return (v, v, v)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        st = self.conv3d.stride
        pa = self.conv3d.padding
        di = self.conv3d.dilation
        b = self.conv3d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv3d_forward(x.contiguous(), self.conv3d.weight.contiguous(), b,
                                   int(st[0]), int(st[1]), int(st[2]),
                                   int(pa[0]), int(pa[1]), int(pa[2]),
                                   int(di[0]), int(di[1]), int(di[2]))