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'''\n#include <torch/extension.h>\ntorch::Tensor bn_train_relu(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps);\ntorch::Tensor bn2_train_add_relu(torch::Tensor a, torch::Tensor d,\n torch::Tensor ga, torch::Tensor ba,\n torch::Tensor gd, torch::Tensor bd,\n double epsa, double epsd);\ntorch::Tensor bn_relu_eval(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,\n torch::Tensor mean, torch::Tensor var, double eps);\ntorch::Tensor bn_add_relu_eval(torch::Tensor out, torch::Tensor id,\n torch::Tensor g2, torch::Tensor b2, torch::Tensor m2, torch::Tensor v2,\n torch::Tensor gd, torch::Tensor bd, torch::Tensor md, torch::Tensor vd,\n double eps2, double epsd);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void stats1_kernel(const float* __restrict__ x, float* __restrict__ stats, int N, int C, int HW, float eps){\n int c = blockIdx.x;\n extern __shared__ float sh[];\n float* ss = sh;\n float* sq = sh + blockDim.x;\n float sum = 0.f, sumsq = 0.f;\n int M = N * HW;\n for (int k = threadIdx.x; k < M; k += blockDim.x) {\n int n = k / HW;\n int hw = k - n * HW;\n float v = x[(n*C + c)*HW + hw];\n sum += v;\n sumsq += v*v;\n }\n ss[threadIdx.x] = sum; sq[threadIdx.x] = sumsq;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (threadIdx.x < s) { ss[threadIdx.x] += ss[threadIdx.x+s]; sq[threadIdx.x] += sq[threadIdx.x+s]; }\n __syncthreads();\n }\n if (threadIdx.x == 0) {\n float mean = ss[0] / (float)M;\n float var = fmaxf(sq[0] / (float)M - mean * mean, 0.0f);\n stats[c] = mean;\n stats[C+c] = rsqrtf(var + eps);\n }\n}\n\n__global__ void stats2_kernel(const float* __restrict__ a, const float* __restrict__ d, float* __restrict__ stats, int N, int C, int HW, float epsa, float epsd){\n int c = blockIdx.x;\n extern __shared__ float sh[];\n float* sa = sh;\n float* sqa = sh + blockDim.x;\n float* sd = sh + 2*blockDim.x;\n float* sqd = sh + 3*blockDim.x;\n float suma = 0.f, sumqa = 0.f, sumd = 0.f, sumqd = 0.f;\n int M = N * HW;\n for (int k = threadIdx.x; k < M; k += blockDim.x) {\n int n = k / HW;\n int hw = k - n * HW;\n int idx = (n*C + c)*HW + hw;\n float va = a[idx];\n float vd = d[idx];\n suma += va; sumqa += va*va;\n sumd += vd; sumqd += vd*vd;\n }\n sa[threadIdx.x] = suma; sqa[threadIdx.x] = sumqa; sd[threadIdx.x] = sumd; sqd[threadIdx.x] = sumqd;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n sa[threadIdx.x] += sa[threadIdx.x+s]; sqa[threadIdx.x] += sqa[threadIdx.x+s];\n sd[threadIdx.x] += sd[threadIdx.x+s]; sqd[threadIdx.x] += sqd[threadIdx.x+s];\n }\n __syncthreads();\n }\n if (threadIdx.x == 0) {\n float invM = 1.0f / (float)M;\n float meana = sa[0] * invM;\n float vara = fmaxf(sqa[0] * invM - meana*meana, 0.0f);\n float meand = sd[0] * invM;\n float vard = fmaxf(sqd[0] * invM - meand*meand, 0.0f);\n stats[c] = meana; stats[C+c] = rsqrtf(vara + epsa);\n stats[2*C+c] = meand; stats[3*C+c] = rsqrtf(vard + epsd);\n }\n}\n\n__global__ void norm_relu_kernel(const float* __restrict__ x, float* __restrict__ y,\n const float* __restrict__ gamma, const float* __restrict__ beta,\n const float* __restrict__ stats, int total, int C, int HW){\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < total; i += stride) {\n int c = (i / HW) % C;\n float v = (x[i] - stats[c]) * stats[C+c] * gamma[c] + beta[c];\n y[i] = v > 0.f ? v : 0.f;\n }\n}\n\n__global__ void norm2_add_relu_kernel(const float* __restrict__ a, const float* __restrict__ d, float* __restrict__ y,\n const float* __restrict__ ga, const float* __restrict__ ba,\n const float* __restrict__ gd, const float* __restrict__ bd,\n const float* __restrict__ stats, int total, int C, int HW){\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < total; i += stride) {\n int c = (i / HW) % C;\n float va = (a[i] - stats[c]) * stats[C+c] * ga[c] + ba[c];\n float vd = (d[i] - stats[2*C+c]) * stats[3*C+c] * gd[c] + bd[c];\n float s = va + vd;\n y[i] = s > 0.f ? s : 0.f;\n }\n}\n\n__global__ void bn_relu_eval_kernel(const float* __restrict__ x, float* __restrict__ y,\n const float* __restrict__ gamma, const float* __restrict__ beta,\n const float* __restrict__ mean, const float* __restrict__ var,\n int total, int C, int HW, float eps) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < total; i += stride) {\n int c = (i / HW) % C;\n float v = (x[i] - mean[c]) * rsqrtf(var[c] + eps) * gamma[c] + beta[c];\n y[i] = v > 0.0f ? v : 0.0f;\n }\n}\n\n__global__ void bn_add_relu_eval_kernel(const float* __restrict__ out, const float* __restrict__ id,\n float* __restrict__ y,\n const float* __restrict__ g2, const float* __restrict__ b2,\n const float* __restrict__ m2, const float* __restrict__ v2,\n const float* __restrict__ gd, const float* __restrict__ bd,\n const float* __restrict__ md, const float* __restrict__ vd,\n int total, int C, int HW, float eps2, float epsd) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (int i = idx; i < total; i += stride) {\n int c = (i / HW) % C;\n float a = (out[i] - m2[c]) * rsqrtf(v2[c] + eps2) * g2[c] + b2[c];\n float d = (id[i] - md[c]) * rsqrtf(vd[c] + epsd) * gd[c] + bd[c];\n float s = a + d;\n y[i] = s > 0.0f ? s : 0.0f;\n }\n}\n\ntorch::Tensor bn_train_relu(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps) {\n int N=x.size(0), C=x.size(1), HW=x.size(2)*x.size(3); int64_t total64=x.numel();\n auto y=torch::empty_like(x);\n auto stats=torch::empty({2*C}, x.options());\n int threads=256;\n stats1_kernel<<<C, threads, 2*threads*sizeof(float)>>>(x.data_ptr<float>(), stats.data_ptr<float>(), N,C,HW,(float)eps);\n int blocks=(int)std::min<int64_t>((total64+threads-1)/threads, 4096);\n norm_relu_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), stats.data_ptr<float>(), (int)total64, C, HW);\n return y;\n}\n\ntorch::Tensor bn2_train_add_relu(torch::Tensor a, torch::Tensor d, torch::Tensor ga, torch::Tensor ba, torch::Tensor gd, torch::Tensor bd, double epsa, double epsd) {\n int N=a.size(0), C=a.size(1), HW=a.size(2)*a.size(3); int64_t total64=a.numel();\n auto y=torch::empty_like(a);\n auto stats=torch::empty({4*C}, a.options());\n int threads=256;\n stats2_kernel<<<C, threads, 4*threads*sizeof(float)>>>(a.data_ptr<float>(), d.data_ptr<float>(), stats.data_ptr<float>(), N,C,HW,(float)epsa,(float)epsd);\n int blocks=(int)std::min<int64_t>((total64+threads-1)/threads, 4096);\n norm2_add_relu_kernel<<<blocks, threads>>>(a.data_ptr<float>(), d.data_ptr<float>(), y.data_ptr<float>(), ga.data_ptr<float>(), ba.data_ptr<float>(), gd.data_ptr<float>(), bd.data_ptr<float>(), stats.data_ptr<float>(), (int)total64, C, HW);\n return y;\n}\n\ntorch::Tensor bn_relu_eval(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, double eps) {\n auto y=torch::empty_like(x); int64_t total64=x.numel(); int C=x.size(1); int HW=x.size(2)*x.size(3); int threads=256; int blocks=(int)std::min<int64_t>((total64+threads-1)/threads,4096);\n bn_relu_eval_kernel<<<blocks,threads>>>(x.data_ptr<float>(),y.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),(int)total64,C,HW,(float)eps); return y;\n}\n\ntorch::Tensor bn_add_relu_eval(torch::Tensor out, torch::Tensor id, torch::Tensor g2, torch::Tensor b2, torch::Tensor m2, torch::Tensor v2, torch::Tensor gd, torch::Tensor bd, torch::Tensor md, torch::Tensor vd, double eps2, double epsd) {\n auto y=torch::empty_like(out); int64_t total64=out.numel(); int C=out.size(1); int HW=out.size(2)*out.size(3); int threads=256; int blocks=(int)std::min<int64_t>((total64+threads-1)/threads,4096);\n bn_add_relu_eval_kernel<<<blocks,threads>>>(out.data_ptr<float>(),id.data_ptr<float>(),y.data_ptr<float>(),g2.data_ptr<float>(),b2.data_ptr<float>(),m2.data_ptr<float>(),v2.data_ptr<float>(),gd.data_ptr<float>(),bd.data_ptr<float>(),md.data_ptr<float>(),vd.data_ptr<float>(),(int)total64,C,HW,(float)eps2,(float)epsd); return y;\n}\n'''\n_ext = load_inline(name='resblock_train_bn_ext_v1', cpp_sources=_cpp_src, cuda_sources=_cuda_src,\n functions=['bn_train_relu','bn2_train_add_relu','bn_relu_eval','bn_add_relu_eval'], verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])\n\nclass ModelNew(nn.Module):\n expansion=1\n def __init__(self, in_channels, out_channels, stride=1):\n super(ModelNew,self).__init__()\n self.conv1=nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=stride,padding=1,bias=False)\n self.bn1=nn.BatchNorm2d(out_channels)\n self.relu=nn.ReLU(inplace=True)\n self.conv2=nn.Conv2d(out_channels,out_channels,kernel_size=3,stride=1,padding=1,bias=False)\n self.bn2=nn.BatchNorm2d(out_channels)\n self.downsample=nn.Sequential(nn.Conv2d(in_channels,out_channels*self.expansion,kernel_size=1,stride=stride,bias=False), nn.BatchNorm2d(out_channels*self.expansion))\n self.stride=stride\n def forward(self,x):\n out=F.conv2d(x,self.conv1.weight,None,self.stride,1)\n if self.training:\n out=_ext.bn_train_relu(out,self.bn1.weight,self.bn1.bias,self.bn1.eps)\n out2=F.conv2d(out,self.conv2.weight,None,1,1)\n ident=F.conv2d(x,self.downsample[0].weight,None,self.stride,0)\n return _ext.bn2_train_add_relu(out2,ident,self.bn2.weight,self.bn2.bias,self.downsample[1].weight,self.downsample[1].bias,self.bn2.eps,self.downsample[1].eps)\n else:\n out=_ext.bn_relu_eval(out,self.bn1.weight,self.bn1.bias,self.bn1.running_mean,self.bn1.running_var,self.bn1.eps)\n out2=F.conv2d(out,self.conv2.weight,None,1,1)\n ident=F.conv2d(x,self.downsample[0].weight,None,self.stride,0)\n return _ext.bn_add_relu_eval(out2,ident,self.bn2.weight,self.bn2.bias,self.bn2.running_mean,self.bn2.running_var,self.downsample[1].weight,self.downsample[1].bias,self.downsample[1].running_mean,self.downsample[1].running_var,self.bn2.eps,self.downsample[1].eps)\n"
}