call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r\"\"\"\ntorch::Tensor relu_maxpool2d_cl(torch::Tensor x);\ntorch::Tensor relu_maxpool2d_last(torch::Tensor x);\n\"\"\"\n_cuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__device__ __forceinline__ float relu_f(float v) { return v > 0.0f ? v : 0.0f; }\n__device__ __forceinline__ float max4(float a, float b, float c, float d) { return fmaxf(fmaxf(a,b), fmaxf(c,d)); }\n\n__global__ void relu_maxpool2d_cl_vec4_kernel(const float4* __restrict__ x4, float4* __restrict__ y4,\n int N, int C4, int H, int W, int OH, int OW) {\n int64_t total = (int64_t)N * OH * OW * C4;\n int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = (int64_t)blockDim.x * gridDim.x;\n for (; idx < total; idx += stride) {\n int cv = idx % C4;\n int64_t t = idx / C4;\n int ow = t % OW;\n t /= OW;\n int oh = t % OH;\n int n = t / OH;\n int ih = oh * 2;\n int iw = ow * 2;\n int64_t base = (((int64_t)n * H + ih) * W + iw) * C4 + cv;\n float4 a = x4[base];\n float4 b = x4[base + C4];\n float4 c = x4[base + (int64_t)W * C4];\n float4 d = x4[base + (int64_t)(W + 1) * C4];\n float4 o;\n o.x = relu_f(max4(a.x,b.x,c.x,d.x));\n o.y = relu_f(max4(a.y,b.y,c.y,d.y));\n o.z = relu_f(max4(a.z,b.z,c.z,d.z));\n o.w = relu_f(max4(a.w,b.w,c.w,d.w));\n y4[idx] = o;\n }\n}\n\n__global__ void relu_maxpool2d_last_nchw_kernel(const float* __restrict__ x, float* __restrict__ y,\n int N, int C, int H, int W, int OH, int OW) {\n int64_t total = (int64_t)N * C * OH * OW;\n int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = (int64_t)blockDim.x * gridDim.x;\n for (; idx < total; idx += stride) {\n int ow = idx % OW;\n int64_t t = idx / OW;\n int oh = t % OH;\n t /= OH;\n int c = t % C;\n int n = t / C;\n int ih = oh * 2;\n int iw = ow * 2;\n int64_t base = ((int64_t)n * H * W * C) + ((int64_t)ih * W + iw) * C + c;\n float m = x[base];\n float v = x[base + C]; if (v > m) m = v;\n v = x[base + (int64_t)W * C]; if (v > m) m = v;\n v = x[base + (int64_t)(W + 1) * C]; if (v > m) m = v;\n y[idx] = m > 0.0f ? m : 0.0f;\n }\n}\n\ntorch::Tensor relu_maxpool2d_cl(torch::Tensor x) {\n const int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);\n const int OH = H / 2, OW = W / 2, C4 = C / 4;\n auto y = torch::empty_strided({N, C, OH, OW}, {C*OH*OW, 1, C*OW, C}, x.options());\n int threads = 256;\n int64_t total = (int64_t)N * OH * OW * C4;\n int blocks = (int)((total + threads - 1) / threads);\n if (blocks > 4096) blocks = 4096;\n relu_maxpool2d_cl_vec4_kernel<<<blocks, threads>>>(reinterpret_cast<const float4*>(x.data_ptr<float>()), reinterpret_cast<float4*>(y.data_ptr<float>()), N, C4, H, W, OH, OW);\n return y;\n}\n\ntorch::Tensor relu_maxpool2d_last(torch::Tensor x) {\n const int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);\n const int OH = H / 2, OW = W / 2;\n auto y = torch::empty({N, C, OH, OW}, x.options());\n int threads = 256;\n int64_t total = (int64_t)N * C * OH * OW;\n int blocks = (int)((total + threads - 1) / threads);\n if (blocks > 4096) blocks = 4096;\n relu_maxpool2d_last_nchw_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, OH, OW);\n return y;\n}\n\"\"\"\n\n_ext = load_inline(name=\"vgg19_fused_pool_ext_v2\", cpp_sources=_cpp_src, cuda_sources=_cuda_src,\n functions=[\"relu_maxpool2d_cl\", \"relu_maxpool2d_last\"], verbose=False,\n extra_cflags=[\"-O3\"], extra_cuda_cflags=[\"-O3\"])\n\nclass ModelNew(nn.Module):\n def __init__(self, num_classes=1000):\n super(ModelNew, self).__init__()\n self.features = nn.Sequential(\n nn.Conv2d(3, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2),\n nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(128, 128, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2),\n nn.Conv2d(128, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2),\n nn.Conv2d(256, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True),\n nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2)\n )\n self.classifier = nn.Sequential(\n nn.Linear(512 * 7 * 7, 4096), nn.ReLU(inplace=True), nn.Dropout(p=0.0),\n nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Dropout(p=0.0),\n nn.Linear(4096, num_classes)\n )\n self.features.to(memory_format=torch.channels_last)\n def forward(self, x):\n x = x.contiguous(memory_format=torch.channels_last)\n f = self.features\n x = F.relu(f[0](x), inplace=True); x = f[2](x); x = _ext.relu_maxpool2d_cl(x)\n x = F.relu(f[5](x), inplace=True); x = f[7](x); x = _ext.relu_maxpool2d_cl(x)\n x = F.relu(f[10](x), inplace=True); x = F.relu(f[12](x), inplace=True); x = F.relu(f[14](x), inplace=True); x = f[16](x); x = _ext.relu_maxpool2d_cl(x)\n x = F.relu(f[19](x), inplace=True); x = F.relu(f[21](x), inplace=True); x = F.relu(f[23](x), inplace=True); x = f[25](x); x = _ext.relu_maxpool2d_cl(x)\n x = F.relu(f[28](x), inplace=True); x = F.relu(f[30](x), inplace=True); x = F.relu(f[32](x), inplace=True); x = f[34](x); x = _ext.relu_maxpool2d_last(x)\n x = torch.flatten(x, 1)\n c = self.classifier\n x = F.relu(c[0](x), inplace=True); x = F.relu(c[3](x), inplace=True); x = c[6](x)\n return x\n"
}