call → compile_kernel
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>\nvoid conv_gelu_gap4_launcher(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor out);\ntorch::Tensor conv_gelu_gap(torch::Tensor x, torch::Tensor w, torch::Tensor b) {\n auto out = torch::empty({x.size(0), w.size(0)}, x.options());\n conv_gelu_gap4_launcher(x, w, b, out);\n return out;\n}\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__device__ __forceinline__ float gelu_exact(float v) {\n return 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));\n}\n\n__global__ __launch_bounds__(256, 2)\nvoid conv_gelu_gap4_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n float* __restrict__ out,\n int N, int H, int W, int O) {\n int og = blockIdx.x * 4;\n int n = blockIdx.y;\n int tid = threadIdx.x;\n constexpr int C = 8;\n constexpr int K2 = 72;\n int Ho = H - 2;\n int Wo = W - 2;\n int S = Ho * Wo;\n __shared__ float sw[4 * K2];\n __shared__ float sh[4][256];\n for (int i = tid; i < 4*K2; i += blockDim.x) {\n int oo = i / K2;\n int kk = i - oo * K2;\n sw[i] = w[(og + oo) * K2 + kk];\n }\n __syncthreads();\n float b0 = b[og+0], b1 = b[og+1], b2 = b[og+2], b3 = b[og+3];\n float sum0 = 0.f, sum1 = 0.f, sum2 = 0.f, sum3 = 0.f;\n for (int p = tid; p < S; p += blockDim.x) {\n int yy = p / Wo;\n int xx = p - yy * Wo;\n float v0 = b0, v1 = b1, v2 = b2, v3 = b3;\n #pragma unroll\n for (int c = 0; c < C; ++c) {\n const float* xp = x + ((n * C + c) * H + yy) * W + xx;\n float a0 = xp[0];\n float a1 = xp[1];\n float a2 = xp[2];\n float a3 = xp[W];\n float a4 = xp[W+1];\n float a5 = xp[W+2];\n float a6 = xp[2*W];\n float a7 = xp[2*W+1];\n float a8 = xp[2*W+2];\n const float* w0 = sw + 0*K2 + c*9;\n const float* w1 = sw + 1*K2 + c*9;\n const float* w2 = sw + 2*K2 + c*9;\n const float* w3 = sw + 3*K2 + c*9;\n v0 = fmaf(a0, w0[0], v0); v1 = fmaf(a0, w1[0], v1); v2 = fmaf(a0, w2[0], v2); v3 = fmaf(a0, w3[0], v3);\n v0 = fmaf(a1, w0[1], v0); v1 = fmaf(a1, w1[1], v1); v2 = fmaf(a1, w2[1], v2); v3 = fmaf(a1, w3[1], v3);\n v0 = fmaf(a2, w0[2], v0); v1 = fmaf(a2, w1[2], v1); v2 = fmaf(a2, w2[2], v2); v3 = fmaf(a2, w3[2], v3);\n v0 = fmaf(a3, w0[3], v0); v1 = fmaf(a3, w1[3], v1); v2 = fmaf(a3, w2[3], v2); v3 = fmaf(a3, w3[3], v3);\n v0 = fmaf(a4, w0[4], v0); v1 = fmaf(a4, w1[4], v1); v2 = fmaf(a4, w2[4], v2); v3 = fmaf(a4, w3[4], v3);\n v0 = fmaf(a5, w0[5], v0); v1 = fmaf(a5, w1[5], v1); v2 = fmaf(a5, w2[5], v2); v3 = fmaf(a5, w3[5], v3);\n v0 = fmaf(a6, w0[6], v0); v1 = fmaf(a6, w1[6], v1); v2 = fmaf(a6, w2[6], v2); v3 = fmaf(a6, w3[6], v3);\n v0 = fmaf(a7, w0[7], v0); v1 = fmaf(a7, w1[7], v1); v2 = fmaf(a7, w2[7], v2); v3 = fmaf(a7, w3[7], v3);\n v0 = fmaf(a8, w0[8], v0); v1 = fmaf(a8, w1[8], v1); v2 = fmaf(a8, w2[8], v2); v3 = fmaf(a8, w3[8], v3);\n }\n sum0 += gelu_exact(v0);\n sum1 += gelu_exact(v1);\n sum2 += gelu_exact(v2);\n sum3 += gelu_exact(v3);\n }\n sh[0][tid] = sum0; sh[1][tid] = sum1; sh[2][tid] = sum2; sh[3][tid] = sum3;\n __syncthreads();\n for (int s = 128; s > 32; s >>= 1) {\n if (tid < s) {\n sh[0][tid] += sh[0][tid+s];\n sh[1][tid] += sh[1][tid+s];\n sh[2][tid] += sh[2][tid+s];\n sh[3][tid] += sh[3][tid+s];\n }\n __syncthreads();\n }\n if (tid < 32) {\n float r0 = sh[0][tid] + sh[0][tid+32];\n float r1 = sh[1][tid] + sh[1][tid+32];\n float r2 = sh[2][tid] + sh[2][tid+32];\n float r3 = sh[3][tid] + sh[3][tid+32];\n #pragma unroll\n for (int off=16; off>0; off>>=1) {\n r0 += __shfl_down_sync(0xffffffff, r0, off);\n r1 += __shfl_down_sync(0xffffffff, r1, off);\n r2 += __shfl_down_sync(0xffffffff, r2, off);\n r3 += __shfl_down_sync(0xffffffff, r3, off);\n }\n if (tid == 0) {\n float inv = 1.0f / (float)S;\n out[n*O + og+0] = r0 * inv;\n out[n*O + og+1] = r1 * inv;\n out[n*O + og+2] = r2 * inv;\n out[n*O + og+3] = r3 * inv;\n }\n }\n}\n\nvoid conv_gelu_gap4_launcher(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor out) {\n int N=(int)x.size(0), H=(int)x.size(2), W=(int)x.size(3), O=(int)w.size(0);\n dim3 block(256);\n dim3 grid(O/4, N);\n conv_gelu_gap4_kernel<<<grid, block>>>(x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), N, H, W, O);\n}\n'''\n\n_ext = load_inline(name='conv_gelu_gap_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['conv_gelu_gap'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n def forward(self, x):\n return _ext.conv_gelu_gap(x.contiguous(), self.conv.weight, self.conv.bias)\n"
}