call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\ncpp_src=r'''\n#include <torch/extension.h>\ntorch::Tensor bn_relu_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, double eps, bool training);\ntorch::Tensor avgpool_fc_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b);\n'''\ncuda_src=r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n__global__ void bn_relu_train_kernel(const float* __restrict__ x,const float* __restrict__ g,const float* __restrict__ be,float* __restrict__ y,int N,int C,int H,int W,float eps){int c=blockIdx.x,tid=threadIdx.x,M=N*H*W; extern __shared__ float sh[]; float sum=0,sum2=0; for(int idx=tid;idx<M;idx+=blockDim.x){int w=idx%W; int t=idx/W; int h=t%H; int n=t/H; float v=x[((n*C+c)*H+h)*W+w]; sum+=v; sum2+=v*v;} sh[tid]=sum; sh[blockDim.x+tid]=sum2; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s){sh[tid]+=sh[tid+s]; sh[blockDim.x+tid]+=sh[blockDim.x+tid+s];} __syncthreads();} float mean=sh[0]/M; float var=sh[blockDim.x]/M-mean*mean; float inv=rsqrtf(var+eps); float gg=g[c],bb=be[c]; for(int idx=tid;idx<M;idx+=blockDim.x){int w=idx%W; int t=idx/W; int h=t%H; int n=t/H; long long off=((long long)(n*C+c)*H+h)*W+w; float v=(x[off]-mean)*inv*gg+bb; y[off]=v>0?v:0;}}\n__global__ void bn_relu_eval_kernel(const float* x,const float* g,const float* be,const float* rm,const float* rv,float* y,int total,int C,int H,int W,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total)return; int c=(idx/(H*W))%C; float v=(x[idx]-rm[c])*rsqrtf(rv[c]+eps)*g[c]+be[c]; y[idx]=v>0?v:0;}\ntorch::Tensor bn_relu_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, double eps, bool training){int N=x.size(0),C=x.size(1),H=x.size(2),W=x.size(3); auto y=torch::empty_like(x); if(training) bn_relu_train_kernel<<<C,256,512*sizeof(float)>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),y.data_ptr<float>(),N,C,H,W,(float)eps); else {int total=N*C*H*W; bn_relu_eval_kernel<<<(total+255)/256,256>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),y.data_ptr<float>(),total,C,H,W,(float)eps);} return y;}\n__global__ void avg_kernel(const float* x,float* a,int N,int C,int H,int W){int idx=blockIdx.x,c=idx%C,n=idx/C,tid=threadIdx.x; extern __shared__ float sh[]; float s=0; int HW=H*W; const float* xp=x+((long long)n*C+c)*HW; for(int i=tid;i<HW;i+=blockDim.x)s+=xp[i]; sh[tid]=s; __syncthreads(); for(int off=32;off>0;off>>=1){if(tid<off)sh[tid]+=sh[tid+off]; __syncthreads();} if(tid==0)a[idx]=sh[0]/HW;}\n__global__ void fc_kernel(const float* a,const float* w,const float* b,float* y,int N,int C,int K){int k=blockIdx.x,n=blockIdx.y,tid=threadIdx.x; __shared__ float sh[256]; float s=0; const float* an=a+(long long)n*C; const float* wk=w+(long long)k*C; for(int c=tid;c<C;c+=256)s+=an[c]*wk[c]; sh[tid]=s; __syncthreads(); for(int off=128;off>0;off>>=1){if(tid<off)sh[tid]+=sh[tid+off]; __syncthreads();} if(tid==0)y[n*K+k]=sh[0]+b[k];}\ntorch::Tensor avgpool_fc_forward(torch::Tensor x,torch::Tensor w,torch::Tensor b){int N=x.size(0),C=x.size(1),H=x.size(2),W=x.size(3),K=w.size(0); auto a=torch::empty({N,C},x.options()); auto y=torch::empty({N,K},x.options()); avg_kernel<<<N*C,64,64*sizeof(float)>>>(x.data_ptr<float>(),a.data_ptr<float>(),N,C,H,W); fc_kernel<<<dim3(K,N),256>>>(a.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),y.data_ptr<float>(),N,C,K); return y;}\n'''\n_ext=load_inline(name='mobilenet_hybrid_ext_v1', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['bn_relu_forward','avgpool_fc_forward'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,num_classes=1000,input_channels=3,alpha=1.0):\n super().__init__()\n def conv_bn(inp,oup,stride): return nn.Sequential(nn.Conv2d(inp,oup,3,stride,1,bias=False),nn.BatchNorm2d(oup),nn.ReLU(inplace=True))\n def conv_dw(inp,oup,stride): return nn.Sequential(nn.Conv2d(inp,inp,3,stride,1,groups=inp,bias=False),nn.BatchNorm2d(inp),nn.ReLU(inplace=True),nn.Conv2d(inp,oup,1,1,0,bias=False),nn.BatchNorm2d(oup),nn.ReLU(inplace=True))\n self.model=nn.Sequential(conv_bn(input_channels,int(32*alpha),2),conv_dw(int(32*alpha),int(64*alpha),1),conv_dw(int(64*alpha),int(128*alpha),2),conv_dw(int(128*alpha),int(128*alpha),1),conv_dw(int(128*alpha),int(256*alpha),2),conv_dw(int(256*alpha),int(256*alpha),1),conv_dw(int(256*alpha),int(512*alpha),2),conv_dw(int(512*alpha),int(512*alpha),1),conv_dw(int(512*alpha),int(512*alpha),1),conv_dw(int(512*alpha),int(512*alpha),1),conv_dw(int(512*alpha),int(512*alpha),1),conv_dw(int(512*alpha),int(512*alpha),1),conv_dw(int(512*alpha),int(1024*alpha),2),conv_dw(int(1024*alpha),int(1024*alpha),1),nn.AvgPool2d(7))\n self.fc=nn.Linear(int(1024*alpha),num_classes)\n def cbn(self,x,bn): return _ext.bn_relu_forward(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,bn.eps,self.training)\n def forward(self,x):\n for i,l in enumerate(list(self.model.children())[:-1]):\n if i==0: x=l(x)\n elif i>=13:\n x=self.cbn(l[0](x),l[1]); x=self.cbn(l[3](x),l[4])\n else: x=l(x)\n return _ext.avgpool_fc_forward(x,self.fc.weight,self.fc.bias)\n"
}