callrun_correctness
{
"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 lstm_forward_cuda(torch::Tensor x, torch::Tensor h0, torch::Tensor c0,\n std::vector<torch::Tensor> w_ih, std::vector<torch::Tensor> w_hh,\n std::vector<torch::Tensor> b_ih, std::vector<torch::Tensor> b_hh,\n torch::Tensor fc_w, torch::Tensor fc_b, int64_t num_layers);\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 <cooperative_groups.h>\n#include <vector>\nnamespace cg = cooperative_groups;\n#define HIDDEN 256\n#define CHUNK 64\n#define VEC 4\n__device__ __forceinline__ float sigmoid_dev(float x) { return 1.0f / (1.0f + expf(-x)); }\n\n__global__ void lstm_layer_kernel(const float* __restrict__ inp, float* __restrict__ out_full,\n float* __restrict__ out_final, const float* __restrict__ h0,\n const float* __restrict__ c0, const float* __restrict__ w_ih,\n const float* __restrict__ w_hh, const float* __restrict__ b_ih,\n const float* __restrict__ b_hh, float* __restrict__ gates,\n float* __restrict__ hstate, float* __restrict__ cstate,\n int B, int T, int K, int dir, int steps, int final_mode) {\n cg::grid_group grid = cg::this_grid();\n int tx = threadIdx.x;\n int oj = tx / VEC;\n int lane = tx & (VEC - 1);\n int b = blockIdx.x, gate = blockIdx.y, chunk = blockIdx.z;\n int j = chunk * CHUNK + oj;\n if (gate == 0 && lane == 0) {\n hstate[b * HIDDEN + j] = h0[b * HIDDEN + j];\n cstate[b * HIDDEN + j] = c0[b * HIDDEN + j];\n }\n grid.sync();\n for (int s = 0; s < steps; ++s) {\n int t = (dir == 0) ? s : (T - 1 - s);\n int gj = gate * HIDDEN + j;\n float sum = (lane == 0) ? (b_ih[gj] + b_hh[gj]) : 0.0f;\n const float* xptr = inp + ((b * T + t) * K);\n const float* wip = w_ih + gj * K;\n for (int k = lane; k < K; k += VEC) sum = fmaf(xptr[k], wip[k], sum);\n const float* hp = hstate + b * HIDDEN;\n const float* whp = w_hh + gj * HIDDEN;\n for (int k = lane; k < HIDDEN; k += VEC) sum = fmaf(hp[k], whp[k], sum);\n sum += __shfl_down_sync(0xffffffff, sum, 2, VEC);\n sum += __shfl_down_sync(0xffffffff, sum, 1, VEC);\n if (lane == 0) gates[b * 4 * HIDDEN + gj] = sum;\n grid.sync();\n if (gate == 0 && lane == 0) {\n float i = sigmoid_dev(gates[b * 4 * HIDDEN + j]);\n float f = sigmoid_dev(gates[b * 4 * HIDDEN + HIDDEN + j]);\n float g = tanhf(gates[b * 4 * HIDDEN + 2 * HIDDEN + j]);\n float o = sigmoid_dev(gates[b * 4 * HIDDEN + 3 * HIDDEN + j]);\n float c = fmaf(f, cstate[b * HIDDEN + j], i * g);\n float h = o * tanhf(c);\n cstate[b * HIDDEN + j] = c; hstate[b * HIDDEN + j] = h;\n if (final_mode) {\n if ((dir == 0 && s == steps - 1) || (dir == 1 && s == 0)) out_final[b * (2 * HIDDEN) + dir * HIDDEN + j] = h;\n } else out_full[(b * T + t) * (2 * HIDDEN) + dir * HIDDEN + j] = h;\n }\n grid.sync();\n }\n}\n\n__global__ void fc_kernel(const float* __restrict__ inp, const float* __restrict__ w,\n const float* __restrict__ bias, float* __restrict__ out,\n int B, int O, int F) {\n int b = blockIdx.x, o = blockIdx.y, tid = threadIdx.x;\n __shared__ float smem[256];\n float sum = 0.0f;\n const float* xp = inp + b * F; const float* wp = w + o * F;\n for (int k = tid; k < F; k += 256) sum = fmaf(xp[k], wp[k], sum);\n smem[tid] = sum; __syncthreads();\n for (int stride = 128; stride > 0; stride >>= 1) { if (tid < stride) smem[tid] += smem[tid + stride]; __syncthreads(); }\n if (tid == 0) out[b * O + o] = smem[0] + bias[o];\n}\n\nstatic void launch_lstm(const torch::Tensor& inp, torch::Tensor& out_full, torch::Tensor& out_final,\n const torch::Tensor& h0, const torch::Tensor& c0, const torch::Tensor& w_ih,\n const torch::Tensor& w_hh, const torch::Tensor& b_ih, const torch::Tensor& b_hh,\n torch::Tensor& gates, torch::Tensor& hstate, torch::Tensor& cstate,\n int idx, int B, int T, int K, int dir, int steps, int final_mode) {\n const float* inp_p=inp.data_ptr<float>(); float* ofp=final_mode?nullptr:out_full.data_ptr<float>(); float* ofinp=final_mode?out_final.data_ptr<float>():nullptr;\n const float* h0p=h0.data_ptr<float>()+idx*B*HIDDEN; const float* c0p=c0.data_ptr<float>()+idx*B*HIDDEN;\n const float* wip=w_ih.data_ptr<float>(); const float* whp=w_hh.data_ptr<float>(); const float* bip=b_ih.data_ptr<float>(); const float* bhp=b_hh.data_ptr<float>();\n float* gp=gates.data_ptr<float>(); float* hp=hstate.data_ptr<float>(); float* cp=cstate.data_ptr<float>();\n dim3 grid(B,4,(HIDDEN+CHUNK-1)/CHUNK); dim3 block(CHUNK*VEC);\n void* args[]={&inp_p,&ofp,&ofinp,&h0p,&c0p,&wip,&whp,&bip,&bhp,&gp,&hp,&cp,&B,&T,&K,&dir,&steps,&final_mode};\n cudaError_t err=cudaLaunchCooperativeKernel((void*)lstm_layer_kernel,grid,block,args,0,at::cuda::getCurrentCUDAStream());\n TORCH_CHECK(err==cudaSuccess,\"cooperative LSTM launch failed: \",cudaGetErrorString(err));\n C10_CUDA_KERNEL_LAUNCH_CHECK();\n}\n\ntorch::Tensor lstm_forward_cuda(torch::Tensor x, torch::Tensor h0, torch::Tensor c0,\n std::vector<torch::Tensor> w_ih, std::vector<torch::Tensor> w_hh,\n std::vector<torch::Tensor> b_ih, std::vector<torch::Tensor> b_hh,\n torch::Tensor fc_w, torch::Tensor fc_b, int64_t num_layers) {\n int B=x.size(0), T=x.size(1), K=x.size(2), O=fc_w.size(0); auto opts=x.options();\n auto buf1=torch::empty({B,T,2*HIDDEN},opts), buf2=torch::empty({B,T,2*HIDDEN},opts), final_out=torch::empty({B,2*HIDDEN},opts);\n auto gates=torch::empty({B,4*HIDDEN},opts), hstate=torch::empty({B,HIDDEN},opts), cstate=torch::empty({B,HIDDEN},opts);\n torch::Tensor current=x;\n for(int l=0;l<(int)num_layers;++l){ bool last=(l==(int)num_layers-1); torch::Tensor out_full=(l%2==0)?buf1:buf2;\n for(int dir=0;dir<2;++dir){ int idx=l*2+dir; int steps=(last&&dir==1)?1:T;\n launch_lstm(current,out_full,final_out,h0,c0,w_ih[idx],w_hh[idx],b_ih[idx],b_hh[idx],gates,hstate,cstate,idx,B,T,K,dir,steps,last?1:0); }\n if(!last){ current=out_full; K=2*HIDDEN; }}\n auto y=torch::empty({B,O},opts);\n fc_kernel<<<dim3(B,O),256,0,at::cuda::getCurrentCUDAStream()>>>(final_out.data_ptr<float>(),fc_w.data_ptr<float>(),fc_b.data_ptr<float>(),y.data_ptr<float>(),B,O,2*HIDDEN);\n C10_CUDA_KERNEL_LAUNCH_CHECK(); return y;\n}\n'''\n_ext = load_inline(name='lstm_custom_coop_v3', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['lstm_forward_cuda'], with_cuda=True, extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.0):\n super().__init__(); self.lstm=nn.LSTM(input_size,hidden_size,num_layers,batch_first=True,dropout=dropout,bidirectional=True); self.fc=nn.Linear(hidden_size*2,output_size); self.num_layers=num_layers\n def forward(self,x,h0,c0):\n w_ih=[]; w_hh=[]; b_ih=[]; b_hh=[]\n for l in range(self.num_layers):\n for suf in ['', '_reverse']:\n w_ih.append(getattr(self.lstm,f'weight_ih_l{l}{suf}')); w_hh.append(getattr(self.lstm,f'weight_hh_l{l}{suf}')); b_ih.append(getattr(self.lstm,f'bias_ih_l{l}{suf}')); b_hh.append(getattr(self.lstm,f'bias_hh_l{l}{suf}'))\n return _ext.lstm_forward_cuda(x,h0,c0,w_ih,w_hh,b_ih,b_hh,self.fc.weight,self.fc.bias,self.num_layers)\n"
}