callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n_cpp_src=r'''\n#include <torch/extension.h>\n#include <vector>\nstd::vector<torch::Tensor> pf_forward(torch::Tensor z, torch::Tensor w0, torch::Tensor u0, torch::Tensor b0, torch::Tensor wu0, torch::Tensor w1, torch::Tensor u1, torch::Tensor b1, torch::Tensor wu1, torch::Tensor w2, torch::Tensor u2, torch::Tensor b2, torch::Tensor wu2, torch::Tensor w3, torch::Tensor u3, torch::Tensor b3, torch::Tensor wu3, torch::Tensor w4, torch::Tensor u4, torch::Tensor b4, torch::Tensor wu4, torch::Tensor w5, torch::Tensor u5, torch::Tensor b5, torch::Tensor wu5, torch::Tensor w6, torch::Tensor u6, torch::Tensor b6, torch::Tensor wu6, torch::Tensor w7, torch::Tensor u7, torch::Tensor b7, torch::Tensor wu7);\n'''\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#include <vector>\n__device__ __forceinline__ float warp_sum(float v){for(int o=16;o>0;o>>=1)v+=__shfl_down_sync(0xffffffff,v,o);return v;}\n__device__ __forceinline__ const float* pp(int f,const float*p0,const float*p1,const float*p2,const float*p3,const float*p4,const float*p5,const float*p6,const float*p7){switch(f){case 0:return p0;case 1:return p1;case 2:return p2;case 3:return p3;case 4:return p4;case 5:return p5;case 6:return p6;default:return p7;}}\n__global__ void ker(const float* z,float* zo,float* sl,int B,const float*w0,const float*u0,const float*b0,const float*wu0,const float*w1,const float*u1,const float*b1,const float*wu1,const float*w2,const float*u2,const float*b2,const float*wu2,const float*w3,const float*u3,const float*b3,const float*wu3,const float*w4,const float*u4,const float*b4,const float*wu4,const float*w5,const float*u5,const float*b5,const float*wu5,const float*w6,const float*u6,const float*b6,const float*wu6,const float*w7,const float*u7,const float*b7,const float*wu7){int tid=blockIdx.x*blockDim.x+threadIdx.x;int s=tid>>5;int lane=threadIdx.x&31;if(s>=B)return;float zv=z[s*32+lane];float acc=0.f;for(int f=0;f<8;++f){const float*w=pp(f,w0,w1,w2,w3,w4,w5,w6,w7);const float*u=pp(f,u0,u1,u2,u3,u4,u5,u6,u7);const float*b=pp(f,b0,b1,b2,b3,b4,b5,b6,b7);const float*wu=pp(f,wu0,wu1,wu2,wu3,wu4,wu5,wu6,wu7);float lin=warp_sum(zv*w[lane]);lin=__shfl_sync(0xffffffff,lin,0)+b[0];float h=tanhf(lin);zv=zv+u[lane]*h;float hp=1.f-h*h;acc=acc+logf(fabsf(1.f+hp*wu[0])+1.e-8f);}zo[s*32+lane]=zv;if(lane==0)sl[s]=acc;}\nstd::vector<torch::Tensor> pf_forward(torch::Tensor z, torch::Tensor w0, torch::Tensor u0, torch::Tensor b0, torch::Tensor wu0, torch::Tensor w1, torch::Tensor u1, torch::Tensor b1, torch::Tensor wu1, torch::Tensor w2, torch::Tensor u2, torch::Tensor b2, torch::Tensor wu2, torch::Tensor w3, torch::Tensor u3, torch::Tensor b3, torch::Tensor wu3, torch::Tensor w4, torch::Tensor u4, torch::Tensor b4, torch::Tensor wu4, torch::Tensor w5, torch::Tensor u5, torch::Tensor b5, torch::Tensor wu5, torch::Tensor w6, torch::Tensor u6, torch::Tensor b6, torch::Tensor wu6, torch::Tensor w7, torch::Tensor u7, torch::Tensor b7, torch::Tensor wu7){int B=z.size(0);auto zo=torch::empty_like(z);auto sl=torch::empty({B},z.options());int th=256,bl=(B*32+th-1)/th;ker<<<bl,th,0,at::cuda::getCurrentCUDAStream()>>>(z.data_ptr<float>(),zo.data_ptr<float>(),sl.data_ptr<float>(),B,w0.data_ptr<float>(),u0.data_ptr<float>(),b0.data_ptr<float>(),wu0.data_ptr<float>(),w1.data_ptr<float>(),u1.data_ptr<float>(),b1.data_ptr<float>(),wu1.data_ptr<float>(),w2.data_ptr<float>(),u2.data_ptr<float>(),b2.data_ptr<float>(),wu2.data_ptr<float>(),w3.data_ptr<float>(),u3.data_ptr<float>(),b3.data_ptr<float>(),wu3.data_ptr<float>(),w4.data_ptr<float>(),u4.data_ptr<float>(),b4.data_ptr<float>(),wu4.data_ptr<float>(),w5.data_ptr<float>(),u5.data_ptr<float>(),b5.data_ptr<float>(),wu5.data_ptr<float>(),w6.data_ptr<float>(),u6.data_ptr<float>(),b6.data_ptr<float>(),wu6.data_ptr<float>(),w7.data_ptr<float>(),u7.data_ptr<float>(),b7.data_ptr<float>(),wu7.data_ptr<float>());return{zo,sl};}\n'''\n_ext=load_inline(name='planar_flow_ext_pre_v1',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['pf_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,dim,num_flows):\n super().__init__(); self.dim=dim; self.num_flows=num_flows; self.w=nn.ParameterList([nn.Parameter(torch.randn(dim)) for _ in range(num_flows)]); self.u=nn.ParameterList([nn.Parameter(torch.randn(dim)) for _ in range(num_flows)]); self.b=nn.ParameterList([nn.Parameter(torch.zeros(1)) for _ in range(num_flows)])\n def _cu(self,w,u):\n wtu=(w*u).sum(); m=-1+F.softplus(wtu); return u+(m-wtu)*w/(w**2).sum()\n def forward(self,z):\n us=[]; wus=[]\n for i in range(8):\n uh=self._cu(self.w[i],self.u[i]); us.append(uh); wus.append((self.w[i]*uh).sum().reshape(1))\n o=_ext.pf_forward(z.contiguous(),self.w[0],us[0],self.b[0],wus[0],self.w[1],us[1],self.b[1],wus[1],self.w[2],us[2],self.b[2],wus[2],self.w[3],us[3],self.b[3],wus[3],self.w[4],us[4],self.b[4],wus[4],self.w[5],us[5],self.b[5],wus[5],self.w[6],us[6],self.b[6],wus[6],self.w[7],us[7],self.b[7],wus[7]); return o[0],o[1]\n"
}