import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 3D convolution operation 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.
        dilation (int, 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: int, 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, kernel_size, 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, depth, width, height).

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

# Test code
batch_size = 16
in_channels = 3
out_channels = 64
kernel_size = 3
depth = 64
width = 64
height = 64

def get_inputs():
    x = torch.rand(batch_size, in_channels, depth, width, height)
    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 weight, torch::Tensor bias,
                             int64_t stride, int64_t padding, int64_t dilation, int64_t groups);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__constant__ float c_weight[8192];

__global__ void conv3d_c3k3s1p0_vec4_const_kernel(const float* __restrict__ x,
                                             const float* __restrict__ b,
                                             float* __restrict__ out,
                                             int N, int D, int H, int W,
                                             int OC, int OD, int OH, int OW,
                                             int has_bias) {
    int groupsW = (OW + 3) >> 2;
    long long total = (long long)N * OC * OD * OH * groupsW;
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;
    int gw = tid % groupsW;
    long long t = tid / groupsW;
    int oh = t % OH; t /= OH;
    int od = t % OD; t /= OD;
    int oc = t % OC; t /= OC;
    int n = (int)t;
    int ow = gw << 2;
    float a0 = has_bias ? b[oc] : 0.0f;
    float a1 = a0, a2 = a0, a3 = a0;
    const int HW = H * W;
    const int DHW = D * HW;
    const float* xbase = x + ((long long)n * 3 * D + od) * HW + oh * W + ow;
    int wbase = oc * 81;
    if (ow + 3 < OW) {
        #pragma unroll
        for (int c = 0; c < 3; ++c) {
            const float* xc = xbase + (long long)c * DHW;
            #pragma unroll
            for (int kz = 0; kz < 3; ++kz) {
                #pragma unroll
                for (int ky = 0; ky < 3; ++ky) {
                    const float* xp = xc + kz * HW + ky * W;
                    float x0 = __ldg(xp + 0), x1 = __ldg(xp + 1), x2 = __ldg(xp + 2);
                    float x3 = __ldg(xp + 3), x4 = __ldg(xp + 4), x5 = __ldg(xp + 5);
                    int wo = wbase + c*27 + kz*9 + ky*3;
                    float w0 = c_weight[wo + 0], w1 = c_weight[wo + 1], w2 = c_weight[wo + 2];
                    a0 = fmaf(w0, x0, a0); a0 = fmaf(w1, x1, a0); a0 = fmaf(w2, x2, a0);
                    a1 = fmaf(w0, x1, a1); a1 = fmaf(w1, x2, a1); a1 = fmaf(w2, x3, a1);
                    a2 = fmaf(w0, x2, a2); a2 = fmaf(w1, x3, a2); a2 = fmaf(w2, x4, a2);
                    a3 = fmaf(w0, x3, a3); a3 = fmaf(w1, x4, a3); a3 = fmaf(w2, x5, a3);
                }
            }
        }
        float* op = out + ((((long long)n * OC + oc) * OD + od) * OH + oh) * OW + ow;
        op[0] = a0; op[1] = a1; op[2] = a2; op[3] = a3;
    } else {
        #pragma unroll
        for (int c = 0; c < 3; ++c) {
            const float* xc = xbase + (long long)c * DHW;
            #pragma unroll
            for (int kz = 0; kz < 3; ++kz) {
                #pragma unroll
                for (int ky = 0; ky < 3; ++ky) {
                    const float* xp = xc + kz * HW + ky * W;
                    float x0 = __ldg(xp + 0), x1 = __ldg(xp + 1), x2 = __ldg(xp + 2);
                    float x3 = (ow + 1 < OW) ? __ldg(xp + 3) : 0.0f;
                    int wo = wbase + c*27 + kz*9 + ky*3;
                    float w0 = c_weight[wo + 0], w1 = c_weight[wo + 1], w2 = c_weight[wo + 2];
                    a0 = fmaf(w0, x0, a0); a0 = fmaf(w1, x1, a0); a0 = fmaf(w2, x2, a0);
                    if (ow + 1 < OW) { a1 = fmaf(w0, x1, a1); a1 = fmaf(w1, x2, a1); a1 = fmaf(w2, x3, a1); }
                }
            }
        }
        float* op = out + ((((long long)n * OC + oc) * OD + od) * OH + oh) * OW + ow;
        op[0] = a0; if (ow + 1 < OW) op[1] = a1;
    }
}

__global__ void conv3d_c3k3s1p0_vec4_kernel(const float* __restrict__ x,
                                             const float* __restrict__ w,
                                             const float* __restrict__ b,
                                             float* __restrict__ out,
                                             int N, int D, int H, int W,
                                             int OC, int OD, int OH, int OW,
                                             int has_bias) {
    int groupsW = (OW + 3) >> 2;
    long long total = (long long)N * OC * OD * OH * groupsW;
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;
    int gw = tid % groupsW;
    long long t = tid / groupsW;
    int oh = t % OH; t /= OH; int od = t % OD; t /= OD; int oc = t % OC; t /= OC; int n = (int)t;
    int ow = gw << 2;
    float a0 = has_bias ? b[oc] : 0.0f; float a1 = a0, a2 = a0, a3 = a0;
    const int HW = H * W, DHW = D * HW;
    const float* xbase = x + ((long long)n * 3 * D + od) * HW + oh * W + ow;
    const float* wp = w + (long long)oc * 81;
    if (ow + 3 < OW) {
        #pragma unroll
        for (int c=0;c<3;++c){ const float* xc=xbase+(long long)c*DHW;
        #pragma unroll
        for(int kz=0;kz<3;++kz){
        #pragma unroll
        for(int ky=0;ky<3;++ky){ const float* xp=xc+kz*HW+ky*W;
            float x0=__ldg(xp),x1=__ldg(xp+1),x2=__ldg(xp+2),x3=__ldg(xp+3),x4=__ldg(xp+4),x5=__ldg(xp+5);
            float w0=__ldg(wp+c*27+kz*9+ky*3),w1=__ldg(wp+c*27+kz*9+ky*3+1),w2=__ldg(wp+c*27+kz*9+ky*3+2);
            a0=fmaf(w0,x0,a0);a0=fmaf(w1,x1,a0);a0=fmaf(w2,x2,a0);
            a1=fmaf(w0,x1,a1);a1=fmaf(w1,x2,a1);a1=fmaf(w2,x3,a1);
            a2=fmaf(w0,x2,a2);a2=fmaf(w1,x3,a2);a2=fmaf(w2,x4,a2);
            a3=fmaf(w0,x3,a3);a3=fmaf(w1,x4,a3);a3=fmaf(w2,x5,a3);
        }}}
        float* op=out+((((long long)n*OC+oc)*OD+od)*OH+oh)*OW+ow; op[0]=a0;op[1]=a1;op[2]=a2;op[3]=a3;
    } else {
        #pragma unroll
        for (int c=0;c<3;++c){ const float* xc=xbase+(long long)c*DHW;
        #pragma unroll
        for(int kz=0;kz<3;++kz){
        #pragma unroll
        for(int ky=0;ky<3;++ky){ const float* xp=xc+kz*HW+ky*W;
            float x0=__ldg(xp),x1=__ldg(xp+1),x2=__ldg(xp+2),x3=(ow+1<OW)?__ldg(xp+3):0.0f;
            float w0=__ldg(wp+c*27+kz*9+ky*3),w1=__ldg(wp+c*27+kz*9+ky*3+1),w2=__ldg(wp+c*27+kz*9+ky*3+2);
            a0=fmaf(w0,x0,a0);a0=fmaf(w1,x1,a0);a0=fmaf(w2,x2,a0); if(ow+1<OW){a1=fmaf(w0,x1,a1);a1=fmaf(w1,x2,a1);a1=fmaf(w2,x3,a1);}        
        }}}
        float* op=out+((((long long)n*OC+oc)*OD+od)*OH+oh)*OW+ow; op[0]=a0; if(ow+1<OW) op[1]=a1;
    }
}

__global__ void conv3d_generic_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ out,int N,int C,int D,int H,int W,int OC,int KD,int KH,int KW,int OD,int OH,int OW,int stride,int padding,int dilation,int has_bias){
 long long total=(long long)N*OC*OD*OH*OW; long long idx=(long long)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; t/=OC; int n=(int)t; float sum=has_bias?b[oc]:0.0f;
 for(int c=0;c<C;++c) for(int kz=0;kz<KD;++kz){int iz=od*stride-padding+kz*dilation; if((unsigned)iz>=(unsigned)D)continue; for(int ky=0;ky<KH;++ky){int iy=oh*stride-padding+ky*dilation; if((unsigned)iy>=(unsigned)H)continue; for(int kx=0;kx<KW;++kx){int ix=ow*stride-padding+kx*dilation; if((unsigned)ix>=(unsigned)W)continue; float xv=x[(((long long)n*C+c)*D+iz)*H*W+iy*W+ix]; float wv=w[(((long long)oc*C+c)*KD+kz)*KH*KW+ky*KW+kx]; sum=fmaf(wv,xv,sum);}}}
 out[idx]=sum;
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride64, int64_t padding64, int64_t dilation64, int64_t groups64) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA"); TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    x=x.contiguous(); weight=weight.contiguous();
    int N=(int)x.size(0), C=(int)x.size(1), D=(int)x.size(2), H=(int)x.size(3), W=(int)x.size(4);
    int OC=(int)weight.size(0), KD=(int)weight.size(2), KH=(int)weight.size(3), KW=(int)weight.size(4);
    int stride=(int)stride64, padding=(int)padding64, dilation=(int)dilation64;
    int OD=(D+2*padding-dilation*(KD-1)-1)/stride+1, OH=(H+2*padding-dilation*(KH-1)-1)/stride+1, OW=(W+2*padding-dilation*(KW-1)-1)/stride+1;
    auto out=torch::empty({N,OC,OD,OH,OW}, x.options());
    int has_bias=(bias.defined() && bias.numel()>0)?1:0; const float* bp=has_bias?bias.contiguous().data_ptr<float>():nullptr;
    const int threads=256;
    if(C==3 && KD==3 && KH==3 && KW==3 && stride==1 && padding==0 && dilation==1 && groups64==1 && x.scalar_type()==torch::kFloat32){
        int groupsW=(OW+3)>>2; long long total=(long long)N*OC*OD*OH*groupsW; int blocks=(int)((total+threads-1)/threads);
        if (OC*81 <= 8192) {
            cudaMemcpyToSymbol(c_weight, weight.data_ptr<float>(), (size_t)OC*81*sizeof(float), 0, cudaMemcpyDeviceToDevice);
            conv3d_c3k3s1p0_vec4_const_kernel<<<blocks,threads>>>(x.data_ptr<float>(), bp, out.data_ptr<float>(), N,D,H,W,OC,OD,OH,OW,has_bias);
        } else {
            conv3d_c3k3s1p0_vec4_kernel<<<blocks,threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp, out.data_ptr<float>(), N,D,H,W,OC,OD,OH,OW,has_bias);
        }
    } else {
        long long total=(long long)N*OC*OD*OH*OW; int blocks=(int)((total+threads-1)/threads);
        conv3d_generic_kernel<<<blocks,threads>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bp,out.data_ptr<float>(),N,C,D,H,W,OC,KD,KH,KW,OD,OH,OW,stride,padding,dilation,has_bias);
    }
    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, 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, kernel_size, kernel_size), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.stride = stride; self.padding = padding; self.dilation = dilation; self.groups = groups
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = 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, b, int(self.stride), int(self.padding), int(self.dilation), int(self.groups))