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 conv_bn_forward(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b,\n torch::Tensor bn_w, torch::Tensor bn_b,\n torch::Tensor running_mean, torch::Tensor running_var,\n double eps, double scale, bool training, double momentum);\n\"\"\"\n\ncuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <ATen/cuda/CUDAContext.h>\n\n#define OC_PER 8\n#define TW 16\n#define TH 8\n\n__global__ void conv_eval_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ cb,\n const float* __restrict__ bw, const float* __restrict__ bb,\n const float* __restrict__ rm, const float* __restrict__ rv,\n float* __restrict__ out,\n int N, int Cin, int H, int W, int Cout, int K, int Hout, int Wout,\n float eps, float scale) {\n extern __shared__ float sh[];\n int tile_w = blockIdx.x % ((Wout + TW - 1) / TW);\n int tile_h = blockIdx.x / ((Wout + TW - 1) / TW);\n int ocg = blockIdx.y * OC_PER;\n int n = blockIdx.z;\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n int tid = ty * blockDim.x + tx;\n int in_tw = TW + K - 1;\n int in_th = TH + K - 1;\n int load_elems = Cin * in_th * in_tw;\n int base_oh = tile_h * TH;\n int base_ow = tile_w * TW;\n for (int idx = tid; idx < load_elems; idx += blockDim.x * blockDim.y) {\n int t = idx;\n int iwl = t % in_tw; t /= in_tw;\n int ihl = t % in_th; t /= in_th;\n int ci = t;\n int ih = base_oh + ihl;\n int iw = base_ow + iwl;\n float v = 0.0f;\n if (ih < H && iw < W) v = x[((n * Cin + ci) * H + ih) * W + iw];\n sh[idx] = v;\n }\n __syncthreads();\n int oh = base_oh + ty;\n int ow = base_ow + tx;\n if (oh < Hout && ow < Wout) {\n float acc[OC_PER];\n #pragma unroll\n for (int o=0;o<OC_PER;o++) {\n int oc = ocg + o;\n acc[o] = (oc < Cout && cb != nullptr) ? cb[oc] : 0.0f;\n }\n for (int ci=0; ci<Cin; ++ci) {\n for (int kh=0; kh<K; ++kh) {\n #pragma unroll\n for (int kw=0; kw<3; ++kw) {\n if (kw >= K) continue;\n float xv = sh[(ci * in_th + (ty + kh)) * in_tw + (tx + kw)];\n #pragma unroll\n for (int o=0;o<OC_PER;o++) {\n int oc = ocg + o;\n if (oc < Cout) acc[o] = fmaf(xv, w[((oc * Cin + ci) * K + kh) * K + kw], acc[o]);\n }\n }\n }\n }\n int pix = oh * Wout + ow;\n int plane = Hout * Wout;\n #pragma unroll\n for (int o=0;o<OC_PER;o++) {\n int oc = ocg + o;\n if (oc < Cout) {\n float inv = rsqrtf(rv[oc] + eps);\n float a = bw[oc] * inv * scale;\n float b = (bb[oc] - rm[oc] * bw[oc] * inv) * scale;\n out[((n * Cout + oc) * Hout * Wout) + pix] = acc[o] * a + b;\n }\n }\n }\n}\n\n__global__ void conv_store_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ cb,\n float* __restrict__ out,\n int N, int Cin, int H, int W, int Cout, int K, int Hout, int Wout) {\n extern __shared__ float sh[];\n int tile_w = blockIdx.x % ((Wout + TW - 1) / TW);\n int tile_h = blockIdx.x / ((Wout + TW - 1) / TW);\n int ocg = blockIdx.y * OC_PER;\n int n = blockIdx.z;\n int tx = threadIdx.x;\n int ty = threadIdx.y;\n int tid = ty * blockDim.x + tx;\n int in_tw = TW + K - 1;\n int in_th = TH + K - 1;\n int load_elems = Cin * in_th * in_tw;\n int base_oh = tile_h * TH;\n int base_ow = tile_w * TW;\n for (int idx = tid; idx < load_elems; idx += blockDim.x * blockDim.y) {\n int t = idx;\n int iwl = t % in_tw; t /= in_tw;\n int ihl = t % in_th; t /= in_th;\n int ci = t;\n int ih = base_oh + ihl;\n int iw = base_ow + iwl;\n float v = 0.0f;\n if (ih < H && iw < W) v = x[((n * Cin + ci) * H + ih) * W + iw];\n sh[idx] = v;\n }\n __syncthreads();\n int oh = base_oh + ty;\n int ow = base_ow + tx;\n if (oh < Hout && ow < Wout) {\n float acc[OC_PER];\n #pragma unroll\n for (int o=0;o<OC_PER;o++) { int oc=ocg+o; acc[o] = (oc<Cout && cb != nullptr) ? cb[oc] : 0.0f; }\n for (int ci=0; ci<Cin; ++ci) {\n for (int kh=0; kh<K; ++kh) {\n #pragma unroll\n for (int kw=0; kw<3; ++kw) { if (kw>=K) continue; float xv=sh[(ci*in_th+(ty+kh))*in_tw+(tx+kw)];\n #pragma unroll\n for (int o=0;o<OC_PER;o++) { int oc=ocg+o; if (oc<Cout) acc[o]=fmaf(xv,w[((oc*Cin+ci)*K+kh)*K+kw],acc[o]); }\n }\n }\n }\n int pix=oh*Wout+ow; int plane=Hout*Wout;\n #pragma unroll\n for (int o=0;o<OC_PER;o++) { int oc=ocg+o; if (oc<Cout) out[((n*Cout+oc)*plane)+pix]=acc[o]; }\n }\n}\n\n__global__ void reduce_stats_kernel(const float* __restrict__ y, float* __restrict__ mean, float* __restrict__ var,\n int N, int C, int HW) {\n int c = blockIdx.x;\n float sum = 0.0f, ss = 0.0f;\n int M = N * HW;\n for (int i = threadIdx.x; i < M; i += blockDim.x) {\n int n = i / HW;\n int p = i - n * HW;\n float v = y[(n * C + c) * HW + p];\n sum += v;\n ss += v * v;\n }\n __shared__ float ssum[256];\n __shared__ float sss[256];\n ssum[threadIdx.x]=sum; sss[threadIdx.x]=ss;\n __syncthreads();\n for (int stride=128; stride>0; stride >>= 1) {\n if (threadIdx.x < stride) { ssum[threadIdx.x] += ssum[threadIdx.x+stride]; sss[threadIdx.x] += sss[threadIdx.x+stride]; }\n __syncthreads();\n }\n if (threadIdx.x == 0) {\n float m = ssum[0] / (float)M;\n float vv = sss[0] / (float)M - m * m;\n if (vv < 0.0f) vv = 0.0f;\n mean[c] = m; var[c] = vv;\n }\n}\n\n__global__ void norm_update_kernel(float* __restrict__ y, const float* __restrict__ mean, const float* __restrict__ var,\n const float* __restrict__ bw, const float* __restrict__ bb,\n float* __restrict__ rm, float* __restrict__ rv,\n int total, int C, int HW, int M, float eps, float scale, float momentum) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n int p = idx % HW;\n int t = idx / HW;\n int c = t % C;\n float inv = rsqrtf(var[c] + eps);\n y[idx] = ((y[idx] - mean[c]) * inv * bw[c] + bb[c]) * scale;\n }\n // update running stats using first C threads of first block\n if (blockIdx.x == 0 && threadIdx.x < C) {\n int c = threadIdx.x;\n float unbiased = var[c] * ((float)M / (float)(M - 1));\n rm[c] = rm[c] * (1.0f - momentum) + mean[c] * momentum;\n rv[c] = rv[c] * (1.0f - momentum) + unbiased * momentum;\n }\n}\n\ntorch::Tensor conv_bn_forward(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b,\n torch::Tensor bn_w, torch::Tensor bn_b,\n torch::Tensor running_mean, torch::Tensor running_var,\n double eps_d, double scale_d, bool training, double momentum_d) {\n int N = x.size(0); int Cin = x.size(1); int H = x.size(2); int W = x.size(3);\n int Cout = w.size(0); int K = w.size(2); int Hout = H - K + 1; int Wout = W - K + 1;\n auto out = torch::empty({N, Cout, Hout, Wout}, x.options());\n dim3 block(TW, TH);\n int tiles = ((Wout + TW - 1) / TW) * ((Hout + TH - 1) / TH);\n dim3 grid(tiles, (Cout + OC_PER - 1) / OC_PER, N);\n size_t shmem = (size_t)Cin * (TH + K - 1) * (TW + K - 1) * sizeof(float);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (!training) {\n conv_eval_kernel<<<grid, block, shmem, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), conv_b.defined()?conv_b.data_ptr<float>():nullptr,\n bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), out.data_ptr<float>(),\n N,Cin,H,W,Cout,K,Hout,Wout,(float)eps_d,(float)scale_d);\n } else {\n conv_store_kernel<<<grid, block, shmem, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), conv_b.defined()?conv_b.data_ptr<float>():nullptr,\n out.data_ptr<float>(), N,Cin,H,W,Cout,K,Hout,Wout);\n auto mean = torch::empty({Cout}, x.options());\n auto var = torch::empty({Cout}, x.options());\n int HW = Hout * Wout; int M = N * HW; int total = N * Cout * HW;\n reduce_stats_kernel<<<Cout, 256, 0, stream>>>(out.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), N, Cout, HW);\n norm_update_kernel<<<(total + 255)/256, 256, 0, stream>>>(out.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(),\n bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),\n total, Cout, HW, M, (float)eps_d, (float)scale_d, (float)momentum_d);\n }\n return out;\n}\n\"\"\"\n\n_ext = load_inline(name=\"conv_bn_fused_ext_v1\", cpp_sources=cpp_src, cuda_sources=cuda_src,\n functions=[\"conv_bn_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, scaling_factor):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bn = nn.BatchNorm2d(out_channels)\n self.scaling_factor = scaling_factor\n def forward(self, x):\n return _ext.conv_bn_forward(x.contiguous(), self.conv.weight, self.conv.bias, self.bn.weight, self.bn.bias,\n self.bn.running_mean, self.bn.running_var, float(self.bn.eps), float(self.scaling_factor),\n bool(self.training), float(self.bn.momentum if self.bn.momentum is not None else 0.1))\n"
}