callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n_cpp_src='torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor gn_weight, torch::Tensor gn_bias, int64_t groups, double eps);'\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 conv_swish_sum_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ y,float* __restrict__ psum,float* __restrict__ psq,int N,int Cin,int Din,int Hin,int Win,int Cout,int Dout,int Hout,int Wout,int groups,int cpg,int spatial,int M,int bpg,int has_bias){int tid=threadIdx.x;int ng=blockIdx.y;int tile=blockIdx.x;int n=ng/groups;int g=ng-n*groups;int j=tile*blockDim.x+tid;float val=0.f,v2=0.f;if(j<M){int lc=j/spatial;int c=g*cpg+lc;int s=j-lc*spatial;int ow=s%Wout;int t=s/Wout;int oh=t%Hout;int od=t/Hout;val=has_bias?bias[c]:0.f;int nd=(od&1)?2:1,nh=(oh&1)?2:1,nw=(ow&1)?2:1;int kd0=(od&1)?0:1,kh0=(oh&1)?0:1,kw0=(ow&1)?0:1;#pragma unroll\nfor(int ci=0;ci<3;++ci){for(int ad=0;ad<nd;++ad){int kd=kd0+ad*2;int id=(od+1-kd)>>1;for(int ah=0;ah<nh;++ah){int kh=kh0+ah*2;int ih=(oh+1-kh)>>1;const float* xp=x+(((n*Cin+ci)*Din+id)*Hin+ih)*Win;for(int aw=0;aw<nw;++aw){int kw=kw0+aw*2;int iw=(ow+1-kw)>>1;val=fmaf(xp[iw],w[((((ci*Cout+c)*3+kd)*3+kh)*3+kw)],val);}}}}float sig=1.f/(1.f+expf(-val));val=val*sig;y[((n*Cout+c)*spatial)+s]=val;v2=val*val;}__shared__ float sm[512];float* ss=sm;float* sq=sm+256;ss[tid]=val;sq[tid]=v2;__syncthreads();for(int st=blockDim.x>>1;st>0;st>>=1){if(tid<st){ss[tid]+=ss[tid+st];sq[tid]+=sq[tid+st];}__syncthreads();}if(tid==0){int p=ng*bpg+tile;psum[p]=ss[0];psq[p]=sq[0];}}\n__global__ void reduce_stats_kernel(const float* __restrict__ psum,const float* __restrict__ psq,float* __restrict__ mean,float* __restrict__ invstd,int bpg,int M,float eps){int tid=threadIdx.x,ng=blockIdx.x;float s=0.f,q=0.f;for(int i=tid;i<bpg;i+=blockDim.x){s+=psum[ng*bpg+i];q+=psq[ng*bpg+i];}__shared__ float sm[512];float* ss=sm;float* sq=sm+256;ss[tid]=s;sq[tid]=q;__syncthreads();for(int st=blockDim.x>>1;st>0;st>>=1){if(tid<st){ss[tid]+=ss[tid+st];sq[tid]+=sq[tid+st];}__syncthreads();}if(tid==0){float m=ss[0]/(float)M;float var=sq[0]/(float)M-m*m;if(var<0)var=0;mean[ng]=m;invstd[ng]=rsqrtf(var+eps);}}\n__global__ void norm_hardswish_kernel(float* __restrict__ y,const float* __restrict__ mean,const float* __restrict__ invstd,const float* __restrict__ gamma,const float* __restrict__ beta,int N,int Cout,int groups,int cpg,int spatial,int M){int ng=blockIdx.y,tile=blockIdx.x,tid=threadIdx.x;int j=tile*blockDim.x+tid;if(j>=M)return;int n=ng/groups;int g=ng-n*groups;int lc=j/spatial;int c=g*cpg+lc;int s=j-lc*spatial;int idx=((n*Cout+c)*spatial)+s;float z=(y[idx]-mean[ng])*invstd[ng];z=fmaf(z,gamma[c],beta[c]);float r=fminf(fmaxf(z+3.f,0.f),6.f);y[idx]=z*r*(1.f/6.f);}\ntorch::Tensor fused_forward(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,torch::Tensor gn_weight,torch::Tensor gn_bias,int64_t groups,double eps){int N=x.size(0),Cin=x.size(1),Din=x.size(2),Hin=x.size(3),Win=x.size(4);int Cout=weight.size(1);int Dout=(Din-1)*2+1,Hout=(Hin-1)*2+1,Wout=(Win-1)*2+1;int spatial=Dout*Hout*Wout,cpg=Cout/(int)groups,M=cpg*spatial;const int threads=256;int bpg=(M+threads-1)/threads;auto y=torch::empty({N,Cout,Dout,Hout,Wout},x.options());auto opts=x.options().dtype(torch::kFloat32);auto psum=torch::empty({N*(int)groups,bpg},opts),psq=torch::empty({N*(int)groups,bpg},opts),mean=torch::empty({N*(int)groups},opts),invstd=torch::empty({N*(int)groups},opts);dim3 grid(bpg,N*(int)groups);auto stream=at::cuda::getCurrentCUDAStream();int has_bias=(bias.defined()&&bias.numel()>0);const float* bp=has_bias?bias.data_ptr<float>():nullptr;conv_swish_sum_kernel<<<grid,threads,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bp,y.data_ptr<float>(),psum.data_ptr<float>(),psq.data_ptr<float>(),N,Cin,Din,Hin,Win,Cout,Dout,Hout,Wout,(int)groups,cpg,spatial,M,bpg,has_bias);reduce_stats_kernel<<<N*(int)groups,threads,0,stream>>>(psum.data_ptr<float>(),psq.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),bpg,M,(float)eps);norm_hardswish_kernel<<<grid,threads,0,stream>>>(y.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),gn_weight.data_ptr<float>(),gn_bias.data_ptr<float>(),N,Cout,(int)groups,cpg,spatial,M);return y;}\n'''\n_ext=load_inline(name='fused_deconv_gn_hswish_ext_v4',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['fused_forward'],extra_cuda_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,in_channels,out_channels,kernel_size,stride,padding,groups,eps,bias=True):\n super().__init__();self.conv_transpose=nn.ConvTranspose3d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,bias=bias);self.group_norm=nn.GroupNorm(groups,out_channels,eps=eps);self.groups=groups;self.eps=eps\n def forward(self,x):\n b=self.conv_transpose.bias\n if b is None:b=torch.empty((0,),device=x.device,dtype=x.dtype)\n return _ext.fused_forward(x.contiguous(),self.conv_transpose.weight.contiguous(),b.contiguous(),self.group_norm.weight.contiguous(),self.group_norm.bias.contiguous(),self.groups,self.eps)\n"
}