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\ncpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\ntorch::Tensor tail_forward(torch::Tensor x, std::vector<torch::Tensor> p, double eps);\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 <vector>\n\n__global__ void pool_kernel(const float* __restrict__ x, float* __restrict__ y, int N, int C, int HW) {\n int idx = blockIdx.x;\n int n = idx / C;\n int c = idx - n*C;\n float sum = 0.0f;\n const float* base = x + ((n*C + c)*HW);\n for (int i = threadIdx.x; i < HW; i += blockDim.x) sum += base[i];\n __shared__ float smem[256];\n smem[threadIdx.x] = sum;\n __syncthreads();\n for (int s=128; s>0; s>>=1) { if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x+s]; __syncthreads(); }\n if (threadIdx.x == 0) y[idx] = smem[0] / (float)HW;\n}\n\n__global__ void linear_kernel(const float* __restrict__ x, const float* __restrict__ w,\n const float* __restrict__ bias, const float* __restrict__ bn_w,\n const float* __restrict__ bn_b, const float* __restrict__ bn_m,\n const float* __restrict__ bn_v, float* __restrict__ y,\n int N, int IN, int OUT, int act, int has_bias, int has_bn, float eps) {\n int idx = blockIdx.x;\n int n = idx / OUT;\n int o = idx - n*OUT;\n float acc = 0.0f;\n const float* xp = x + n*IN;\n const float* wp = w + o*IN;\n for (int i = threadIdx.x; i < IN; i += blockDim.x) acc += xp[i] * wp[i];\n __shared__ float smem[256];\n smem[threadIdx.x] = acc;\n __syncthreads();\n for (int s=128; s>0; s>>=1) { if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x+s]; __syncthreads(); }\n if (threadIdx.x == 0) {\n float v = smem[0];\n if (has_bias) v += bias[o];\n if (has_bn) v = (v - bn_m[o]) * rsqrtf(bn_v[o] + eps) * bn_w[o] + bn_b[o];\n if (act == 1) v = v > 0.0f ? v : 0.0f;\n else if (act == 2) v = 1.0f / (1.0f + expf(-v));\n y[idx] = v;\n }\n}\n\n__global__ void depthwise_kernel(const float* __restrict__ x, const float* __restrict__ w,\n const float* __restrict__ bn_w, const float* __restrict__ bn_b,\n const float* __restrict__ bn_m, const float* __restrict__ bn_v,\n float* __restrict__ y, int total, float eps) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total) {\n int c = i % total; // overwritten below for N*C with C known not passed? no\n }\n}\n\n__global__ void depthwise_kernel2(const float* __restrict__ x, const float* __restrict__ w,\n const float* __restrict__ bn_w, const float* __restrict__ bn_b,\n const float* __restrict__ bn_m, const float* __restrict__ bn_v,\n float* __restrict__ y, int N, int C, float eps) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n int total = N*C;\n if (i < total) {\n int c = i - (i / C) * C;\n float v = x[i] * w[c*9 + 4];\n v = (v - bn_m[c]) * rsqrtf(bn_v[c] + eps) * bn_w[c] + bn_b[c];\n y[i] = v > 0.0f ? v : 0.0f;\n }\n}\n\nstatic inline torch::Tensor linear(torch::Tensor x, torch::Tensor w, torch::Tensor bias,\n torch::Tensor bw, torch::Tensor bb, torch::Tensor bm, torch::Tensor bv,\n int act, int has_bias, int has_bn, double eps) {\n int64_t N = x.size(0), IN = x.size(1), OUT = w.size(0);\n auto y = torch::empty({N, OUT}, x.options());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n linear_kernel<<<N*OUT, 256, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(),\n has_bias?bias.data_ptr<float>():nullptr, has_bn?bw.data_ptr<float>():nullptr,\n has_bn?bb.data_ptr<float>():nullptr, has_bn?bm.data_ptr<float>():nullptr,\n has_bn?bv.data_ptr<float>():nullptr, y.data_ptr<float>(), (int)N, (int)IN, (int)OUT,\n act, has_bias, has_bn, (float)eps);\n return y;\n}\n\nstatic inline torch::Tensor depthwise(torch::Tensor x, torch::Tensor w, torch::Tensor bw, torch::Tensor bb, torch::Tensor bm, torch::Tensor bv, double eps) {\n int64_t N = x.size(0), C = x.size(1);\n auto y = torch::empty_like(x);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n int total = (int)(N*C);\n depthwise_kernel2<<<(total+255)/256, 256, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), bw.data_ptr<float>(), bb.data_ptr<float>(), bm.data_ptr<float>(), bv.data_ptr<float>(), y.data_ptr<float>(), (int)N, (int)C, (float)eps);\n return y;\n}\n\ntorch::Tensor tail_forward(torch::Tensor x, std::vector<torch::Tensor> p, double eps) {\n int64_t N=x.size(0), C=x.size(1), HW=x.size(2)*x.size(3);\n auto v = torch::empty({N,C}, x.options());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n pool_kernel<<<N*C, 256, 0, stream>>>(x.data_ptr<float>(), v.data_ptr<float>(), (int)N, (int)C, (int)HW);\n int k=0;\n torch::Tensor empty;\n // mbconv1 se and output\n v = linear(v, p[k++], empty, empty, empty, empty, empty, 1, 0, 0, eps);\n v = linear(v, p[k++], empty, empty, empty, empty, empty, 2, 0, 0, eps);\n v = linear(v, p[k++], empty, p[k], p[k+1], p[k+2], p[k+3], 0, 0, 1, eps); k+=4;\n // four remaining blocks\n for (int b=0; b<4; ++b) {\n v = linear(v, p[k++], empty, p[k], p[k+1], p[k+2], p[k+3], 1, 0, 1, eps); k+=4;\n v = depthwise(v, p[k++], p[k], p[k+1], p[k+2], p[k+3], eps); k+=4;\n v = linear(v, p[k++], empty, empty, empty, empty, empty, 1, 0, 0, eps);\n v = linear(v, p[k++], empty, empty, empty, empty, empty, 2, 0, 0, eps);\n v = linear(v, p[k++], empty, p[k], p[k+1], p[k+2], p[k+3], 0, 0, 1, eps); k+=4;\n }\n // final conv+bn+relu\n v = linear(v, p[k++], empty, p[k], p[k+1], p[k+2], p[k+3], 1, 0, 1, eps); k+=4;\n // fc with bias\n v = linear(v, p[k++], p[k++], empty, empty, empty, empty, 0, 1, 0, eps);\n return v;\n}\n'''\n\n_ext = load_inline(name='effb2_tail_ext_v1', cpp_sources=[cpp_src], cuda_sources=[cuda_src], functions=['tail_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.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)\n self.bn1 = nn.BatchNorm2d(32)\n self.relu = nn.ReLU(inplace=True)\n self.mbconv1 = self._make_mbconv_block(32, 96, 1, 3)\n self.mbconv2 = self._make_mbconv_block(96, 144, 2, 6)\n self.mbconv3 = self._make_mbconv_block(144, 192, 2, 6)\n self.mbconv4 = self._make_mbconv_block(192, 288, 2, 6)\n self.mbconv5 = self._make_mbconv_block(288, 384, 1, 6)\n self.conv_final = nn.Conv2d(384, 1408, kernel_size=1, stride=1, padding=0, bias=False)\n self.bn_final = nn.BatchNorm2d(1408)\n self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n self.fc = nn.Linear(1408, num_classes)\n def _make_mbconv_block(self, in_channels, out_channels, stride, expand_ratio):\n layers=[]; expanded_channels=in_channels*expand_ratio\n if expand_ratio != 1:\n layers += [nn.Conv2d(in_channels, expanded_channels, 1, 1, 0, bias=False), nn.BatchNorm2d(expanded_channels), nn.ReLU(inplace=True)]\n layers += [nn.Conv2d(expanded_channels, expanded_channels, 3, stride, 1, groups=expanded_channels, bias=False), nn.BatchNorm2d(expanded_channels), nn.ReLU(inplace=True)]\n layers += [nn.AdaptiveAvgPool2d((1,1)), nn.Conv2d(expanded_channels, expanded_channels//4, 1, bias=False), nn.ReLU(inplace=True), nn.Conv2d(expanded_channels//4, expanded_channels, 1, bias=False), nn.Sigmoid()]\n layers += [nn.Conv2d(expanded_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels)]\n return nn.Sequential(*layers)\n def _params(self):\n p=[]\n b=self.mbconv1\n p += [b[7].weight, b[9].weight, b[11].weight, b[12].weight, b[12].bias, b[12].running_mean, b[12].running_var]\n for b in [self.mbconv2,self.mbconv3,self.mbconv4,self.mbconv5]:\n p += [b[0].weight, b[1].weight, b[1].bias, b[1].running_mean, b[1].running_var]\n p += [b[3].weight, b[4].weight, b[4].bias, b[4].running_mean, b[4].running_var]\n p += [b[7].weight, b[9].weight, b[11].weight, b[12].weight, b[12].bias, b[12].running_mean, b[12].running_var]\n p += [self.conv_final.weight, self.bn_final.weight, self.bn_final.bias, self.bn_final.running_mean, self.bn_final.running_var]\n p += [self.fc.weight, self.fc.bias]\n return p\n def forward(self, x):\n x = self.relu(self.bn1(self.conv1(x)))\n b=self.mbconv1\n x = b[2](b[1](b[0](x)))\n x = b[5](b[4](b[3](x)))\n return _ext.tail_forward(x, self._params(), 1e-5)\n"
}