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\ntorch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b,\n int stride, int padding, int dilation, int groups, bool has_bias);\n\ntorch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b,\n int stride, int padding, int dilation, int groups, bool has_bias) {\n return conv2d_forward_cuda(x, w, b, stride, padding, dilation, groups, has_bias);\n}\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#ifndef __CUDA_ARCH__\n#define __CUDA_ARCH__ 0\n#endif\n\n__global__ __launch_bounds__(64, 4)\nvoid conv3x3_s1p0_c16_oc128_kernel(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 has_bias) {\n // fixed: C=16, OC=128, K=3, stride=1, pad=0\n const int Hout = H - 2;\n const int Wout = W - 2;\n const int tx = threadIdx.x; // 0..15\n const int ty = threadIdx.y; // 0..3\n const int ow = blockIdx.x * 16 + tx;\n const int oh_tile = blockIdx.y & 255; // for Hout up to 1024 -> 256 tiles of height 4\n const int n = blockIdx.y >> 8;\n const int oh = oh_tile * 4 + ty;\n const int oc0 = blockIdx.z * 16;\n if (n >= N || oh >= Hout || ow >= Wout) return;\n\n float acc0=0.f, acc1=0.f, acc2=0.f, acc3=0.f;\n float acc4=0.f, acc5=0.f, acc6=0.f, acc7=0.f;\n float acc8=0.f, acc9=0.f, acc10=0.f, acc11=0.f;\n float acc12=0.f, acc13=0.f, acc14=0.f, acc15=0.f;\n if (has_bias) {\n acc0=b[oc0+0]; acc1=b[oc0+1]; acc2=b[oc0+2]; acc3=b[oc0+3];\n acc4=b[oc0+4]; acc5=b[oc0+5]; acc6=b[oc0+6]; acc7=b[oc0+7];\n acc8=b[oc0+8]; acc9=b[oc0+9]; acc10=b[oc0+10]; acc11=b[oc0+11];\n acc12=b[oc0+12]; acc13=b[oc0+13]; acc14=b[oc0+14]; acc15=b[oc0+15];\n }\n\n const int in_base_n = n * 16 * H * W;\n #pragma unroll\n for (int c = 0; c < 16; ++c) {\n const float* xp = x + in_base_n + c * H * W + oh * W + ow;\n #pragma unroll\n for (int kh = 0; kh < 3; ++kh) {\n #pragma unroll\n for (int kw = 0; kw < 3; ++kw) {\n float xv = __ldg(xp + kh * W + kw);\n int wi = ((oc0 * 16 + c) * 3 + kh) * 3 + kw;\n acc0 = fmaf(xv, __ldg(w + wi + 0*16*3*3), acc0);\n acc1 = fmaf(xv, __ldg(w + wi + 1*16*3*3), acc1);\n acc2 = fmaf(xv, __ldg(w + wi + 2*16*3*3), acc2);\n acc3 = fmaf(xv, __ldg(w + wi + 3*16*3*3), acc3);\n acc4 = fmaf(xv, __ldg(w + wi + 4*16*3*3), acc4);\n acc5 = fmaf(xv, __ldg(w + wi + 5*16*3*3), acc5);\n acc6 = fmaf(xv, __ldg(w + wi + 6*16*3*3), acc6);\n acc7 = fmaf(xv, __ldg(w + wi + 7*16*3*3), acc7);\n acc8 = fmaf(xv, __ldg(w + wi + 8*16*3*3), acc8);\n acc9 = fmaf(xv, __ldg(w + wi + 9*16*3*3), acc9);\n acc10 = fmaf(xv, __ldg(w + wi + 10*16*3*3), acc10);\n acc11 = fmaf(xv, __ldg(w + wi + 11*16*3*3), acc11);\n acc12 = fmaf(xv, __ldg(w + wi + 12*16*3*3), acc12);\n acc13 = fmaf(xv, __ldg(w + wi + 13*16*3*3), acc13);\n acc14 = fmaf(xv, __ldg(w + wi + 14*16*3*3), acc14);\n acc15 = fmaf(xv, __ldg(w + wi + 15*16*3*3), acc15);\n }\n }\n }\n\n const long spatial = (long)Hout * Wout;\n const long out_base = ((long)n * 128 + oc0) * spatial + (long)oh * Wout + ow;\n y[out_base + 0*spatial] = acc0;\n y[out_base + 1*spatial] = acc1;\n y[out_base + 2*spatial] = acc2;\n y[out_base + 3*spatial] = acc3;\n y[out_base + 4*spatial] = acc4;\n y[out_base + 5*spatial] = acc5;\n y[out_base + 6*spatial] = acc6;\n y[out_base + 7*spatial] = acc7;\n y[out_base + 8*spatial] = acc8;\n y[out_base + 9*spatial] = acc9;\n y[out_base + 10*spatial] = acc10;\n y[out_base + 11*spatial] = acc11;\n y[out_base + 12*spatial] = acc12;\n y[out_base + 13*spatial] = acc13;\n y[out_base + 14*spatial] = acc14;\n y[out_base + 15*spatial] = acc15;\n}\n\n__global__ void conv_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,\n const float* __restrict__ b, float* __restrict__ y,\n int N, int C, int H, int W, int OC, int K,\n int Hout, int Wout, int stride, int pad, int dil, int groups, int has_bias) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n long total = (long)N * OC * Hout * Wout;\n if (idx >= total) return;\n int ow = idx % Wout;\n long t = idx / Wout;\n int oh = t % Hout;\n t /= Hout;\n int oc = t % OC;\n int n = t / OC;\n int ocpg = OC / groups;\n int cpg = C / groups;\n int g = oc / ocpg;\n float acc = has_bias ? b[oc] : 0.f;\n for (int ci = 0; ci < cpg; ++ci) {\n int c = g * cpg + ci;\n for (int kh = 0; kh < K; ++kh) {\n int ih = oh * stride - pad + kh * dil;\n if ((unsigned)ih >= (unsigned)H) continue;\n for (int kw = 0; kw < K; ++kw) {\n int iw = ow * stride - pad + kw * dil;\n if ((unsigned)iw >= (unsigned)W) continue;\n float xv = x[((long)n*C + c)*H*W + (long)ih*W + iw];\n float wv = w[(((long)oc*cpg + ci)*K + kh)*K + kw];\n acc = fmaf(xv, wv, acc);\n }\n }\n }\n y[idx] = acc;\n}\n\ntorch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b,\n int stride, int padding, int dilation, int groups, bool has_bias) {\n x = x.contiguous();\n w = w.contiguous();\n const int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);\n const int OC = w.size(0), K = w.size(2);\n const int Hout = (H + 2*padding - dilation*(K-1) - 1) / stride + 1;\n const int Wout = (W + 2*padding - dilation*(K-1) - 1) / stride + 1;\n auto y = torch::empty({N, OC, Hout, Wout}, x.options());\n if (C == 16 && OC == 128 && K == 3 && stride == 1 && padding == 0 && dilation == 1 && groups == 1 && Hout > 0 && Wout > 0 && Hout <= 1024) {\n dim3 block(16,4,1);\n dim3 grid((Wout + 15)/16, ((Hout + 3)/4) * N, 8);\n conv3x3_s1p0_c16_oc128_kernel<<<grid, block>>>(x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? b.data_ptr<float>() : nullptr, y.data_ptr<float>(), N, H, W, has_bias ? 1 : 0);\n } else {\n int threads = 256;\n long total = (long)N * OC * Hout * Wout;\n int blocks = (int)((total + threads - 1) / threads);\n conv_generic_kernel<<<blocks, threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? b.data_ptr<float>() : nullptr, y.data_ptr<float>(), N, C, H, W, OC, K, Hout, Wout, stride, padding, dilation, groups, has_bias ? 1 : 0);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv2d_custom_ext_v1', cpp_sources=cpp_src, cuda_sources=cuda_src,\n functions=['conv2d_forward'], with_cuda=True,\n extra_cflags=['-O3'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):\n super(ModelNew, self).__init__()\n self.conv2d = nn.Conv2d(in_channels, out_channels, (kernel_size, kernel_size), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)\n self.stride = stride\n self.padding = padding\n self.dilation = dilation\n self.groups = groups\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n b = self.conv2d.bias if self.conv2d.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)\n return _ext.conv2d_forward(x, self.conv2d.weight, b, self.stride, self.padding, self.dilation, self.groups, self.conv2d.bias is not None)\n"
}