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#include <vector>\ntorch::Tensor vgg_forward(torch::Tensor x, std::vector<torch::Tensor> p);\n'''\n_cuda = r'''\n#include <torch/extension.h>\n#include <ATen/ATen.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n\n__global__ void relu_pool2d_cl_kernel(const float* __restrict__ x, float* __restrict__ y, int total, int C, int H, int W, int Ho, int Wo) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n int c = idx % C;\n int t = idx / C;\n int ow = t % Wo;\n t /= Wo;\n int oh = t % Ho;\n int n = t / Ho;\n int base = ((n * H + (oh << 1)) * W + (ow << 1)) * C + c;\n float m = fmaxf(fmaxf(x[base], x[base + C]), fmaxf(x[base + W*C], x[base + (W+1)*C]));\n y[idx] = fmaxf(m, 0.0f);\n}\n\nstatic torch::Tensor relu_pool2d_cl(torch::Tensor x) {\n int64_t N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);\n int64_t Ho = H / 2, Wo = W / 2;\n auto y = torch::empty_strided({N,C,Ho,Wo}, {Ho*Wo*C, 1, Wo*C, C}, x.options());\n int total = (int)(N*C*Ho*Wo);\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n relu_pool2d_cl_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), total, (int)C, (int)H, (int)W, (int)Ho, (int)Wo);\n return y;\n}\n\nstatic inline torch::Tensor conv_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b) {\n x = at::conv2d(x, w, b, {1,1}, {1,1}, {1,1}, 1);\n return at::relu_(x);\n}\nstatic inline torch::Tensor conv_only(torch::Tensor x, torch::Tensor w, torch::Tensor b) {\n return at::conv2d(x, w, b, {1,1}, {1,1}, {1,1}, 1);\n}\nstatic inline torch::Tensor lin_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b) {\n x = at::linear(x, w, b);\n return at::relu_(x);\n}\n\ntorch::Tensor vgg_forward(torch::Tensor x, std::vector<torch::Tensor> p) {\n x = x.contiguous(c10::MemoryFormat::ChannelsLast);\n int i = 0;\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_only(x, p[i++], p[i++]); x = relu_pool2d_cl(x);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_only(x, p[i++], p[i++]); x = relu_pool2d_cl(x);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_only(x, p[i++], p[i++]); x = relu_pool2d_cl(x);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_only(x, p[i++], p[i++]); x = relu_pool2d_cl(x);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_relu(x, p[i++], p[i++]);\n x = conv_only(x, p[i++], p[i++]); x = relu_pool2d_cl(x);\n x = at::flatten(x, 1, -1);\n x = lin_relu(x, p[i++], p[i++]);\n x = lin_relu(x, p[i++], p[i++]);\n x = at::linear(x, p[i++], p[i++]);\n return x;\n}\n'''\n_ext = load_inline(name='vgg16_cpp_forward_ext1', cpp_sources=[_cpp], cuda_sources=[_cuda], functions=['vgg_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)\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), 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), 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), nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), 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), nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True), 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), nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.ReLU(inplace=True), 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), nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Dropout(p=0.0), nn.Linear(4096, num_classes)\n )\n self.features.to(memory_format=torch.channels_last)\n def forward(self, x):\n f=self.features; c=self.classifier\n params=[f[0].weight,f[0].bias,f[2].weight,f[2].bias,f[5].weight,f[5].bias,f[7].weight,f[7].bias,f[10].weight,f[10].bias,f[12].weight,f[12].bias,f[14].weight,f[14].bias,f[17].weight,f[17].bias,f[19].weight,f[19].bias,f[21].weight,f[21].bias,f[24].weight,f[24].bias,f[26].weight,f[26].bias,f[28].weight,f[28].bias,c[0].weight,c[0].bias,c[3].weight,c[3].bias,c[6].weight,c[6].bias]\n return _ext.vgg_forward(x, params)\n"
}