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_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                             int stride, int padding, int dilation, int groups, bool has_bias);
'''

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

__global__ __launch_bounds__(128, 2)
void conv3d_c3k357_oc64_kernel(const float* __restrict__ x,
                               const float* __restrict__ w,
                               const float* __restrict__ b,
                               float* __restrict__ y,
                               int M, int OW, int OH, int OD,
                               int IW, int IH, int ID,
                               bool has_bias) {
    int tx = threadIdx.x;
    int ocy = threadIdx.y;
    int ocg = blockIdx.y * 4 + ocy;
    int oc0 = ocg * 4;
    int base_p0 = blockIdx.x * (blockDim.x * PVAL) + tx;

    float acc[PVAL][4];
#pragma unroll
    for (int p=0; p<PVAL; ++p) {
#pragma unroll
        for (int j=0; j<4; ++j) acc[p][j] = has_bias ? __ldg(b + oc0 + j) : 0.0f;
    }

    int in_base[PVAL];
    int spatial_idx[PVAL];
    int n_idx[PVAL];
    bool valid[PVAL];
#pragma unroll
    for (int pp=0; pp<PVAL; ++pp) {
        int m = base_p0 + pp * 32;
        valid[pp] = (m < M);
        int t = m;
        int ow = t % OW; t /= OW;
        int oh = t % OH; t /= OH;
        int od = t % OD; t /= OD;
        int n = t;
        in_base[pp] = ((n * 3 * ID + od) * IH + oh) * IW + ow;
        spatial_idx[pp] = m - n * (OD*OH*OW);
        n_idx[pp] = n;
    }

    const int IDHW = ID * IH * IW;
    const int IHW = IH * IW;
#pragma unroll
    for (int c=0; c<3; ++c) {
#pragma unroll
        for (int kd=0; kd<3; ++kd) {
#pragma unroll
            for (int kh=0; kh<5; ++kh) {
#pragma unroll
                for (int kw=0; kw<7; ++kw) {
                    int kk = ((c*3 + kd)*5 + kh)*7 + kw;
                    const float ww0 = __ldg(w + (oc0+0)*315 + kk);
                    const float ww1 = __ldg(w + (oc0+1)*315 + kk);
                    const float ww2 = __ldg(w + (oc0+2)*315 + kk);
                    const float ww3 = __ldg(w + (oc0+3)*315 + kk);
                    int off = c*IDHW + kd*IHW + kh*IW + kw;
#pragma unroll
                    for (int pp=0; pp<PVAL; ++pp) {
                        float xv = valid[pp] ? __ldg(x + in_base[pp] + off) : 0.0f;
                        acc[pp][0] = fmaf(xv, ww0, acc[pp][0]);
                        acc[pp][1] = fmaf(xv, ww1, acc[pp][1]);
                        acc[pp][2] = fmaf(xv, ww2, acc[pp][2]);
                        acc[pp][3] = fmaf(xv, ww3, acc[pp][3]);
                    }
                }
            }
        }
    }
    const long long plane = (long long)OD * OH * OW;
#pragma unroll
    for (int pp=0; pp<PVAL; ++pp) {
        if (valid[pp]) {
            long long base = ((long long)n_idx[pp] * 64 + oc0) * plane + spatial_idx[pp];
            y[base + 0*plane] = acc[pp][0];
            y[base + 1*plane] = acc[pp][1];
            y[base + 2*plane] = acc[pp][2];
            y[base + 3*plane] = acc[pp][3];
        }
    }
}

__global__ void conv3d_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,
                                      const float* __restrict__ b, float* __restrict__ y,
                                      int N, int C, int OC, int ID, int IH, int IW,
                                      int KD, int KH, int KW, int OD, int OH, int OW,
                                      int stride, int pad, int dil, bool has_bias) {
    long long total = (long long)N*OC*OD*OH*OW;
    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW; long long t = idx / OW;
    int oh = t % OH; t /= OH; int od = t % OD; t /= OD; int oc = t % OC; int n = t / OC;
    float acc = has_bias ? b[oc] : 0.0f;
    for (int c=0;c<C;++c) for (int kd=0;kd<KD;++kd) { int id=od*stride+kd*dil-pad; if ((unsigned)id>=(unsigned)ID) continue;
      for (int kh=0;kh<KH;++kh) { int ih=oh*stride+kh*dil-pad; if ((unsigned)ih>=(unsigned)IH) continue;
        for (int kw=0;kw<KW;++kw) { int iw=ow*stride+kw*dil-pad; if ((unsigned)iw>=(unsigned)IW) continue;
          float xv=x[(((long long)n*C+c)*ID+id)*IH*IW+ih*IW+iw];
          float ww=w[(((long long)oc*C+c)*KD+kd)*KH*KW+kh*KW+kw];
          acc=fmaf(xv,ww,acc);
        }}}
    y[idx]=acc;
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                             int stride, int padding, int dilation, int groups, bool has_bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors expected");
    auto xc=x.contiguous(); auto wc=w.contiguous(); torch::Tensor bc=bias; if (has_bias) bc=bias.contiguous();
    int N=xc.size(0), C=xc.size(1), ID=xc.size(2), IH=xc.size(3), IW=xc.size(4);
    int OC=wc.size(0), KD=wc.size(2), KH=wc.size(3), KW=wc.size(4);
    int OD=(ID+2*padding-dilation*(KD-1)-1)/stride+1, OH=(IH+2*padding-dilation*(KH-1)-1)/stride+1, OW=(IW+2*padding-dilation*(KW-1)-1)/stride+1;
    auto y=torch::empty({N,OC,OD,OH,OW}, xc.options());
    const float* bp=has_bias?bc.data_ptr<float>():nullptr; cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    if (N==16&&C==3&&OC==64&&ID==64&&IH==64&&IW==64&&KD==3&&KH==5&&KW==7&&stride==1&&padding==0&&dilation==1&&groups==1) {
        int M=N*OD*OH*OW; dim3 block(32,4,1); dim3 grid((M+32*PVAL-1)/(32*PVAL),4,1);
        conv3d_c3k357_oc64_kernel<<<grid,block,0,stream>>>(xc.data_ptr<float>(),wc.data_ptr<float>(),bp,y.data_ptr<float>(),M,OW,OH,OD,IW,IH,ID,has_bias);
    } else {
        long long total=(long long)N*OC*OD*OH*OW; int block=256; int grid=(total+block-1)/block;
        conv3d_generic_kernel<<<grid,block,0,stream>>>(xc.data_ptr<float>(),wc.data_ptr<float>(),bp,y.data_ptr<float>(),N,C,OC,ID,IH,IW,KD,KH,KW,OD,OH,OW,stride,padding,dilation,has_bias);
    }
    return y;
}
'''
_ext=load_inline(name='conv3d_asym_ext_v2', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['conv3d_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)

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)
        self.stride=stride if isinstance(stride,int) else stride[0]; self.padding=padding if isinstance(padding,int) else padding[0]; self.dilation=dilation if isinstance(dilation,int) else dilation[0]; self.groups=groups
    def forward(self,x:torch.Tensor)->torch.Tensor:
        bias=self.conv3d.bias if self.conv3d.bias is not None else torch.empty(0,device=x.device,dtype=x.dtype)
        return _ext.conv3d_forward(x,self.conv3d.weight,bias,int(self.stride),int(self.padding),int(self.dilation),int(self.groups),self.conv3d.bias is not None)