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'''\n#include <torch/extension.h>\ntorch::Tensor avgpool7(torch::Tensor x);\ntorch::Tensor bn_relu_train(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void avgpool7_kernel(const float* __restrict__ x, float* __restrict__ y, int NC) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= NC) return;\n const float* p = x + ((long long)idx) * 49LL;\n float s = 0.0f;\n #pragma unroll\n for (int i = 0; i < 49; ++i) s += p[i];\n y[idx] = s * (1.0f / 49.0f);\n}\n\ntorch::Tensor avgpool7(torch::Tensor x) {\n auto xc = x.contiguous();\n int64_t N = xc.size(0), C = xc.size(1);\n auto y = torch::empty({N, C}, xc.options());\n int NC = (int)(N * C);\n int threads = 256;\n int blocks = (NC + threads - 1) / threads;\n avgpool7_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(xc.data_ptr<float>(), y.data_ptr<float>(), NC);\n return y;\n}\n\n__global__ void bn_relu_train_kernel(const float* __restrict__ x, const float* __restrict__ gamma,\n const float* __restrict__ beta, float* __restrict__ y,\n int N, int C, int HW, float eps) {\n int c = blockIdx.x;\n int tid = threadIdx.x;\n extern __shared__ float sh[];\n float* ss = sh;\n float* ss2 = sh + blockDim.x;\n float sum = 0.0f;\n float sum2 = 0.0f;\n int M = N * HW;\n for (int t = tid; t < M; t += blockDim.x) {\n int n = t / HW;\n int s = t - n * HW;\n float v = x[((long long)n * C + c) * HW + s];\n sum += v;\n sum2 += v * v;\n }\n ss[tid] = sum;\n ss2[tid] = sum2;\n __syncthreads();\n for (int off = blockDim.x >> 1; off > 0; off >>= 1) {\n if (tid < off) {\n ss[tid] += ss[tid + off];\n ss2[tid] += ss2[tid + off];\n }\n __syncthreads();\n }\n float mean = ss[0] / (float)M;\n float var = ss2[0] / (float)M - mean * mean;\n var = var < 0.0f ? 0.0f : var;\n float inv = rsqrtf(var + eps);\n float g = gamma[c];\n float be = beta[c];\n for (int t = tid; t < M; t += blockDim.x) {\n int n = t / HW;\n int s = t - n * HW;\n long long idx = ((long long)n * C + c) * HW + s;\n float z = (x[idx] - mean) * inv * g + be;\n y[idx] = z > 0.0f ? z : 0.0f;\n }\n}\n\ntorch::Tensor bn_relu_train(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps) {\n TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, \"x must be cuda float\");\n auto xc = x.contiguous();\n int N = (int)xc.size(0), C = (int)xc.size(1), H = (int)xc.size(2), W = (int)xc.size(3);\n int HW = H * W;\n auto y = torch::empty_like(xc);\n int threads = 256;\n size_t shmem = threads * 2 * sizeof(float);\n bn_relu_train_kernel<<<C, threads, shmem, at::cuda::getCurrentCUDAStream()>>>(\n xc.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), y.data_ptr<float>(), N, C, HW, (float)eps);\n return y;\n}\n'''\n\n_ext = load_inline(\n name='mobilenetv1_fused_ext_v3',\n cpp_sources=_cpp_src,\n cuda_sources=_cuda_src,\n functions=['avgpool7','bn_relu_train'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, num_classes=1000, input_channels=3, alpha=1.0):\n super(ModelNew, self).__init__()\n torch.backends.cudnn.benchmark = True\n def conv_bn(inp, oup, stride):\n return nn.Sequential(nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup), nn.ReLU(inplace=True))\n def conv_dw(inp, oup, stride):\n return nn.Sequential(nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), nn.BatchNorm2d(inp), nn.ReLU(inplace=True), nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), nn.ReLU(inplace=True))\n self.model = nn.Sequential(\n conv_bn(input_channels, int(32 * alpha), 2), conv_dw(int(32 * alpha), int(64 * alpha), 1), conv_dw(int(64 * alpha), int(128 * alpha), 2), conv_dw(int(128 * alpha), int(128 * alpha), 1), conv_dw(int(128 * alpha), int(256 * alpha), 2), conv_dw(int(256 * alpha), int(256 * alpha), 1), conv_dw(int(256 * alpha), int(512 * alpha), 2), conv_dw(int(512 * alpha), int(512 * alpha), 1), conv_dw(int(512 * alpha), int(512 * alpha), 1), conv_dw(int(512 * alpha), int(512 * alpha), 1), conv_dw(int(512 * alpha), int(512 * alpha), 1), conv_dw(int(512 * alpha), int(512 * alpha), 1), conv_dw(int(512 * alpha), int(1024 * alpha), 2), conv_dw(int(1024 * alpha), int(1024 * alpha), 1), nn.AvgPool2d(7))\n self.fc = nn.Linear(int(1024 * alpha), num_classes)\n self._folded = None; self._fold_key = None\n def _apply(self, fn):\n ret = super()._apply(fn); self._folded = None; self._fold_key = None; return ret\n def load_state_dict(self, state_dict, strict=True, assign=False):\n ret = super().load_state_dict(state_dict, strict=strict, assign=assign); self._folded = None; self._fold_key = None; return ret\n def _fold_conv_bn(self, conv, bn):\n invstd = torch.rsqrt(bn.running_var + bn.eps); scale = bn.weight * invstd\n return (conv.weight * scale.reshape([-1,1,1,1])).contiguous(), (bn.bias - bn.running_mean * scale).contiguous()\n def _ensure_folded(self):\n key = (self.fc.weight.device, self.fc.weight.dtype)\n if self._folded is not None and self._fold_key == key: return self._folded\n folded=[]\n with torch.no_grad():\n folded.append(self._fold_conv_bn(self.model[0][0], self.model[0][1]))\n for i in range(1,14):\n block=self.model[i]; folded.append(self._fold_conv_bn(block[0],block[1])); folded.append(self._fold_conv_bn(block[3],block[4]))\n self._folded=folded; self._fold_key=key; return folded\n def _bn_relu(self, x, bn):\n return _ext.bn_relu_train(x, bn.weight, bn.bias, bn.eps)\n def forward(self, x):\n if self.training:\n x = F.conv2d(x, self.model[0][0].weight, None, stride=2, padding=1)\n x = self._bn_relu(x, self.model[0][1])\n for i in range(1,14):\n block=self.model[i]; dw=block[0]\n x = F.conv2d(x, dw.weight, None, stride=dw.stride, padding=1, groups=dw.groups)\n x = self._bn_relu(x, block[1])\n x = F.conv2d(x, block[3].weight, None, stride=1, padding=0)\n x = self._bn_relu(x, block[4])\n else:\n folded=self._ensure_folded(); j=0\n w,b=folded[j]; j+=1; x=F.conv2d(x,w,b,stride=2,padding=1); x=F.relu(x,inplace=True)\n for i in range(1,14):\n block=self.model[i]; dw=block[0]\n w,b=folded[j]; j+=1; x=F.conv2d(x,w,b,stride=dw.stride,padding=1,groups=dw.groups); x=F.relu(x,inplace=True)\n w,b=folded[j]; j+=1; x=F.conv2d(x,w,b,stride=1,padding=0); x=F.relu(x,inplace=True)\n x = _ext.avgpool7(x)\n return F.linear(x, self.fc.weight, self.fc.bias)\n"
}