callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor deconv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,\n int64_t stride, int64_t padding, int64_t dilation, bool has_bias);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n__global__ void deconv3d_s2p1d2_k3_kernel(const float* __restrict__ x,const float* __restrict__ wgt,const float* __restrict__ bias,float* __restrict__ out,int N,int Ci,int Co,int D,int H,int W,int Do,int Ho,int Wo,bool has_bias){\n long total=(long)N*Co*D*H*W; long idx=(long)blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total)return;\n int oxc=idx%W; long t=idx/W; int oyc=t%H; t/=H; int ozc=t%D; t/=D; int co=t%Co; int n=t/Co;\n float acc=has_bias?bias[co]:0.0f; long in_n_base=(long)n*Ci*D*H*W;\n #pragma unroll\n for(int kz=0;kz<3;++kz){int iz=ozc+1-kz; if((unsigned)iz>=(unsigned)D)continue;\n #pragma unroll\n for(int ky=0;ky<3;++ky){int iy=oyc+1-ky; if((unsigned)iy>=(unsigned)H)continue;\n #pragma unroll\n for(int kx=0;kx<3;++kx){int ix=oxc+1-kx; if((unsigned)ix>=(unsigned)W)continue;\n #pragma unroll 1\n for(int ci=0;ci<Ci;++ci){\n float xv=x[in_n_base+((long)ci*D+iz)*H*W+(long)iy*W+ix];\n float wv=wgt[((long)ci*Co+co)*27+kz*9+ky*3+kx];\n acc += xv*wv;\n }\n }\n }\n }\n long out_idx=(((long)n*Co+co)*Do+(2*ozc+1))*Ho*Wo+(long)(2*oyc+1)*Wo+(2*oxc+1); out[out_idx]=acc;\n}\n__global__ void fill_bias_kernel(float* out,const float* bias,long total,int Co,long spatial){long idx=(long)blockIdx.x*blockDim.x+threadIdx.x; if(idx<total){int co=(idx/spatial)%Co; out[idx]=bias[co];}}\n__global__ void deconv3d_generic_kernel(const float* x,const float* wgt,const float* bias,float* out,long total,int N,int Ci,int Co,int D,int H,int W,int K,int stride,int padding,int dilation,int Do,int Ho,int Wo,bool has_bias){long idx=(long)blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total)return; int ow=idx%Wo; long t=idx/Wo; int oh=t%Ho;t/=Ho; int od=t%Do;t/=Do; int co=t%Co; int n=t/Co; float acc=has_bias?bias[co]:0.0f; for(int ci=0;ci<Ci;++ci)for(int kz=0;kz<K;++kz){int rz=od+padding-kz*dilation;if(rz%stride)continue;int iz=rz/stride;if((unsigned)iz>=(unsigned)D)continue;for(int ky=0;ky<K;++ky){int ry=oh+padding-ky*dilation;if(ry%stride)continue;int iy=ry/stride;if((unsigned)iy>=(unsigned)H)continue;for(int kx=0;kx<K;++kx){int rx=ow+padding-kx*dilation;if(rx%stride)continue;int ix=rx/stride;if((unsigned)ix>=(unsigned)W)continue;acc += x[(long)n*Ci*D*H*W+((long)ci*D+iz)*H*W+(long)iy*W+ix]*wgt[(((long)ci*Co+co)*K+kz)*K*K+ky*K+kx];}}} out[idx]=acc;}\ntorch::Tensor deconv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,int64_t stride,int64_t padding,int64_t dilation,bool has_bias){x=x.contiguous();weight=weight.contiguous();int N=x.size(0),Ci=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4),Co=weight.size(1),K=weight.size(2);int Do=(D-1)*stride-2*padding+dilation*(K-1)+1,Ho=(H-1)*stride-2*padding+dilation*(K-1)+1,Wo=(W-1)*stride-2*padding+dilation*(K-1)+1;auto out=torch::empty({N,Co,Do,Ho,Wo},x.options());cudaStream_t stream=at::cuda::getCurrentCUDAStream();int threads=256;long out_total=(long)N*Co*Do*Ho*Wo;if(has_bias){long spatial=(long)Do*Ho*Wo;int blocks=(out_total+threads-1)/threads;fill_bias_kernel<<<blocks,threads,0,stream>>>(out.data_ptr<float>(),bias.data_ptr<float>(),out_total,Co,spatial);}else cudaMemsetAsync(out.data_ptr<float>(),0,out_total*sizeof(float),stream);if(K==3&&stride==2&&padding==1&&dilation==2){long comp=(long)N*Co*D*H*W;int blocks=(comp+threads-1)/threads;deconv3d_s2p1d2_k3_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,out.data_ptr<float>(),N,Ci,Co,D,H,W,Do,Ho,Wo,has_bias);}else{int blocks=(out_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,out.data_ptr<float>(),out_total,N,Ci,Co,D,H,W,K,stride,padding,dilation,Do,Ho,Wo,has_bias);}return out;}\n'''\n_ext=load_inline(name='deconv3d_order_ext',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['deconv3d_forward_cuda'],verbose=False,extra_cuda_cflags=['-O3'],extra_cflags=['-O3'])\nclass ModelNew(nn.Module):\n def __init__(self,in_channels:int,out_channels:int,kernel_size:int,stride:int=1,padding:int=0,dilation:int=1,bias:bool=False):\n super().__init__(); self.conv_transpose3d=nn.ConvTranspose3d(in_channels,out_channels,(kernel_size,kernel_size,kernel_size),stride=stride,padding=padding,dilation=dilation,bias=bias); self.stride=int(stride); self.padding=int(padding); self.dilation=int(dilation)\n def forward(self,x):\n b=self.conv_transpose3d.bias; has=b is not None\n if b is None: b=torch.empty((0,),device=x.device,dtype=x.dtype)\n return _ext.deconv3d_forward_cuda(x,self.conv_transpose3d.weight,b,self.stride,self.padding,self.dilation,has)\n"
}