callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\ntorch::Tensor cvit_forward(torch::Tensor x, std::vector<torch::Tensor> p);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <cmath>\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be CUDA\")\n__inline__ __device__ float warp_reduce_sum(float val){for(int o=16;o>0;o>>=1) val += __shfl_down_sync(0xffffffff,val,o); return val;}\ntemplate<int BLOCK> __inline__ __device__ float block_reduce_sum(float val){\n static __shared__ float shared[32]; int lane=threadIdx.x&31, wid=threadIdx.x>>5;\n val=warp_reduce_sum(val); if(lane==0) shared[wid]=val; __syncthreads();\n val=(threadIdx.x < (BLOCK+31)/32)? shared[lane] : 0.0f; if(wid==0) val=warp_reduce_sum(val); __syncthreads(); return val;\n}\n__global__ void conv_flat_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int B,int C,int H,int W,int E,int Ksz,int OH,int OW){\n int idx=blockIdx.x*blockDim.x+threadIdx.x, total=B*E*OH*OW; if(idx>=total) return; int p=idx%(OH*OW); int e=(idx/(OH*OW))%E; int b=idx/(E*OH*OW); int ph=p/OW, pw=p-ph*OW; float sum=bias[e];\n for(int c=0;c<C;++c) for(int kh=0;kh<Ksz;++kh) for(int kw=0;kw<Ksz;++kw) sum=fmaf(x[((b*C+c)*H+ph*Ksz+kh)*W+pw*Ksz+kw], w[((e*C+c)*Ksz+kh)*Ksz+kw], sum);\n out[(b*E+e)*(OH*OW)+p]=sum;\n}\ntemplate<int BLOCK> __global__ void linear_kernel(const float* __restrict__ inp,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int M,int K,int N,int relu){\n int idx=blockIdx.x,n=idx%N,m=idx/N; if(m>=M) return; float acc=0.0f; const float* inrow=inp+m*K; const float* wrow=w+n*K; for(int k=threadIdx.x;k<K;k+=BLOCK) acc=fmaf(inrow[k],wrow[k],acc); acc=block_reduce_sum<BLOCK>(acc); if(threadIdx.x==0){acc+=bias[n]; if(relu&&acc<0.0f) acc=0.0f; out[m*N+n]=acc;}\n}\n__global__ void init_seq_kernel(const float* __restrict__ proj,const float* __restrict__ cls,float* __restrict__ seq,int B,int D){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=B*2*D) return; int d=idx%D; int s=(idx/D)&1; int b=idx/(2*D); seq[idx]=(s==0)?cls[d]:proj[b*D+d];}\n__global__ void attn_ctx_kernel(const float* __restrict__ qkv,float* __restrict__ ctx,int B,int D,int Hh,int Hd){\n int bid=blockIdx.x, qpos=bid&1, h=(bid>>1)%Hh, b=bid/(2*Hh), tid=threadIdx.x; int rowq=b*2+qpos,row0=b*2,row1=row0+1; float a0=0.0f,a1=0.0f; if(tid<Hd){int off=h*Hd+tid; float q=qkv[rowq*3*D+off]; a0=q*qkv[row0*3*D+D+off]; a1=q*qkv[row1*3*D+D+off];} for(int o=16;o>0;o>>=1){a0+=__shfl_down_sync(0xffffffff,a0,o); a1+=__shfl_down_sync(0xffffffff,a1,o);} float s0=__shfl_sync(0xffffffff,a0,0)/sqrtf((float)Hd); float s1=__shfl_sync(0xffffffff,a1,0)/sqrtf((float)Hd); float m=fmaxf(s0,s1); float e0=expf(s0-m), e1=expf(s1-m); float inv=1.0f/(e0+e1); if(tid<Hd){int off=h*Hd+tid; ctx[rowq*D+off]=(e0*inv)*qkv[row0*3*D+2*D+off]+(e1*inv)*qkv[row1*3*D+2*D+off];}\n}\ntemplate<int BLOCK> __global__ void residual_norm_kernel(const float* __restrict__ a,const float* __restrict__ b,float* __restrict__ out,const float* __restrict__ gamma,const float* __restrict__ beta,int M,int D){\n __shared__ float mean_s, rstd_s; int m=blockIdx.x; float sum=0.0f; for(int d=threadIdx.x;d<D;d+=BLOCK) sum += a[m*D+d]+b[m*D+d]; sum=block_reduce_sum<BLOCK>(sum); if(threadIdx.x==0) mean_s=sum/(float)D; __syncthreads(); float mean=mean_s; float ss=0.0f; for(int d=threadIdx.x;d<D;d+=BLOCK){float t=a[m*D+d]+b[m*D+d]-mean; ss+=t*t;} ss=block_reduce_sum<BLOCK>(ss); if(threadIdx.x==0) rstd_s=rsqrtf(ss/(float)D+1.0e-5f); __syncthreads(); float rstd=rstd_s; for(int d=threadIdx.x;d<D;d+=BLOCK){float val=(a[m*D+d]+b[m*D+d]-mean)*rstd; out[m*D+d]=val*gamma[d]+beta[d];}\n}\ntemplate<int BLOCK> __global__ void fc_kernel(const float* __restrict__ seq,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int B,int D,int N){int idx=blockIdx.x,n=idx%N,b=idx/N; float acc=0.0f; const float* inrow=seq+b*2*D; const float* wrow=w+n*D; for(int d=threadIdx.x;d<D;d+=BLOCK) acc=fmaf(inrow[d],wrow[d],acc); acc=block_reduce_sum<BLOCK>(acc); if(threadIdx.x==0) out[b*N+n]=acc+bias[n];}\n\ntorch::Tensor cvit_forward(torch::Tensor x,std::vector<torch::Tensor> p){\n CHECK_CUDA(x); int B=x.size(0),C=x.size(1),H=x.size(2),W=x.size(3); auto conv_w=p[0],conv_b=p[1],lp_w=p[2],lp_b=p[3],cls=p[4]; int E=conv_w.size(0),Ksz=conv_w.size(2),OH=H/Ksz,OW=W/Ksz,Pn=OH*OW,D=lp_w.size(0); auto opts=x.options(); cudaStream_t stream=at::cuda::getCurrentCUDAStream(); auto conv=torch::empty({B,E*Pn},opts); conv_flat_kernel<<<(B*E*Pn+255)/256,256,0,stream>>>(x.data_ptr<float>(),conv_w.data_ptr<float>(),conv_b.data_ptr<float>(),conv.data_ptr<float>(),B,C,H,W,E,Ksz,OH,OW); auto proj=torch::empty({B,D},opts); linear_kernel<256><<<B*D,256,0,stream>>>(conv.data_ptr<float>(),lp_w.data_ptr<float>(),lp_b.data_ptr<float>(),proj.data_ptr<float>(),B,E*Pn,D,0); auto bufA=torch::empty({B*2,D},opts); init_seq_kernel<<<(B*2*D+255)/256,256,0,stream>>>(proj.data_ptr<float>(),cls.data_ptr<float>(),bufA.data_ptr<float>(),B,D);\n int idx=5,num_layers=(p.size()-7)/12,M=B*2,heads=4,Hd=D/heads; auto qkv=torch::empty({M,3*D},opts),ctx=torch::empty({M,D},opts),tmpD=torch::empty({M,D},opts),bufB=torch::empty({M,D},opts); for(int l=0;l<num_layers;++l){auto inpw=p[idx++],inpb=p[idx++],outw=p[idx++],outb=p[idx++],lin1w=p[idx++],lin1b=p[idx++],lin2w=p[idx++],lin2b=p[idx++],n1w=p[idx++],n1b=p[idx++],n2w=p[idx++],n2b=p[idx++]; linear_kernel<256><<<M*3*D,256,0,stream>>>(bufA.data_ptr<float>(),inpw.data_ptr<float>(),inpb.data_ptr<float>(),qkv.data_ptr<float>(),M,D,3*D,0); attn_ctx_kernel<<<B*heads*2,32,0,stream>>>(qkv.data_ptr<float>(),ctx.data_ptr<float>(),B,D,heads,Hd); linear_kernel<256><<<M*D,256,0,stream>>>(ctx.data_ptr<float>(),outw.data_ptr<float>(),outb.data_ptr<float>(),tmpD.data_ptr<float>(),M,D,D,0); residual_norm_kernel<256><<<M,256,0,stream>>>(bufA.data_ptr<float>(),tmpD.data_ptr<float>(),bufB.data_ptr<float>(),n1w.data_ptr<float>(),n1b.data_ptr<float>(),M,D); int Hidden=lin1w.size(0); auto hid=torch::empty({M,Hidden},opts); linear_kernel<256><<<M*Hidden,256,0,stream>>>(bufB.data_ptr<float>(),lin1w.data_ptr<float>(),lin1b.data_ptr<float>(),hid.data_ptr<float>(),M,D,Hidden,1); linear_kernel<256><<<M*D,256,0,stream>>>(hid.data_ptr<float>(),lin2w.data_ptr<float>(),lin2b.data_ptr<float>(),tmpD.data_ptr<float>(),M,Hidden,D,0); residual_norm_kernel<256><<<M,256,0,stream>>>(bufB.data_ptr<float>(),tmpD.data_ptr<float>(),bufA.data_ptr<float>(),n2w.data_ptr<float>(),n2b.data_ptr<float>(),M,D);} auto fc_w=p[idx++],fc_b=p[idx++]; int Ncls=fc_w.size(0); auto out=torch::empty({B,Ncls},opts); fc_kernel<256><<<B*Ncls,256,0,stream>>>(bufA.data_ptr<float>(),fc_w.data_ptr<float>(),fc_b.data_ptr<float>(),out.data_ptr<float>(),B,D,Ncls); return out;\n}\n'''\n_ext = load_inline(name='cvit_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['cvit_forward'], verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])\nclass ModelNew(nn.Module):\n def __init__(self, num_classes, embed_dim=512, num_heads=8, num_layers=6, mlp_ratio=4.0, patch_size=4, in_channels=3, image_size=32):\n super().__init__(); self.patch_size=patch_size; self.image_size=image_size; self.embed_dim=embed_dim; self.num_heads=num_heads\n self.conv1=nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)\n num_patches=(image_size//patch_size)**2; self.linear_proj=nn.Linear(embed_dim*num_patches, embed_dim)\n self.transformer_layers=nn.ModuleList([nn.TransformerEncoderLayer(d_model=embed_dim,nhead=num_heads,dim_feedforward=int(embed_dim*mlp_ratio),dropout=0.0,batch_first=True) for _ in range(num_layers)])\n self.cls_token=nn.Parameter(torch.zeros(1,1,embed_dim)); self.fc_out=nn.Linear(embed_dim,num_classes)\n def forward(self,x):\n params=[self.conv1.weight,self.conv1.bias,self.linear_proj.weight,self.linear_proj.bias,self.cls_token]\n for layer in self.transformer_layers:\n params += [layer.self_attn.in_proj_weight,layer.self_attn.in_proj_bias,layer.self_attn.out_proj.weight,layer.self_attn.out_proj.bias,layer.linear1.weight,layer.linear1.bias,layer.linear2.weight,layer.linear2.bias,layer.norm1.weight,layer.norm1.bias,layer.norm2.weight,layer.norm2.bias]\n params += [self.fc_out.weight,self.fc_out.bias]\n return _ext.cvit_forward(x,params)\n"
}