callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n_cpp=r'''\n#include <torch/extension.h>\ntorch::Tensor shufflenet_forward(torch::Tensor x,torch::Tensor w1,torch::Tensor b1w,torch::Tensor b1b,torch::Tensor b1m,torch::Tensor b1v,torch::Tensor w2,torch::Tensor b2w,torch::Tensor b2b,torch::Tensor b2m,torch::Tensor b2v,torch::Tensor w3,torch::Tensor b3w,torch::Tensor b3b,torch::Tensor b3m,torch::Tensor b3v,torch::Tensor ws,torch::Tensor bsw,torch::Tensor bsb,torch::Tensor bsm,torch::Tensor bsv,int groups);\n'''\n_cuda=r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <c10/cuda/CUDAGuard.h>\n#include <cublas_v2.h>\n#define CUBLAS_CHECK(e) TORCH_CHECK((e)==CUBLAS_STATUS_SUCCESS,\"cublas\")\n__global__ void stats_kernel(const float* __restrict__ x,float* __restrict__ mean,float* __restrict__ var,float* __restrict__ running_m,float* __restrict__ running_v,int C,int M,int HW){int c=blockIdx.x; extern __shared__ double sh[]; double* ss=sh; double* qq=sh+blockDim.x; double s=0.0,q=0.0; for(int i=threadIdx.x;i<M;i+=blockDim.x){int n=i/HW; int p=i-n*HW; float v=x[(n*C+c)*HW+p]; s+=v; q+=(double)v*v;} ss[threadIdx.x]=s; qq[threadIdx.x]=q; __syncthreads(); for(int off=blockDim.x/2;off>0;off>>=1){ if(threadIdx.x<off){ss[threadIdx.x]+=ss[threadIdx.x+off]; qq[threadIdx.x]+=qq[threadIdx.x+off];} __syncthreads(); } if(threadIdx.x==0){double mu=ss[0]/(double)M; double va=qq[0]/(double)M-mu*mu; if(va<0) va=0; mean[c]=(float)mu; var[c]=(float)va; if(running_m){running_m[c]=running_m[c]*0.9f+(float)mu*0.1f; running_v[c]=running_v[c]*0.9f+(float)(va*((double)M/(double)(M-1)))*0.1f;}}}\n__global__ void bn_apply_kernel(float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ var,int total,int C,int HW,int relu){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total)return; int c=(idx/HW)%C; float v=(y[idx]-mean[c])*rsqrtf(var[c]+1e-5f)*gamma[c]+beta[c]; if(relu&&v<0)v=0; y[idx]=v;}\n__global__ void dw_raw_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ w,int N,int C,int H,int W){int idx=blockIdx.x*blockDim.x+threadIdx.x,total=N*C*H*W;if(idx>=total)return;int HW=H*W,p=idx%HW,c=(idx/HW)%C,n=idx/(C*HW),h=p/W,ww=p-h*W;float sum=0;const float* wk=w+c*9;for(int kh=0;kh<3;kh++){int ih=h+kh-1;if((unsigned)ih<(unsigned)H)for(int kw=0;kw<3;kw++){int iw=ww+kw-1;if((unsigned)iw<(unsigned)W)sum=fmaf(x[(n*C+c)*HW+ih*W+iw],wk[kh*3+kw],sum);}}y[idx]=sum;}\n__global__ void shuffle_kernel(const float* __restrict__ x,float* __restrict__ y,int total,int C,int HW,int groups){int idx=blockIdx.x*blockDim.x+threadIdx.x;if(idx>=total)return;int p=idx%HW,c=(idx/HW)%C,n=idx/(C*HW);int cpg=C/groups,co=(c%cpg)*groups+c/cpg;y[(n*C+co)*HW+p]=x[idx];}\n__global__ void final_train_kernel(float* __restrict__ main,float* __restrict__ sc,const float* __restrict__ g3,const float* __restrict__ b3,const float* __restrict__ m3,const float* __restrict__ v3,const float* __restrict__ gs,const float* __restrict__ bs,const float* __restrict__ ms,const float* __restrict__ vs,int total,int C,int HW){int idx=blockIdx.x*blockDim.x+threadIdx.x;if(idx>=total)return;int c=(idx/HW)%C;float a=(main[idx]-m3[c])*rsqrtf(v3[c]+1e-5f)*g3[c]+b3[c]; if(a<0)a=0;float s=(sc[idx]-ms[c])*rsqrtf(vs[c]+1e-5f)*gs[c]+bs[c];sc[idx]=a+s;}\nstatic inline void gemm(cublasHandle_t h,int M,int N,int K,const float*A,const float*B,float*C){float al=1,be=0;CUBLAS_CHECK(cublasSgemm(h,CUBLAS_OP_N,CUBLAS_OP_N,N,M,K,&al,B,N,A,K,&be,C,N));}\ntorch::Tensor shufflenet_forward(torch::Tensor x,torch::Tensor w1,torch::Tensor b1w,torch::Tensor b1b,torch::Tensor b1m,torch::Tensor b1v,torch::Tensor w2,torch::Tensor b2w,torch::Tensor b2b,torch::Tensor b2m,torch::Tensor b2v,torch::Tensor w3,torch::Tensor b3w,torch::Tensor b3b,torch::Tensor b3m,torch::Tensor b3v,torch::Tensor ws,torch::Tensor bsw,torch::Tensor bsb,torch::Tensor bsm,torch::Tensor bsv,int groups){c10::cuda::CUDAGuard guard(x.device());int N=x.size(0),Cin=x.size(1),H=x.size(2),W=x.size(3),HW=H*W,mid=w1.size(0),Cout=w3.size(0),M=N*HW;auto opt=x.options();auto t1=torch::empty({N,mid,H,W},opt),t2=torch::empty({N,mid,H,W},opt),tsh=torch::empty({N,mid,H,W},opt),t3=torch::empty({N,Cout,H,W},opt),out=torch::empty({N,Cout,H,W},opt);auto m1=torch::empty({Cout},opt),v1=torch::empty({Cout},opt),m2=torch::empty({Cout},opt),v2=torch::empty({Cout},opt),m3=torch::empty({Cout},opt),v3=torch::empty({Cout},opt),ms=torch::empty({Cout},opt),vs=torch::empty({Cout},opt);cublasHandle_t h=at::cuda::getCurrentCUDABlasHandle();CUBLAS_CHECK(cublasSetStream(h,at::cuda::getCurrentCUDAStream()));CUBLAS_CHECK(cublasSetMathMode(h,CUBLAS_DEFAULT_MATH));const float*xp=x.data_ptr<float>();float*t1p=t1.data_ptr<float>();int inpg=Cin/groups,outpg=mid/groups;for(int n=0;n<N;n++)for(int g=0;g<groups;g++)gemm(h,outpg,HW,inpg,w1.data_ptr<float>()+(g*outpg)*inpg,xp+(n*Cin+g*inpg)*HW,t1p+(n*mid+g*outpg)*HW);int th=256,total1=N*mid*HW,totalo=N*Cout*HW;stats_kernel<<<mid,th,th*2*sizeof(double),at::cuda::getCurrentCUDAStream()>>>(t1p,m1.data_ptr<float>(),v1.data_ptr<float>(),b1m.data_ptr<float>(),b1v.data_ptr<float>(),mid,M,HW);bn_apply_kernel<<<(total1+th-1)/th,th,0,at::cuda::getCurrentCUDAStream()>>>(t1p,b1w.data_ptr<float>(),b1b.data_ptr<float>(),m1.data_ptr<float>(),v1.data_ptr<float>(),total1,mid,HW,1);dw_raw_kernel<<<(total1+th-1)/th,th,0,at::cuda::getCurrentCUDAStream()>>>(t1p,t2.data_ptr<float>(),w2.data_ptr<float>(),N,mid,H,W);stats_kernel<<<mid,th,th*2*sizeof(double),at::cuda::getCurrentCUDAStream()>>>(t2.data_ptr<float>(),m2.data_ptr<float>(),v2.data_ptr<float>(),b2m.data_ptr<float>(),b2v.data_ptr<float>(),mid,M,HW);bn_apply_kernel<<<(total1+th-1)/th,th,0,at::cuda::getCurrentCUDAStream()>>>(t2.data_ptr<float>(),b2w.data_ptr<float>(),b2b.data_ptr<float>(),m2.data_ptr<float>(),v2.data_ptr<float>(),total1,mid,HW,0);shuffle_kernel<<<(total1+th-1)/th,th,0,at::cuda::getCurrentCUDAStream()>>>(t2.data_ptr<float>(),tsh.data_ptr<float>(),total1,mid,HW,groups);int in3=mid/groups,out3=Cout/groups;for(int n=0;n<N;n++)for(int g=0;g<groups;g++)gemm(h,out3,HW,in3,w3.data_ptr<float>()+(g*out3)*in3,tsh.data_ptr<float>()+(n*mid+g*in3)*HW,t3.data_ptr<float>()+(n*Cout+g*out3)*HW);for(int n=0;n<N;n++)gemm(h,Cout,HW,Cin,ws.data_ptr<float>(),xp+n*Cin*HW,out.data_ptr<float>()+n*Cout*HW);stats_kernel<<<Cout,th,th*2*sizeof(double),at::cuda::getCurrentCUDAStream()>>>(t3.data_ptr<float>(),m3.data_ptr<float>(),v3.data_ptr<float>(),b3m.data_ptr<float>(),b3v.data_ptr<float>(),Cout,M,HW);stats_kernel<<<Cout,th,th*2*sizeof(double),at::cuda::getCurrentCUDAStream()>>>(out.data_ptr<float>(),ms.data_ptr<float>(),vs.data_ptr<float>(),bsm.data_ptr<float>(),bsv.data_ptr<float>(),Cout,M,HW);final_train_kernel<<<(totalo+th-1)/th,th,0,at::cuda::getCurrentCUDAStream()>>>(t3.data_ptr<float>(),out.data_ptr<float>(),b3w.data_ptr<float>(),b3b.data_ptr<float>(),m3.data_ptr<float>(),v3.data_ptr<float>(),bsw.data_ptr<float>(),bsb.data_ptr<float>(),ms.data_ptr<float>(),vs.data_ptr<float>(),totalo,Cout,HW);return out;}\n'''\n_ext=load_inline(name='shufflenet_train_ext_v1',cpp_sources=_cpp,cuda_sources=_cuda,functions=['shufflenet_forward'],verbose=False,extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'])\nclass ModelNew(nn.Module):\n def __init__(self,in_channels,out_channels,groups=3):\n super().__init__();mid_channels=out_channels//4;self.groups=groups\n self.conv1=nn.Conv2d(in_channels,mid_channels,1,1,0,groups=groups,bias=False);self.bn1=nn.BatchNorm2d(mid_channels)\n self.conv2=nn.Conv2d(mid_channels,mid_channels,3,1,1,groups=mid_channels,bias=False);self.bn2=nn.BatchNorm2d(mid_channels)\n self.conv3=nn.Conv2d(mid_channels,out_channels,1,1,0,groups=groups,bias=False);self.bn3=nn.BatchNorm2d(out_channels)\n self.shortcut=nn.Sequential(nn.Conv2d(in_channels,out_channels,1,1,0,bias=False),nn.BatchNorm2d(out_channels)) if in_channels!=out_channels else nn.Sequential()\n def forward(self,x):\n return _ext.shufflenet_forward(x.contiguous(),self.conv1.weight,self.bn1.weight,self.bn1.bias,self.bn1.running_mean,self.bn1.running_var,self.conv2.weight,self.bn2.weight,self.bn2.bias,self.bn2.running_mean,self.bn2.running_var,self.conv3.weight,self.bn3.weight,self.bn3.bias,self.bn3.running_mean,self.bn3.running_var,self.shortcut[0].weight,self.shortcut[1].weight,self.shortcut[1].bias,self.shortcut[1].running_mean,self.shortcut[1].running_var,self.groups)\n"
}