call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp = r'''\n#include <torch/extension.h>\n\ntorch::Tensor gru_forward(torch::Tensor x, torch::Tensor h0,\n torch::Tensor wih0, torch::Tensor whh0, torch::Tensor bih0, torch::Tensor bhh0,\n torch::Tensor wih1, torch::Tensor whh1, torch::Tensor bih1, torch::Tensor bhh1,\n torch::Tensor wih2, torch::Tensor whh2, torch::Tensor bih2, torch::Tensor bhh2,\n torch::Tensor wih3, torch::Tensor whh3, torch::Tensor bih3, torch::Tensor bhh3,\n torch::Tensor wih4, torch::Tensor whh4, torch::Tensor bih4, torch::Tensor bhh4,\n torch::Tensor wih5, torch::Tensor whh5, torch::Tensor bih5, torch::Tensor bhh5,\n bool batch_first);\n'''\n\n_cuda = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cooperative_groups.h>\n#include <vector>\nnamespace cg = cooperative_groups;\n\n#define HSIZE 256\n#define NLAYER 6\n#define TILE 16\n#define LANES 8\n\n__device__ __forceinline__ float warp_reduce8(float v) {\n unsigned int group = (threadIdx.x & 31) >> 3;\n unsigned int mask = 0xFFu << (group * 8);\n v += __shfl_down_sync(mask, v, 4, 8);\n v += __shfl_down_sync(mask, v, 2, 8);\n v += __shfl_down_sync(mask, v, 1, 8);\n return v;\n}\n\n__device__ __forceinline__ float sigmoidf_stable(float x) {\n return 1.0f / (1.0f + expf(-x));\n}\n\n__global__ void gru_kernel(const float* __restrict__ x, const float* __restrict__ h0, float* __restrict__ hn, float* __restrict__ scratch,\n const float* __restrict__ wih0, const float* __restrict__ whh0, const float* __restrict__ bih0, const float* __restrict__ bhh0,\n const float* __restrict__ wih1, const float* __restrict__ whh1, const float* __restrict__ bih1, const float* __restrict__ bhh1,\n const float* __restrict__ wih2, const float* __restrict__ whh2, const float* __restrict__ bih2, const float* __restrict__ bhh2,\n const float* __restrict__ wih3, const float* __restrict__ whh3, const float* __restrict__ bih3, const float* __restrict__ bhh3,\n const float* __restrict__ wih4, const float* __restrict__ whh4, const float* __restrict__ bih4, const float* __restrict__ bhh4,\n const float* __restrict__ wih5, const float* __restrict__ whh5, const float* __restrict__ bih5, const float* __restrict__ bhh5,\n int S, int B, int I, int batch_first) {\n\n cg::grid_group grid = cg::this_grid();\n const int b = blockIdx.x;\n const int tile = blockIdx.y;\n const int tid = threadIdx.x;\n const int row_in_tile = tid / LANES;\n const int lane = tid & (LANES-1);\n const int row = tile * TILE + row_in_tile;\n\n __shared__ float sh_h[HSIZE];\n __shared__ float sh_in[HSIZE];\n\n // initialize output hidden state from h0\n for (int l = 0; l < NLAYER; ++l) {\n if (tid < TILE) {\n int r = tile * TILE + tid;\n hn[(l * B + b) * HSIZE + r] = h0[(l * B + b) * HSIZE + r];\n }\n }\n grid.sync();\n\n const float* wihs[NLAYER] = {wih0,wih1,wih2,wih3,wih4,wih5};\n const float* whhs[NLAYER] = {whh0,whh1,whh2,whh3,whh4,whh5};\n const float* bihs[NLAYER] = {bih0,bih1,bih2,bih3,bih4,bih5};\n const float* bhhs[NLAYER] = {bhh0,bhh1,bhh2,bhh3,bhh4,bhh5};\n\n for (int l = 0; l < NLAYER; ++l) {\n const float* __restrict__ wih = wihs[l];\n const float* __restrict__ whh = whhs[l];\n const float* __restrict__ bih = bihs[l];\n const float* __restrict__ bhh = bhhs[l];\n int in_dim = (l == 0 ? I : HSIZE);\n int write_buf = l & 1;\n int read_buf = (l - 1) & 1;\n for (int t = 0; t < S; ++t) {\n // load recurrent hidden vector for this layer/batch into shared memory\n if (tid < HSIZE) sh_h[tid] = hn[(l * B + b) * HSIZE + tid];\n // load input vector into shared memory\n if (l == 0) {\n if (tid < I) {\n int idx = batch_first ? ((b * S + t) * I + tid) : ((t * B + b) * I + tid);\n sh_in[tid] = x[idx];\n }\n } else {\n if (tid < HSIZE) sh_in[tid] = scratch[((read_buf * S + t) * B + b) * HSIZE + tid];\n }\n __syncthreads();\n\n float ar = 0.f, az = 0.f, ain = 0.f, ahn = 0.f;\n // input contribution\n for (int k = lane; k < in_dim; k += LANES) {\n float xv = sh_in[k];\n ar += wih[(0 * HSIZE + row) * in_dim + k] * xv;\n az += wih[(1 * HSIZE + row) * in_dim + k] * xv;\n ain += wih[(2 * HSIZE + row) * in_dim + k] * xv;\n }\n // recurrent contribution\n for (int k = lane; k < HSIZE; k += LANES) {\n float hv = sh_h[k];\n ar += whh[(0 * HSIZE + row) * HSIZE + k] * hv;\n az += whh[(1 * HSIZE + row) * HSIZE + k] * hv;\n ahn += whh[(2 * HSIZE + row) * HSIZE + k] * hv;\n }\n ar = warp_reduce8(ar);\n az = warp_reduce8(az);\n ain = warp_reduce8(ain);\n ahn = warp_reduce8(ahn);\n\n if (lane == 0) {\n float hprev = sh_h[row];\n float r = sigmoidf_stable(ar + bih[row] + bhh[row]);\n float z = sigmoidf_stable(az + bih[HSIZE + row] + bhh[HSIZE + row]);\n float n = tanhf(ain + bih[2*HSIZE + row] + r * (ahn + bhh[2*HSIZE + row]));\n float hnew = (1.0f - z) * n + z * hprev;\n hn[(l * B + b) * HSIZE + row] = hnew;\n scratch[((write_buf * S + t) * B + b) * HSIZE + row] = hnew;\n }\n grid.sync();\n }\n grid.sync();\n }\n}\n\ntorch::Tensor gru_forward(torch::Tensor x, torch::Tensor h0,\n torch::Tensor wih0, torch::Tensor whh0, torch::Tensor bih0, torch::Tensor bhh0,\n torch::Tensor wih1, torch::Tensor whh1, torch::Tensor bih1, torch::Tensor bhh1,\n torch::Tensor wih2, torch::Tensor whh2, torch::Tensor bih2, torch::Tensor bhh2,\n torch::Tensor wih3, torch::Tensor whh3, torch::Tensor bih3, torch::Tensor bhh3,\n torch::Tensor wih4, torch::Tensor whh4, torch::Tensor bih4, torch::Tensor bhh4,\n torch::Tensor wih5, torch::Tensor whh5, torch::Tensor bih5, torch::Tensor bhh5,\n bool batch_first) {\n int64_t S = batch_first ? x.size(1) : x.size(0);\n int64_t B = batch_first ? x.size(0) : x.size(1);\n int64_t I = x.size(2);\n auto hn = torch::empty_like(h0);\n auto scratch = torch::empty({2, S, B, (int64_t)HSIZE}, x.options());\n dim3 grid((unsigned)B, HSIZE / TILE, 1);\n dim3 block(TILE * LANES, 1, 1);\n void* args[] = {&x, &h0, &hn, &scratch,\n &wih0,&whh0,&bih0,&bhh0, &wih1,&whh1,&bih1,&bhh1, &wih2,&whh2,&bih2,&bhh2,\n &wih3,&whh3,&bih3,&bhh3, &wih4,&whh4,&bih4,&bhh4, &wih5,&whh5,&bih5,&bhh5,\n &S, &B, &I, &batch_first};\n // Need raw pointers in launch args, not Tensor objects: use a small lambda-style expansion below is not possible with void* args above.\n const float* xp = x.data_ptr<float>();\n const float* h0p = h0.data_ptr<float>();\n float* hnp = hn.data_ptr<float>();\n float* sp = scratch.data_ptr<float>();\n const float* p_wih0=wih0.data_ptr<float>(); const float* p_whh0=whh0.data_ptr<float>(); const float* p_bih0=bih0.data_ptr<float>(); const float* p_bhh0=bhh0.data_ptr<float>();\n const float* p_wih1=wih1.data_ptr<float>(); const float* p_whh1=whh1.data_ptr<float>(); const float* p_bih1=bih1.data_ptr<float>(); const float* p_bhh1=bhh1.data_ptr<float>();\n const float* p_wih2=wih2.data_ptr<float>(); const float* p_whh2=whh2.data_ptr<float>(); const float* p_bih2=bih2.data_ptr<float>(); const float* p_bhh2=bhh2.data_ptr<float>();\n const float* p_wih3=wih3.data_ptr<float>(); const float* p_whh3=whh3.data_ptr<float>(); const float* p_bih3=bih3.data_ptr<float>(); const float* p_bhh3=bhh3.data_ptr<float>();\n const float* p_wih4=wih4.data_ptr<float>(); const float* p_whh4=whh4.data_ptr<float>(); const float* p_bih4=bih4.data_ptr<float>(); const float* p_bhh4=bhh4.data_ptr<float>();\n const float* p_wih5=wih5.data_ptr<float>(); const float* p_whh5=whh5.data_ptr<float>(); const float* p_bih5=bih5.data_ptr<float>(); const float* p_bhh5=bhh5.data_ptr<float>();\n int Si=(int)S, Bi=(int)B, Ii=(int)I, bf=batch_first?1:0;\n void* kargs[] = {&xp, &h0p, &hnp, &sp,\n &p_wih0,&p_whh0,&p_bih0,&p_bhh0, &p_wih1,&p_whh1,&p_bih1,&p_bhh1, &p_wih2,&p_whh2,&p_bih2,&p_bhh2,\n &p_wih3,&p_whh3,&p_bih3,&p_bhh3, &p_wih4,&p_whh4,&p_bih4,&p_bhh4, &p_wih5,&p_whh5,&p_bih5,&p_bhh5,\n &Si, &Bi, &Ii, &bf};\n cudaLaunchCooperativeKernel((void*)gru_kernel, grid, block, kargs, 0, at::cuda::getCurrentCUDAStream());\n return hn;\n}\n'''\n\n_ext = load_inline(name='gru_custom_ext_v1', cpp_sources=_cpp, cuda_sources=_cuda, functions=['gru_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, input_size, hidden_size, num_layers=3, bias=True, batch_first=False):\n super(ModelNew, self).__init__()\n self.gru = nn.GRU(input_size, hidden_size, num_layers, bias, batch_first, dropout=0, bidirectional=False)\n self.batch_first = batch_first\n self.num_layers = num_layers\n self.hidden_size = hidden_size\n self.bias = bias\n def forward(self, x, h0):\n return _ext.gru_forward(x, h0,\n self.gru.weight_ih_l0, self.gru.weight_hh_l0, self.gru.bias_ih_l0, self.gru.bias_hh_l0,\n self.gru.weight_ih_l1, self.gru.weight_hh_l1, self.gru.bias_ih_l1, self.gru.bias_hh_l1,\n self.gru.weight_ih_l2, self.gru.weight_hh_l2, self.gru.bias_ih_l2, self.gru.bias_hh_l2,\n self.gru.weight_ih_l3, self.gru.weight_hh_l3, self.gru.bias_ih_l3, self.gru.bias_hh_l3,\n self.gru.weight_ih_l4, self.gru.weight_hh_l4, self.gru.bias_ih_l4, self.gru.bias_hh_l4,\n self.gru.weight_ih_l5, self.gru.weight_hh_l5, self.gru.bias_ih_l5, self.gru.bias_hh_l5,\n self.batch_first)\n"
}