call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = r'''\n#include <torch/extension.h>\n\ntorch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias,\n int stride_h, int stride_w, int pad_h, int pad_w,\n int dil_h, int dil_w, bool has_bias);\n'''\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <mma.h>\n\nusing namespace nvcuda;\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be CUDA\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n\nconstexpr int WM = 16;\nconstexpr int WN = 16;\nconstexpr int WK = 8;\nconstexpr int WARPS_M = 8;\nconstexpr int OC_TILES = 8; // 128 output channels / 16\nconstexpr int THREADS = WARPS_M * 32;\n\n__global__ __launch_bounds__(THREADS, 2)\nvoid conv2d_tf32_128oc_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ bias,\n float* __restrict__ y,\n int N, int C, int H, int W,\n int OC, int KH, int KW,\n int OH, int OW,\n int stride_h, int stride_w,\n int pad_h, int pad_w,\n int dil_h, int dil_w,\n int has_bias) {\n extern __shared__ float sh[];\n float* As = sh; // WARPS_M * 16 * 8\n float* Bs = As + WARPS_M * WM * WK; // OC_TILES * 16 * 8, col-major per tile: [tile][n][k]\n\n int tid = threadIdx.x;\n int warp = tid >> 5;\n int lane = tid & 31;\n int block_m0 = blockIdx.x * (WARPS_M * WM);\n int m0 = block_m0 + warp * WM;\n\n wmma::fragment<wmma::accumulator, WM, WN, WK, float> acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7;\n wmma::fill_fragment(acc0, 0.0f);\n wmma::fill_fragment(acc1, 0.0f);\n wmma::fill_fragment(acc2, 0.0f);\n wmma::fill_fragment(acc3, 0.0f);\n wmma::fill_fragment(acc4, 0.0f);\n wmma::fill_fragment(acc5, 0.0f);\n wmma::fill_fragment(acc6, 0.0f);\n wmma::fill_fragment(acc7, 0.0f);\n\n int Ktot = C * KH * KW;\n for (int kt = 0; kt < Ktot; kt += WK) {\n // Load B weights for all 128 OCs, K tile. Bs tile layout col_major with ld=8: addr = tile*128 + n*8 + k\n for (int idx = tid; idx < OC_TILES * WN * WK; idx += THREADS) {\n int kk = idx & 7;\n int tmp = idx >> 3;\n int nn = tmp & 15;\n int ot = tmp >> 4;\n int oc = ot * 16 + nn;\n int k = kt + kk;\n float val = 0.0f;\n if (oc < OC && k < Ktot) {\n int t = k;\n int kw = t % KW; t /= KW;\n int kh = t % KH; t /= KH;\n int c = t;\n val = w[((oc * C + c) * KH + kh) * KW + kw];\n }\n Bs[idx] = val;\n }\n // Load A inputs for 8 warps x 16 rows x 8 K, row-major per warp ld=8\n for (int idx = tid; idx < WARPS_M * WM * WK; idx += THREADS) {\n int kk = idx & 7;\n int tmp = idx >> 3;\n int mi = tmp & 15;\n int wp = tmp >> 4;\n int m = block_m0 + wp * WM + mi;\n int k = kt + kk;\n float val = 0.0f;\n if (m < N * OH * OW && k < Ktot) {\n int ow = m % OW;\n int t2 = m / OW;\n int oh = t2 % OH;\n int n = t2 / OH;\n int t = k;\n int kw = t % KW; t /= KW;\n int kh = t % KH; t /= KH;\n int c = t;\n int ih = oh * stride_h - pad_h + kh * dil_h;\n int iw = ow * stride_w - pad_w + kw * dil_w;\n if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) {\n val = x[((n * C + c) * H + ih) * W + iw];\n }\n }\n As[idx] = val;\n }\n __syncthreads();\n\n wmma::fragment<wmma::matrix_a, WM, WN, WK, wmma::precision::tf32, wmma::row_major> a;\n wmma::fragment<wmma::matrix_b, WM, WN, WK, wmma::precision::tf32, wmma::col_major> b;\n wmma::load_matrix_sync(a, As + warp * WM * WK, WK);\n wmma::load_matrix_sync(b, Bs + 0 * WN * WK, WK); wmma::mma_sync(acc0, a, b, acc0);\n wmma::load_matrix_sync(b, Bs + 1 * WN * WK, WK); wmma::mma_sync(acc1, a, b, acc1);\n wmma::load_matrix_sync(b, Bs + 2 * WN * WK, WK); wmma::mma_sync(acc2, a, b, acc2);\n wmma::load_matrix_sync(b, Bs + 3 * WN * WK, WK); wmma::mma_sync(acc3, a, b, acc3);\n wmma::load_matrix_sync(b, Bs + 4 * WN * WK, WK); wmma::mma_sync(acc4, a, b, acc4);\n wmma::load_matrix_sync(b, Bs + 5 * WN * WK, WK); wmma::mma_sync(acc5, a, b, acc5);\n wmma::load_matrix_sync(b, Bs + 6 * WN * WK, WK); wmma::mma_sync(acc6, a, b, acc6);\n wmma::load_matrix_sync(b, Bs + 7 * WN * WK, WK); wmma::mma_sync(acc7, a, b, acc7);\n __syncthreads();\n }\n\n // Store fragments to shared then scatter to NCHW output\n float* Cs = As; // reuse, need WARPS_M * 8 * 16*16 floats = 64KB\n auto store_one = [&](auto &acc, int ot) {\n float* ptr = Cs + (warp * OC_TILES + ot) * WM * WN;\n wmma::store_matrix_sync(ptr, acc, WN, wmma::mem_row_major);\n };\n store_one(acc0,0); store_one(acc1,1); store_one(acc2,2); store_one(acc3,3);\n store_one(acc4,4); store_one(acc5,5); store_one(acc6,6); store_one(acc7,7);\n __syncthreads();\n\n // Each thread writes elements\n int total_vals = WARPS_M * OC_TILES * WM * WN;\n int Mtot = N * OH * OW;\n for (int idx = tid; idx < total_vals; idx += THREADS) {\n int col = idx & 15;\n int tmp = idx >> 4;\n int row = tmp & 15;\n tmp >>= 4;\n int ot = tmp & 7;\n int wp = tmp >> 3;\n int m = block_m0 + wp * WM + row;\n int oc = ot * 16 + col;\n if (m < Mtot && oc < OC) {\n int ow = m % OW;\n int t2 = m / OW;\n int oh = t2 % OH;\n int n = t2 / OH;\n float val = Cs[((wp * OC_TILES + ot) * WM + row) * WN + col];\n if (has_bias) val += bias[oc];\n y[((n * OC + oc) * OH + oh) * OW + ow] = val;\n }\n }\n}\n\n// Generic fallback direct kernel for non-128 OC (not optimized)\n__global__ void conv2d_direct_kernel(const float* x, const float* w, const float* bias, float* y,\n int N,int C,int H,int W,int OC,int KH,int KW,int OH,int OW,\n int sh,int sw,int ph,int pw,int dh,int dw,int has_bias) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = N*OC*OH*OW;\n if (idx >= total) return;\n int ow = idx % OW; int t = idx / OW;\n int oh = t % OH; t /= OH;\n int oc = t % OC; int n = t / OC;\n float sum = has_bias ? bias[oc] : 0.0f;\n for (int c=0;c<C;c++) for (int kh=0;kh<KH;kh++) for (int kw=0;kw<KW;kw++) {\n int ih = oh*sh - ph + kh*dh;\n int iw = ow*sw - pw + kw*dw;\n if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) sum += x[((n*C+c)*H+ih)*W+iw] * w[((oc*C+c)*KH+kh)*KW+kw];\n }\n y[idx] = sum;\n}\n\ntorch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias,\n int stride_h, int stride_w, int pad_h, int pad_w,\n int dil_h, int dil_w, bool has_bias) {\n CHECK_CUDA(x); CHECK_CUDA(w); CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(w);\n int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);\n int OC = w.size(0), KH = w.size(2), KW = w.size(3);\n int OH = (H + 2*pad_h - dil_h*(KH-1) - 1) / stride_h + 1;\n int OW = (W + 2*pad_w - dil_w*(KW-1) - 1) / stride_w + 1;\n auto y = torch::empty({N, OC, OH, OW}, x.options());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (OC == 128 && x.scalar_type() == torch::kFloat32) {\n int Mtot = N * OH * OW;\n dim3 grid((Mtot + WARPS_M*WM - 1) / (WARPS_M*WM));\n size_t smem = (WARPS_M*OC_TILES*WM*WN) * sizeof(float); // enough for C store (larger than A+B)\n conv2d_tf32_128oc_kernel<<<grid, THREADS, smem, stream>>>(\n x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(),\n N,C,H,W,OC,KH,KW,OH,OW,stride_h,stride_w,pad_h,pad_w,dil_h,dil_w,has_bias ? 1 : 0);\n } else {\n int total = N*OC*OH*OW;\n conv2d_direct_kernel<<<(total+255)/256,256,0,stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr, y.data_ptr<float>(),\n N,C,H,W,OC,KH,KW,OH,OW,stride_h,stride_w,pad_h,pad_w,dil_h,dil_w,has_bias ? 1 : 0);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv2d_tf32_ext_v1', cpp_sources=cpp_source, cuda_sources=cuda_source,\n functions=['conv2d_forward_cuda'], 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: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), dilation: tuple = (1, 1), groups: int = 1, bias: bool = False):\n super(ModelNew, self).__init__()\n self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)\n self.stride = stride; self.padding = padding; self.dilation = dilation; 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_cuda(x.contiguous(), self.conv2d.weight.contiguous(), b, int(self.stride[0]), int(self.stride[1]), int(self.padding[0]), int(self.padding[1]), int(self.dilation[0]), int(self.dilation[1]), self.conv2d.bias is not None)\n"
}