callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\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\n__global__ void conv3x3_c8_oc16_row_kernel(\n const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n float* __restrict__ y,\n int N, int H, int W, int K, float inv_div)\n{\n // specialized for C=8, kernel=3, stride=1, no padding. One block computes one output row\n // for one batch item and a group of 16 output channels. Threads cover output columns.\n const int OH = H - 2;\n const int OW = W - 2;\n const int ow = threadIdx.x;\n if (ow >= OW) return;\n const int oh = blockIdx.x;\n const int z = blockIdx.y;\n const int ocg = blockIdx.z;\n const int n = z;\n const int oc0 = ocg * 16;\n\n float a0=0.f,a1=0.f,a2=0.f,a3=0.f,a4=0.f,a5=0.f,a6=0.f,a7=0.f;\n float a8=0.f,a9=0.f,a10=0.f,a11=0.f,a12=0.f,a13=0.f,a14=0.f,a15=0.f;\n\n #pragma unroll\n for (int c=0; c<8; ++c) {\n #pragma unroll\n for (int r=0; r<3; ++r) {\n const float* xp = x + ((n*8 + c)*H + (oh+r))*W + ow;\n #pragma unroll\n for (int s=0; s<3; ++s) {\n float xv = xp[s];\n int kk = (c*3 + r)*3 + s;\n const float* wp = w + oc0*72 + kk;\n a0 = fmaf(xv, wp[0*72], a0);\n a1 = fmaf(xv, wp[1*72], a1);\n a2 = fmaf(xv, wp[2*72], a2);\n a3 = fmaf(xv, wp[3*72], a3);\n a4 = fmaf(xv, wp[4*72], a4);\n a5 = fmaf(xv, wp[5*72], a5);\n a6 = fmaf(xv, wp[6*72], a6);\n a7 = fmaf(xv, wp[7*72], a7);\n a8 = fmaf(xv, wp[8*72], a8);\n a9 = fmaf(xv, wp[9*72], a9);\n a10 = fmaf(xv, wp[10*72], a10);\n a11 = fmaf(xv, wp[11*72], a11);\n a12 = fmaf(xv, wp[12*72], a12);\n a13 = fmaf(xv, wp[13*72], a13);\n a14 = fmaf(xv, wp[14*72], a14);\n a15 = fmaf(xv, wp[15*72], a15);\n }\n }\n }\n\n const int plane = OH * OW;\n float* yp = y + (n*K)*plane + oh*OW + ow;\n #define STORE(i, acc) do { if (oc0 + (i) < K) { float v = ((acc) + b[oc0+(i)]) * inv_div; v = v >= 0.f ? v : v * 0.01f; yp[((oc0+(i))*plane)] = v; } } while(0)\n STORE(0,a0); STORE(1,a1); STORE(2,a2); STORE(3,a3);\n STORE(4,a4); STORE(5,a5); STORE(6,a6); STORE(7,a7);\n STORE(8,a8); STORE(9,a9); STORE(10,a10); STORE(11,a11);\n STORE(12,a12); STORE(13,a13); STORE(14,a14); STORE(15,a15);\n #undef STORE\n}\n\n__global__ void conv_generic_kernel(\n const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b,\n float* __restrict__ y, int N, int C, int H, int W, int K, int R, int S, float inv_div)\n{\n int OH = H - R + 1, OW = W - S + 1;\n long long total = (long long)N * K * OH * OW;\n long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n int ow = idx % OW; idx /= OW;\n int oh = idx % OH; idx /= OH;\n int oc = idx % K; int n = idx / K;\n float acc = 0.f;\n for (int c=0; c<C; ++c)\n for (int r=0; r<R; ++r)\n for (int s=0; s<S; ++s)\n acc = fmaf(x[((n*C+c)*H + oh+r)*W + ow+s], w[((oc*C+c)*R+r)*S+s], acc);\n float v = (acc + b[oc]) * inv_div;\n y[((n*K+oc)*OH+oh)*OW+ow] = v >= 0.f ? v : v * 0.01f;\n}\n\ntorch::Tensor conv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double divisor) {\n TORCH_CHECK(x.is_cuda(), \"x must be CUDA\");\n TORCH_CHECK(weight.is_cuda(), \"weight must be CUDA\");\n TORCH_CHECK(x.scalar_type() == torch::kFloat32, \"only float32 supported\");\n x = x.contiguous();\n weight = weight.contiguous();\n bias = bias.contiguous();\n const int N = (int)x.size(0);\n const int C = (int)x.size(1);\n const int H = (int)x.size(2);\n const int W = (int)x.size(3);\n const int K = (int)weight.size(0);\n const int R = (int)weight.size(2);\n const int S = (int)weight.size(3);\n const int OH = H - R + 1;\n const int OW = W - S + 1;\n auto y = torch::empty({N, K, OH, OW}, x.options());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n float inv = 1.0f / (float)divisor;\n if (C == 8 && R == 3 && S == 3 && OW <= 128 && K <= 1024) {\n dim3 block(128);\n dim3 grid(OH, N, (K + 15) / 16);\n conv3x3_c8_oc16_row_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), N, H, W, K, inv);\n } else {\n long long total = (long long)N * K * OH * OW;\n int threads = 256;\n int blocks = (int)((total + threads - 1) / threads);\n conv_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), N,C,H,W,K,R,S,inv);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_div_lrelu_ext_v1', cpp_sources='', cuda_sources=cuda_src, functions=['conv_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, divisor):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.divisor = divisor\n def forward(self, x):\n return _ext.conv_forward(x, self.conv.weight, self.conv.bias, float(self.divisor))\n"
}