call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor bn_relu2d(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps, bool training);\ntorch::Tensor tail_forward(torch::Tensor x,\n torch::Tensor mb1_se1_w, torch::Tensor mb1_se2_w, torch::Tensor mb1_out_w, torch::Tensor mb1_bn_w, torch::Tensor mb1_bn_b, torch::Tensor mb1_bn_m, torch::Tensor mb1_bn_v, double mb1_eps,\n torch::Tensor b2_exp_w, torch::Tensor b2_bn1_w, torch::Tensor b2_bn1_b, torch::Tensor b2_bn1_m, torch::Tensor b2_bn1_v, double b2_bn1_eps, torch::Tensor b2_dw_w, torch::Tensor b2_bn2_w, torch::Tensor b2_bn2_b, torch::Tensor b2_bn2_m, torch::Tensor b2_bn2_v, double b2_bn2_eps, torch::Tensor b2_se1_w, torch::Tensor b2_se2_w, torch::Tensor b2_out_w, torch::Tensor b2_bn3_w, torch::Tensor b2_bn3_b, torch::Tensor b2_bn3_m, torch::Tensor b2_bn3_v, double b2_bn3_eps,\n torch::Tensor b3_exp_w, torch::Tensor b3_bn1_w, torch::Tensor b3_bn1_b, torch::Tensor b3_bn1_m, torch::Tensor b3_bn1_v, double b3_bn1_eps, torch::Tensor b3_dw_w, torch::Tensor b3_bn2_w, torch::Tensor b3_bn2_b, torch::Tensor b3_bn2_m, torch::Tensor b3_bn2_v, double b3_bn2_eps, torch::Tensor b3_se1_w, torch::Tensor b3_se2_w, torch::Tensor b3_out_w, torch::Tensor b3_bn3_w, torch::Tensor b3_bn3_b, torch::Tensor b3_bn3_m, torch::Tensor b3_bn3_v, double b3_bn3_eps,\n torch::Tensor b4_exp_w, torch::Tensor b4_bn1_w, torch::Tensor b4_bn1_b, torch::Tensor b4_bn1_m, torch::Tensor b4_bn1_v, double b4_bn1_eps, torch::Tensor b4_dw_w, torch::Tensor b4_bn2_w, torch::Tensor b4_bn2_b, torch::Tensor b4_bn2_m, torch::Tensor b4_bn2_v, double b4_bn2_eps, torch::Tensor b4_se1_w, torch::Tensor b4_se2_w, torch::Tensor b4_out_w, torch::Tensor b4_bn3_w, torch::Tensor b4_bn3_b, torch::Tensor b4_bn3_m, torch::Tensor b4_bn3_v, double b4_bn3_eps,\n torch::Tensor b5_exp_w, torch::Tensor b5_bn1_w, torch::Tensor b5_bn1_b, torch::Tensor b5_bn1_m, torch::Tensor b5_bn1_v, double b5_bn1_eps, torch::Tensor b5_dw_w, torch::Tensor b5_bn2_w, torch::Tensor b5_bn2_b, torch::Tensor b5_bn2_m, torch::Tensor b5_bn2_v, double b5_bn2_eps, torch::Tensor b5_se1_w, torch::Tensor b5_se2_w, torch::Tensor b5_out_w, torch::Tensor b5_bn3_w, torch::Tensor b5_bn3_b, torch::Tensor b5_bn3_m, torch::Tensor b5_bn3_v, double b5_bn3_eps,\n torch::Tensor final_w, torch::Tensor final_bn_w, torch::Tensor final_bn_b, torch::Tensor final_bn_m, torch::Tensor final_bn_v, double final_eps, torch::Tensor fc_w, torch::Tensor fc_b, bool training);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#define CHECK_INPUT(x) TORCH_CHECK(x.is_cuda() && x.is_contiguous(), #x \" must be cuda contiguous\")\n\n__global__ void bnr_eval_k(const float* x,float* y,int N,int C,int HW,const float* g,const float* be,const float* m,const float* v,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=N)return; int c=(idx/HW)%C; float z=(x[idx]-m[c])*(g[c]*rsqrtf(v[c]+eps))+be[c]; y[idx]=z>0?z:0;}\n__global__ void bnr_train_k(const float* x,float* y,int B,int C,int HW,const float* g,const float* be,float eps){int c=blockIdx.x; __shared__ float s[256], q[256]; float a=0,b=0; for(int n=0;n<B;n++){const float* p=x+(n*C+c)*HW; for(int i=threadIdx.x;i<HW;i+=blockDim.x){float z=p[i]; a+=z; b+=z*z;}} s[threadIdx.x]=a; q[threadIdx.x]=b; __syncthreads(); for(int st=128;st;st>>=1){if(threadIdx.x<st){s[threadIdx.x]+=s[threadIdx.x+st]; q[threadIdx.x]+=q[threadIdx.x+st];} __syncthreads();} float cnt=(float)(B*HW); float mean=s[0]/cnt; float var=fmaxf(q[0]/cnt-mean*mean,0.0f); float sc=g[c]*rsqrtf(var+eps); for(int n=0;n<B;n++){float* po=y+(n*C+c)*HW; const float* pi=x+(n*C+c)*HW; for(int i=threadIdx.x;i<HW;i+=blockDim.x){float z=(pi[i]-mean)*sc+be[c]; po[i]=z>0?z:0;}}}\ntorch::Tensor bn_relu2d(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps, bool training){CHECK_INPUT(x); auto y=torch::empty_like(x); int B=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); if(training) bnr_train_k<<<C,256>>>(x.data_ptr<float>(),y.data_ptr<float>(),B,C,HW,weight.data_ptr<float>(),bias.data_ptr<float>(),(float)eps); else {int N=x.numel(); bnr_eval_k<<<(N+255)/256,256>>>(x.data_ptr<float>(),y.data_ptr<float>(),N,C,HW,weight.data_ptr<float>(),bias.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),(float)eps);} return y;}\n\n__global__ void avg_k(const float* x,float* y,int B,int C,int HW){int bc=blockIdx.x,b=bc/C,c=bc-b*C; __shared__ float s[256]; float a=0; const float* p=x+(b*C+c)*HW; for(int i=threadIdx.x;i<HW;i+=256)a+=p[i]; s[threadIdx.x]=a; __syncthreads(); for(int st=128;st;st>>=1){if(threadIdx.x<st)s[threadIdx.x]+=s[threadIdx.x+st]; __syncthreads();} if(!threadIdx.x)y[bc]=s[0]/(float)HW;}\n__global__ void mv_k(const float* in,const float* w,const float* bias,float* out,int B,int IN,int OUT,const float* g,const float* be,const float* m,const float* va,float eps,int use_bn,int use_bias,int act){int bo=blockIdx.x,b=bo/OUT,o=bo-b*OUT; __shared__ float s[256]; float a=0; for(int i=threadIdx.x;i<IN;i+=256)a+=in[b*IN+i]*w[o*IN+i]; s[threadIdx.x]=a; __syncthreads(); for(int st=128;st;st>>=1){if(threadIdx.x<st)s[threadIdx.x]+=s[threadIdx.x+st]; __syncthreads();} if(!threadIdx.x){float z=s[0]; if(use_bias)z+=bias[o]; if(use_bn)z=(z-m[o])*(g[o]*rsqrtf(va[o]+eps))+be[o]; if(act==1)z=z>0?z:0; else if(act==2)z=1.f/(1.f+expf(-z)); out[b*OUT+o]=z;}}\n__global__ void mv_bn_train_k(const float* in,const float* w,float* out,int IN,int OUT,const float* g,const float* be,float eps,int act){int o=blockIdx.x; __shared__ float s0[256],s1[256]; float a0=0,a1=0; for(int i=threadIdx.x;i<IN;i+=256){float ww=w[o*IN+i]; a0+=in[i]*ww; a1+=in[IN+i]*ww;} s0[threadIdx.x]=a0; s1[threadIdx.x]=a1; __syncthreads(); for(int st=128;st;st>>=1){if(threadIdx.x<st){s0[threadIdx.x]+=s0[threadIdx.x+st]; s1[threadIdx.x]+=s1[threadIdx.x+st];} __syncthreads();} if(!threadIdx.x){float v0=s0[0],v1=s1[0],mean=0.5f*(v0+v1),var=0.5f*((v0-mean)*(v0-mean)+(v1-mean)*(v1-mean)); float sc=g[o]*rsqrtf(var+eps); v0=(v0-mean)*sc+be[o]; v1=(v1-mean)*sc+be[o]; if(act==1){v0=v0>0?v0:0; v1=v1>0?v1:0;} out[o]=v0; out[OUT+o]=v1;}}\n__global__ void dw_eval_k(const float* in,const float* w,float* out,int N,int C,const float* g,const float* be,const float* m,const float* va,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx<N){int c=idx%C; float z=in[idx]*w[c*9+4]; z=(z-m[c])*(g[c]*rsqrtf(va[c]+eps))+be[c]; out[idx]=z>0?z:0;}}\n__global__ void dw_train_k(const float* in,const float* w,float* out,int C,const float* g,const float* be,float eps){int c=blockIdx.x*blockDim.x+threadIdx.x; if(c<C){float v0=in[c]*w[c*9+4],v1=in[C+c]*w[c*9+4]; float mean=.5f*(v0+v1),var=.5f*((v0-mean)*(v0-mean)+(v1-mean)*(v1-mean)); float sc=g[c]*rsqrtf(var+eps); v0=(v0-mean)*sc+be[c]; v1=(v1-mean)*sc+be[c]; out[c]=v0>0?v0:0; out[C+c]=v1>0?v1:0;}}\nstatic torch::Tensor matvec(torch::Tensor inp,torch::Tensor w,const float* bias,int use_bias,torch::Tensor g,torch::Tensor be,torch::Tensor m,torch::Tensor va,double eps,int use_bn,int act,bool tr){int B=inp.size(0),IN=inp.size(1),OUT=w.size(0); auto out=torch::empty({B,OUT},inp.options()); if(use_bn&&tr&&B==2) mv_bn_train_k<<<OUT,256>>>(inp.data_ptr<float>(),w.data_ptr<float>(),out.data_ptr<float>(),IN,OUT,g.data_ptr<float>(),be.data_ptr<float>(),(float)eps,act); else mv_k<<<B*OUT,256>>>(inp.data_ptr<float>(),w.data_ptr<float>(),bias,out.data_ptr<float>(),B,IN,OUT,use_bn?g.data_ptr<float>():nullptr,use_bn?be.data_ptr<float>():nullptr,use_bn?m.data_ptr<float>():nullptr,use_bn?va.data_ptr<float>():nullptr,(float)eps,use_bn,use_bias,act); return out;}\nstatic torch::Tensor dw(torch::Tensor inp,torch::Tensor w,torch::Tensor g,torch::Tensor be,torch::Tensor m,torch::Tensor va,double eps,bool tr){int B=inp.size(0),C=inp.size(1); auto out=torch::empty_like(inp); if(tr&&B==2) dw_train_k<<<(C+255)/256,256>>>(inp.data_ptr<float>(),w.data_ptr<float>(),out.data_ptr<float>(),C,g.data_ptr<float>(),be.data_ptr<float>(),(float)eps); else dw_eval_k<<<(B*C+255)/256,256>>>(inp.data_ptr<float>(),w.data_ptr<float>(),out.data_ptr<float>(),B*C,C,g.data_ptr<float>(),be.data_ptr<float>(),m.data_ptr<float>(),va.data_ptr<float>(),(float)eps); return out;}\nstatic torch::Tensor block(torch::Tensor cur,torch::Tensor ew,torch::Tensor b1w,torch::Tensor b1b,torch::Tensor b1m,torch::Tensor b1v,double e1,torch::Tensor dww,torch::Tensor b2w,torch::Tensor b2b,torch::Tensor b2m,torch::Tensor b2v,double e2,torch::Tensor s1w,torch::Tensor s2w,torch::Tensor ow,torch::Tensor b3w,torch::Tensor b3b,torch::Tensor b3m,torch::Tensor b3v,double e3,bool tr){auto a=matvec(cur,ew,nullptr,0,b1w,b1b,b1m,b1v,e1,1,1,tr); auto d=dw(a,dww,b2w,b2b,b2m,b2v,e2,tr); auto s1=matvec(d,s1w,nullptr,0,b1w,b1b,b1m,b1v,0,0,1,tr); auto s2=matvec(s1,s2w,nullptr,0,b1w,b1b,b1m,b1v,0,0,2,tr); return matvec(s2,ow,nullptr,0,b3w,b3b,b3m,b3v,e3,1,0,tr);}\n\ntorch::Tensor tail_forward(torch::Tensor x,\n torch::Tensor mb1s1,torch::Tensor mb1s2,torch::Tensor mb1o,torch::Tensor mb1bw,torch::Tensor mb1bb,torch::Tensor mb1bm,torch::Tensor mb1bv,double mb1e,\n torch::Tensor b2ew,torch::Tensor b2b1w,torch::Tensor b2b1b,torch::Tensor b2b1m,torch::Tensor b2b1v,double b2e1,torch::Tensor b2dw,torch::Tensor b2b2w,torch::Tensor b2b2b,torch::Tensor b2b2m,torch::Tensor b2b2v,double b2e2,torch::Tensor b2s1,torch::Tensor b2s2,torch::Tensor b2ow,torch::Tensor b2b3w,torch::Tensor b2b3b,torch::Tensor b2b3m,torch::Tensor b2b3v,double b2e3,\n torch::Tensor b3ew,torch::Tensor b3b1w,torch::Tensor b3b1b,torch::Tensor b3b1m,torch::Tensor b3b1v,double b3e1,torch::Tensor b3dw,torch::Tensor b3b2w,torch::Tensor b3b2b,torch::Tensor b3b2m,torch::Tensor b3b2v,double b3e2,torch::Tensor b3s1,torch::Tensor b3s2,torch::Tensor b3ow,torch::Tensor b3b3w,torch::Tensor b3b3b,torch::Tensor b3b3m,torch::Tensor b3b3v,double b3e3,\n torch::Tensor b4ew,torch::Tensor b4b1w,torch::Tensor b4b1b,torch::Tensor b4b1m,torch::Tensor b4b1v,double b4e1,torch::Tensor b4dw,torch::Tensor b4b2w,torch::Tensor b4b2b,torch::Tensor b4b2m,torch::Tensor b4b2v,double b4e2,torch::Tensor b4s1,torch::Tensor b4s2,torch::Tensor b4ow,torch::Tensor b4b3w,torch::Tensor b4b3b,torch::Tensor b4b3m,torch::Tensor b4b3v,double b4e3,\n torch::Tensor b5ew,torch::Tensor b5b1w,torch::Tensor b5b1b,torch::Tensor b5b1m,torch::Tensor b5b1v,double b5e1,torch::Tensor b5dw,torch::Tensor b5b2w,torch::Tensor b5b2b,torch::Tensor b5b2m,torch::Tensor b5b2v,double b5e2,torch::Tensor b5s1,torch::Tensor b5s2,torch::Tensor b5ow,torch::Tensor b5b3w,torch::Tensor b5b3b,torch::Tensor b5b3m,torch::Tensor b5b3v,double b5e3,\n torch::Tensor fw,torch::Tensor fbw,torch::Tensor fbb,torch::Tensor fbm,torch::Tensor fbv,double fe,torch::Tensor fcw,torch::Tensor fcb,bool tr){CHECK_INPUT(x); int B=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); auto p=torch::empty({B,C},x.options()); avg_k<<<B*C,256>>>(x.data_ptr<float>(),p.data_ptr<float>(),B,C,HW); auto s1=matvec(p,mb1s1,nullptr,0,mb1bw,mb1bb,mb1bm,mb1bv,0,0,1,tr); auto s2=matvec(s1,mb1s2,nullptr,0,mb1bw,mb1bb,mb1bm,mb1bv,0,0,2,tr); auto cur=matvec(s2,mb1o,nullptr,0,mb1bw,mb1bb,mb1bm,mb1bv,mb1e,1,0,tr); cur=block(cur,b2ew,b2b1w,b2b1b,b2b1m,b2b1v,b2e1,b2dw,b2b2w,b2b2b,b2b2m,b2b2v,b2e2,b2s1,b2s2,b2ow,b2b3w,b2b3b,b2b3m,b2b3v,b2e3,tr); cur=block(cur,b3ew,b3b1w,b3b1b,b3b1m,b3b1v,b3e1,b3dw,b3b2w,b3b2b,b3b2m,b3b2v,b3e2,b3s1,b3s2,b3ow,b3b3w,b3b3b,b3b3m,b3b3v,b3e3,tr); cur=block(cur,b4ew,b4b1w,b4b1b,b4b1m,b4b1v,b4e1,b4dw,b4b2w,b4b2b,b4b2m,b4b2v,b4e2,b4s1,b4s2,b4ow,b4b3w,b4b3b,b4b3m,b4b3v,b4e3,tr); cur=block(cur,b5ew,b5b1w,b5b1b,b5b1m,b5b1v,b5e1,b5dw,b5b2w,b5b2b,b5b2m,b5b2v,b5e2,b5s1,b5s2,b5ow,b5b3w,b5b3b,b5b3m,b5b3v,b5e3,tr); cur=matvec(cur,fw,nullptr,0,fbw,fbb,fbm,fbv,fe,1,1,tr); return matvec(cur,fcw,fcb.data_ptr<float>(),1,fbw,fbb,fbm,fbv,0,0,0,tr);}\n'''\n_ext = load_inline(name='effb2_tail_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['bn_relu2d','tail_forward'], verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])\n\nclass ModelNew(nn.Module):\n def __init__(self, num_classes=1000):\n super().__init__(); self.conv1=nn.Conv2d(3,32,3,2,1,bias=False); self.bn1=nn.BatchNorm2d(32); self.relu=nn.ReLU(inplace=True)\n self.mbconv1=self._make_mbconv_block(32,96,1,3); self.mbconv2=self._make_mbconv_block(96,144,2,6); self.mbconv3=self._make_mbconv_block(144,192,2,6); self.mbconv4=self._make_mbconv_block(192,288,2,6); self.mbconv5=self._make_mbconv_block(288,384,1,6)\n self.conv_final=nn.Conv2d(384,1408,1,1,0,bias=False); self.bn_final=nn.BatchNorm2d(1408); self.avgpool=nn.AdaptiveAvgPool2d((1,1)); self.fc=nn.Linear(1408,num_classes)\n def _make_mbconv_block(self,in_channels,out_channels,stride,expand_ratio):\n layers=[]; ec=in_channels*expand_ratio\n if expand_ratio!=1: layers += [nn.Conv2d(in_channels,ec,1,1,0,bias=False), nn.BatchNorm2d(ec), nn.ReLU(inplace=True)]\n layers += [nn.Conv2d(ec,ec,3,stride,1,groups=ec,bias=False), nn.BatchNorm2d(ec), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1,1)), nn.Conv2d(ec,ec//4,1,bias=False), nn.ReLU(inplace=True), nn.Conv2d(ec//4,ec,1,bias=False), nn.Sigmoid(), nn.Conv2d(ec,out_channels,1,bias=False), nn.BatchNorm2d(out_channels)]\n return nn.Sequential(*layers)\n def _bnr(self,x,bn): return _ext.bn_relu2d(x.contiguous(),bn.weight,bn.bias,bn.running_mean,bn.running_var,bn.eps,self.training)\n def forward(self,x):\n x=F.conv2d(x,self.conv1.weight,None,stride=2,padding=1); x=self._bnr(x,self.bn1)\n x=F.conv2d(x,self.mbconv1[0].weight,None); x=self._bnr(x,self.mbconv1[1])\n x=F.conv2d(x,self.mbconv1[3].weight,None,padding=1,groups=96); x=self._bnr(x,self.mbconv1[4])\n b1,b2,b3,b4,b5=self.mbconv1,self.mbconv2,self.mbconv3,self.mbconv4,self.mbconv5\n return _ext.tail_forward(x.contiguous(), b1[7].weight,b1[9].weight,b1[11].weight,b1[12].weight,b1[12].bias,b1[12].running_mean,b1[12].running_var,b1[12].eps,\n b2[0].weight,b2[1].weight,b2[1].bias,b2[1].running_mean,b2[1].running_var,b2[1].eps,b2[3].weight,b2[4].weight,b2[4].bias,b2[4].running_mean,b2[4].running_var,b2[4].eps,b2[7].weight,b2[9].weight,b2[11].weight,b2[12].weight,b2[12].bias,b2[12].running_mean,b2[12].running_var,b2[12].eps,\n b3[0].weight,b3[1].weight,b3[1].bias,b3[1].running_mean,b3[1].running_var,b3[1].eps,b3[3].weight,b3[4].weight,b3[4].bias,b3[4].running_mean,b3[4].running_var,b3[4].eps,b3[7].weight,b3[9].weight,b3[11].weight,b3[12].weight,b3[12].bias,b3[12].running_mean,b3[12].running_var,b3[12].eps,\n b4[0].weight,b4[1].weight,b4[1].bias,b4[1].running_mean,b4[1].running_var,b4[1].eps,b4[3].weight,b4[4].weight,b4[4].bias,b4[4].running_mean,b4[4].running_var,b4[4].eps,b4[7].weight,b4[9].weight,b4[11].weight,b4[12].weight,b4[12].bias,b4[12].running_mean,b4[12].running_var,b4[12].eps,\n b5[0].weight,b5[1].weight,b5[1].bias,b5[1].running_mean,b5[1].running_var,b5[1].eps,b5[3].weight,b5[4].weight,b5[4].bias,b5[4].running_mean,b5[4].running_var,b5[4].eps,b5[7].weight,b5[9].weight,b5[11].weight,b5[12].weight,b5[12].bias,b5[12].running_mean,b5[12].running_var,b5[12].eps,\n self.conv_final.weight,self.bn_final.weight,self.bn_final.bias,self.bn_final.running_mean,self.bn_final.running_var,self.bn_final.eps,self.fc.weight,self.fc.bias,self.training)\n"
}