call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\n\nvoid conv_relu_bias_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias, torch::Tensor y);\n\ntorch::Tensor conv_relu_bias(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n auto N = x.size(0);\n auto H = x.size(2);\n auto W = x.size(3);\n auto Co = weight.size(0);\n auto R = weight.size(2);\n auto S = weight.size(3);\n auto Ho = H - R + 1;\n auto Wo = W - S + 1;\n auto y = torch::empty({N, Co, Ho, Wo}, x.options());\n conv_relu_bias_cuda(x, weight, conv_bias, add_bias, y);\n return y;\n}\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define GROUP 6\n\n__global__ __launch_bounds__(256, 2)\nvoid conv3x3_g6_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ cb,\n const float* __restrict__ ab,\n float* __restrict__ y,\n int N, int Ci, int Co, int H, int W, int Ho, int Wo) {\n int groups_w = Wo / GROUP;\n long long total = (long long)N * Co * Ho * groups_w;\n long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n if (tid >= total) return;\n\n int gw = tid % groups_w;\n long long t = tid / groups_w;\n int oh = t % Ho;\n t /= Ho;\n int co = t % Co;\n int n = t / Co;\n int ow = gw * GROUP;\n\n float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f, s4 = 0.f, s5 = 0.f;\n\n for (int ci = 0; ci < Ci; ++ci) {\n const float* xp0 = x + ((long long)n * Ci + ci) * H * W + oh * W + ow;\n const float* wp = w + ((long long)co * Ci + ci) * 9;\n #pragma unroll\n for (int r = 0; r < 3; ++r) {\n float w0 = __ldg(wp + r*3 + 0);\n float w1 = __ldg(wp + r*3 + 1);\n float w2 = __ldg(wp + r*3 + 2);\n const float* xp = xp0 + r * W;\n float x0 = __ldg(xp + 0);\n float x1 = __ldg(xp + 1);\n float x2 = __ldg(xp + 2);\n float x3 = __ldg(xp + 3);\n float x4 = __ldg(xp + 4);\n float x5 = __ldg(xp + 5);\n float x6 = __ldg(xp + 6);\n float x7 = __ldg(xp + 7);\n s0 = fmaf(x0, w0, s0); s0 = fmaf(x1, w1, s0); s0 = fmaf(x2, w2, s0);\n s1 = fmaf(x1, w0, s1); s1 = fmaf(x2, w1, s1); s1 = fmaf(x3, w2, s1);\n s2 = fmaf(x2, w0, s2); s2 = fmaf(x3, w1, s2); s2 = fmaf(x4, w2, s2);\n s3 = fmaf(x3, w0, s3); s3 = fmaf(x4, w1, s3); s3 = fmaf(x5, w2, s3);\n s4 = fmaf(x4, w0, s4); s4 = fmaf(x5, w1, s4); s4 = fmaf(x6, w2, s4);\n s5 = fmaf(x5, w0, s5); s5 = fmaf(x6, w1, s5); s5 = fmaf(x7, w2, s5);\n }\n }\n float b = __ldg(cb + co);\n float ba = __ldg(ab + co);\n long long out = ((long long)n * Co + co) * Ho * Wo + (long long)oh * Wo + ow;\n s0 += b; s1 += b; s2 += b; s3 += b; s4 += b; s5 += b;\n y[out+0] = (s0 > 0.f ? s0 : 0.f) + ba;\n y[out+1] = (s1 > 0.f ? s1 : 0.f) + ba;\n y[out+2] = (s2 > 0.f ? s2 : 0.f) + ba;\n y[out+3] = (s3 > 0.f ? s3 : 0.f) + ba;\n y[out+4] = (s4 > 0.f ? s4 : 0.f) + ba;\n y[out+5] = (s5 > 0.f ? s5 : 0.f) + ba;\n}\n\n__global__ void conv3x3_scalar_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ cb,\n const float* __restrict__ ab,\n float* __restrict__ y,\n int N, int Ci, int Co, int H, int W, int Ho, int Wo) {\n long long total = (long long)N * Co * Ho * Wo;\n long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n if (tid >= total) return;\n int ow = tid % Wo;\n long long t = tid / Wo;\n int oh = t % Ho; t /= Ho;\n int co = t % Co;\n int n = t / Co;\n float sum = 0.f;\n for (int ci=0; ci<Ci; ++ci) {\n for (int r=0; r<3; ++r) {\n const float* xp = x + ((long long)n*Ci+ci)*H*W + (oh+r)*W + ow;\n const float* wp = w + ((long long)co*Ci+ci)*9 + r*3;\n sum = fmaf(__ldg(xp+0), __ldg(wp+0), sum);\n sum = fmaf(__ldg(xp+1), __ldg(wp+1), sum);\n sum = fmaf(__ldg(xp+2), __ldg(wp+2), sum);\n }\n }\n sum += __ldg(cb+co);\n y[tid] = (sum > 0.f ? sum : 0.f) + __ldg(ab+co);\n}\n\nvoid conv_relu_bias_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias, torch::Tensor y) {\n int N=x.size(0), Ci=x.size(1), H=x.size(2), W=x.size(3);\n int Co=weight.size(0), Ho=y.size(2), Wo=y.size(3);\n int threads = 256;\n if (weight.size(2)==3 && weight.size(3)==3 && (Wo % GROUP)==0) {\n long long total = (long long)N * Co * Ho * (Wo / GROUP);\n int blocks = (int)((total + threads - 1) / threads);\n conv3x3_g6_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), y.data_ptr<float>(), N, Ci, Co, H, W, Ho, Wo);\n } else {\n long long total = (long long)N * Co * Ho * Wo;\n int blocks = (int)((total + threads - 1) / threads);\n conv3x3_scalar_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), y.data_ptr<float>(), N, Ci, Co, H, W, Ho, Wo);\n }\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_ext_v1', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['conv_relu_bias'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n def forward(self, x):\n return _ext.conv_relu_bias(x.contiguous(), self.conv.weight.contiguous(), self.conv.bias.contiguous(), self.bias.contiguous())\n"
}