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>\n#include <vector>\nat::Tensor gru_forward_cuda(at::Tensor x, at::Tensor h0, std::vector<at::Tensor> w_ih, std::vector<at::Tensor> w_hh, std::vector<at::Tensor> b_ih, std::vector<at::Tensor> b_hh, bool batch_first);\nat::Tensor gru_forward(at::Tensor x, at::Tensor h0, std::vector<at::Tensor> w_ih, std::vector<at::Tensor> w_hh, std::vector<at::Tensor> b_ih, std::vector<at::Tensor> b_hh, bool batch_first){return gru_forward_cuda(x,h0,w_ih,w_hh,b_ih,b_hh,batch_first);}'''\ncuda_src=r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\nstatic __device__ __forceinline__ float sig(float x){return 1.0f/(1.0f+expf(-x));}\nstatic __device__ __forceinline__ float red16(float v){unsigned m=((threadIdx.x&16)?0xffff0000u:0x0000ffffu); v+=__shfl_down_sync(m,v,8,16);v+=__shfl_down_sync(m,v,4,16);v+=__shfl_down_sync(m,v,2,16);v+=__shfl_down_sync(m,v,1,16);return v;}\nstatic __device__ __forceinline__ void barr(int* bar,int n,int& s){__syncthreads();__threadfence();if(threadIdx.x==0){int t=atomicAdd(&bar[0],1);if(t==n-1){bar[0]=0;__threadfence();atomicExch(&bar[1],s);}else while(atomicAdd(&bar[1],0)!=s){}}__syncthreads();s=1-s;}\n__global__ void kern(const float* __restrict__ x,const float* __restrict__ h0,const float** __restrict__ wih,const float** __restrict__ whh,const float** __restrict__ bih,const float** __restrict__ bhh,float* __restrict__ state,float* __restrict__ ns,float* __restrict__ out,int* __restrict__ bar,int S,int B,int I,int H,int L,int bf){int tiles=H/16,bid=blockIdx.x,b=bid/tiles,tile=bid-b*tiles,tid=threadIdx.x,lane=tid&15,grp=tid>>4,h=tile*16+grp,sense=1,nblk=gridDim.x; if(b<B&&h<H)for(int l=0;l<L;l++)if(lane==0){int idx=(l*B+b)*H+h;state[idx]=h0[idx];ns[idx]=h0[idx];} barr(bar,nblk,sense); for(int t=0;t<S;t++){for(int l=0;l<L;l++){if(b<B&&h<H){const float *wi=wih[l],*wh=whh[l],*bi=bih[l],*bh=bhh[l];int K=(l==0?I:H);float ir=0,iz=0,in=0,hr=0,hz=0,hn=0;for(int k=lane;k<K;k+=16){float xv=(l==0?(bf?x[(b*S+t)*I+k]:x[(t*B+b)*I+k]):state[((l-1)*B+b)*H+k]);ir=fmaf(wi[(h)*K+k],xv,ir);iz=fmaf(wi[(H+h)*K+k],xv,iz);in=fmaf(wi[(2*H+h)*K+k],xv,in);}const float* st=state+(l*B+b)*H;for(int k=lane;k<H;k+=16){float hv=st[k];hr=fmaf(wh[(h)*H+k],hv,hr);hz=fmaf(wh[(H+h)*H+k],hv,hz);hn=fmaf(wh[(2*H+h)*H+k],hv,hn);}ir=red16(ir);iz=red16(iz);in=red16(in);hr=red16(hr);hz=red16(hz);hn=red16(hn);if(lane==0){float hp=st[h];float r=sig(ir+bi[h]+hr+bh[h]);float z=sig(iz+bi[H+h]+hz+bh[H+h]);float nv=tanhf(in+bi[2*H+h]+r*(hn+bh[2*H+h]));ns[(l*B+b)*H+h]=fmaf(z,hp-nv,nv);}}barr(bar,nblk,sense);if(b<B&&h<H&&lane==0){float v=ns[(l*B+b)*H+h];state[(l*B+b)*H+h]=v;if(l==L-1){if(bf)out[(b*S+t)*H+h]=v;else out[(t*B+b)*H+h]=v;}}barr(bar,nblk,sense);}}}\nat::Tensor gru_forward_cuda(at::Tensor x,at::Tensor h0,std::vector<at::Tensor>w_ih,std::vector<at::Tensor>w_hh,std::vector<at::Tensor>b_ih,std::vector<at::Tensor>b_hh,bool bf){int64_t S=bf?x.size(1):x.size(0),B=bf?x.size(0):x.size(1),I=x.size(2),L=h0.size(0),H=h0.size(2);auto out=torch::empty(bf?std::vector<int64_t>{B,S,H}:std::vector<int64_t>{S,B,H},x.options());auto state=torch::empty({L,B,H},x.options()),ns=torch::empty({L,B,H},x.options());auto bar=torch::zeros({2},x.options().dtype(torch::kInt32));auto opts=torch::TensorOptions().dtype(torch::kInt64).device(x.device());auto wihp=torch::empty({L},opts),whhp=torch::empty({L},opts),bihp=torch::empty({L},opts),bhhp=torch::empty({L},opts);std::vector<int64_t>a(L),bb(L),c(L),d(L);for(int i=0;i<L;i++){a[i]=(int64_t)w_ih[i].data_ptr<float>();bb[i]=(int64_t)w_hh[i].data_ptr<float>();c[i]=(int64_t)b_ih[i].data_ptr<float>();d[i]=(int64_t)b_hh[i].data_ptr<float>();}cudaMemcpy(wihp.data_ptr<int64_t>(),a.data(),L*8,cudaMemcpyHostToDevice);cudaMemcpy(whhp.data_ptr<int64_t>(),bb.data(),L*8,cudaMemcpyHostToDevice);cudaMemcpy(bihp.data_ptr<int64_t>(),c.data(),L*8,cudaMemcpyHostToDevice);cudaMemcpy(bhhp.data_ptr<int64_t>(),d.data(),L*8,cudaMemcpyHostToDevice);kern<<<B*(H/16),256>>>(x.data_ptr<float>(),h0.data_ptr<float>(),(const float**)wihp.data_ptr<int64_t>(),(const float**)whhp.data_ptr<int64_t>(),(const float**)bihp.data_ptr<int64_t>(),(const float**)bhhp.data_ptr<int64_t>(),state.data_ptr<float>(),ns.data_ptr<float>(),out.data_ptr<float>(),bar.data_ptr<int>(),S,B,I,H,L,bf?1:0);return out;}'''\n_ext=load_inline(name='gru_custom_ext_v3',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['gru_forward'],extra_cflags=['-O3'],extra_cuda_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,input_size,hidden_size,num_layers=3,bias=True,batch_first=False):\n super().__init__();self.gru=nn.GRU(input_size,hidden_size,num_layers,bias,batch_first,dropout=0,bidirectional=False);self.num_layers=num_layers;self.bias=bias;self.batch_first=batch_first\n def forward(self,x,h0):\n return _ext.gru_forward(x,h0,[getattr(self.gru,f'weight_ih_l{i}') for i in range(self.num_layers)],[getattr(self.gru,f'weight_hh_l{i}') for i in range(self.num_layers)],[getattr(self.gru,f'bias_ih_l{i}') for i in range(self.num_layers)],[getattr(self.gru,f'bias_hh_l{i}') for i in range(self.num_layers)],self.batch_first)\n"
}