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#include <vector>\n\nvoid scale_weight_bias_cuda(torch::Tensor weight, torch::Tensor bias, torch::Tensor scale, torch::Tensor sw, torch::Tensor sb);\nvoid bn_train_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, double eps, double momentum);\nvoid bn_eval_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, double eps);\n\nstd::vector<torch::Tensor> scale_params(torch::Tensor weight, torch::Tensor bias, torch::Tensor scale) {\n at::NoGradGuard no_grad;\n auto sw = torch::empty_like(weight);\n auto sb = torch::empty_like(bias);\n scale_weight_bias_cuda(weight, bias, scale, sw, sb);\n return {sw, sb};\n}\n\ntorch::Tensor forward_cached(torch::Tensor x, torch::Tensor sw, torch::Tensor sb,\n torch::Tensor bn_weight, torch::Tensor bn_bias,\n torch::Tensor running_mean, torch::Tensor running_var,\n bool training, double eps, double momentum) {\n at::NoGradGuard no_grad;\n auto y = torch::addmm(sb, x, sw.t());\n if (training) bn_train_cuda(y, bn_weight, bn_bias, running_mean, running_var, eps, momentum);\n else bn_eval_cuda(y, bn_weight, bn_bias, running_mean, running_var, eps);\n return y;\n}\n'''\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n\ntemplate <typename scalar_t>\n__global__ void scale_weight_bias_kernel(const scalar_t* __restrict__ w, const scalar_t* __restrict__ b,\n const scalar_t* __restrict__ s, scalar_t* __restrict__ sw,\n scalar_t* __restrict__ sb, int C, int K, long total) {\n long idx = blockIdx.x * blockDim.x + threadIdx.x;\n long stride = (long)blockDim.x * gridDim.x;\n for (long i = idx; i < total; i += stride) {\n int o = (int)(i / K);\n sw[i] = w[i] * s[o];\n }\n for (int o = idx; o < C; o += stride) sb[o] = b[o] * s[o];\n}\n\nvoid scale_weight_bias_cuda(torch::Tensor weight, torch::Tensor bias, torch::Tensor scale, torch::Tensor sw, torch::Tensor sb) {\n int C = weight.size(0), K = weight.size(1); long total = (long)C*K;\n int threads=256, blocks=(int)((total+threads-1)/threads); if(blocks>4096) blocks=4096;\n AT_DISPATCH_FLOATING_TYPES(weight.scalar_type(), \"scale_weight_bias_cuda\", [&] {\n scale_weight_bias_kernel<scalar_t><<<blocks,threads>>>(weight.data_ptr<scalar_t>(), bias.data_ptr<scalar_t>(), scale.data_ptr<scalar_t>(), sw.data_ptr<scalar_t>(), sb.data_ptr<scalar_t>(), C, K, total);\n });\n}\n\n#define COLS 32\n#define ROW_GROUPS 8\n#define ROWS_PER_CHUNK 256\n#define NORM_ROWS 256\n\ntemplate <typename scalar_t>\n__global__ void partial_sum_kernel(const scalar_t* __restrict__ y, float* __restrict__ ps, float* __restrict__ pss, int N, int C, int chunks) {\n int cb = blockIdx.x * COLS;\n int ch = blockIdx.y;\n int tid = threadIdx.x;\n int col = tid & 31;\n int grp = tid >> 5;\n int c = cb + col;\n int r0 = ch * ROWS_PER_CHUNK;\n int r1 = min(N, r0 + ROWS_PER_CHUNK);\n float sum = 0.0f, ss = 0.0f;\n if (c < C) {\n for (int r = r0 + grp; r < r1; r += ROW_GROUPS) {\n float v = (float)y[(long)r * C + c];\n sum += v; ss += v * v;\n }\n }\n __shared__ float shs[256];\n __shared__ float shq[256];\n shs[tid] = sum; shq[tid] = ss;\n __syncthreads();\n if (tid < COLS && c < C) {\n float s = 0.0f, q = 0.0f;\n #pragma unroll\n for (int g=0; g<ROW_GROUPS; ++g) { s += shs[g*COLS + tid]; q += shq[g*COLS + tid]; }\n ps[(long)ch * C + c] = s;\n pss[(long)ch * C + c] = q;\n }\n}\n\ntemplate <typename scalar_t>\n__global__ void normalize_train_kernel(scalar_t* __restrict__ y, const float* __restrict__ ps, const float* __restrict__ pss,\n const scalar_t* __restrict__ gamma, const scalar_t* __restrict__ beta,\n scalar_t* __restrict__ running_mean, scalar_t* __restrict__ running_var,\n int N, int C, int chunks, float eps, float momentum) {\n int cb = blockIdx.x * COLS;\n int row_tile = blockIdx.y;\n int tid = threadIdx.x;\n int col = tid & 31;\n int grp = tid >> 5;\n int c = cb + col;\n __shared__ float mean_sh[COLS];\n __shared__ float coef_sh[COLS];\n __shared__ float beta_sh[COLS];\n if (tid < COLS && c < C) {\n float s = 0.0f, q = 0.0f;\n for (int ch=0; ch<chunks; ++ch) { s += ps[(long)ch*C + c]; q += pss[(long)ch*C + c]; }\n float mean = s / (float)N;\n float var = q / (float)N - mean * mean;\n if (var < 0.0f) var = 0.0f;\n mean_sh[col] = mean;\n coef_sh[col] = ((float)gamma[c]) * rsqrtf(var + eps);\n beta_sh[col] = (float)beta[c];\n if (row_tile == 0) {\n running_mean[c] = (scalar_t)(((1.0f - momentum) * (float)running_mean[c]) + momentum * mean);\n float unb = var * ((float)N / (float)(N - 1));\n running_var[c] = (scalar_t)(((1.0f - momentum) * (float)running_var[c]) + momentum * unb);\n }\n }\n __syncthreads();\n if (c < C) {\n int r0 = row_tile * NORM_ROWS;\n int r1 = min(N, r0 + NORM_ROWS);\n float mean = mean_sh[col], coef = coef_sh[col], be = beta_sh[col];\n for (int r = r0 + grp; r < r1; r += ROW_GROUPS) {\n long idx = (long)r * C + c;\n float v = (float)y[idx];\n y[idx] = (scalar_t)((v - mean) * coef + be);\n }\n }\n}\n\ntemplate <typename scalar_t>\n__global__ void bn_eval_kernel(scalar_t* __restrict__ y, const scalar_t* __restrict__ gamma, const scalar_t* __restrict__ beta,\n const scalar_t* __restrict__ running_mean, const scalar_t* __restrict__ running_var,\n int N, int C, float eps) {\n long total = (long)N*C;\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n long stride = (long)gridDim.x * blockDim.x;\n for (long i=idx; i<total; i+=stride) {\n int c = (int)(i % C);\n float v = (float)y[i];\n float out = (v - (float)running_mean[c]) * rsqrtf((float)running_var[c] + eps) * (float)gamma[c] + (float)beta[c];\n y[i] = (scalar_t)out;\n }\n}\n\nvoid bn_train_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, double eps, double momentum) {\n int N = y.size(0), C = y.size(1);\n int chunks = (N + ROWS_PER_CHUNK - 1) / ROWS_PER_CHUNK;\n auto opts = y.options().dtype(torch::kFloat32);\n auto ps = torch::empty({chunks, C}, opts);\n auto pss = torch::empty({chunks, C}, opts);\n dim3 grid1((C + COLS - 1)/COLS, chunks);\n dim3 block(256);\n dim3 grid2((C + COLS - 1)/COLS, (N + NORM_ROWS - 1)/NORM_ROWS);\n AT_DISPATCH_FLOATING_TYPES(y.scalar_type(), \"bn_train_cuda\", [&] {\n partial_sum_kernel<scalar_t><<<grid1, block>>>(y.data_ptr<scalar_t>(), ps.data_ptr<float>(), pss.data_ptr<float>(), N, C, chunks);\n normalize_train_kernel<scalar_t><<<grid2, block>>>(y.data_ptr<scalar_t>(), ps.data_ptr<float>(), pss.data_ptr<float>(), bn_weight.data_ptr<scalar_t>(), bn_bias.data_ptr<scalar_t>(), running_mean.data_ptr<scalar_t>(), running_var.data_ptr<scalar_t>(), N, C, chunks, (float)eps, (float)momentum);\n });\n}\n\nvoid bn_eval_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, double eps) {\n int N=y.size(0), C=y.size(1); long total=(long)N*C; int threads=256; int blocks=(int)((total+threads-1)/threads); if(blocks>4096) blocks=4096;\n AT_DISPATCH_FLOATING_TYPES(y.scalar_type(), \"bn_eval_cuda\", [&] {\n bn_eval_kernel<scalar_t><<<blocks,threads>>>(y.data_ptr<scalar_t>(), bn_weight.data_ptr<scalar_t>(), bn_bias.data_ptr<scalar_t>(), running_mean.data_ptr<scalar_t>(), running_var.data_ptr<scalar_t>(), N, C, (float)eps);\n });\n}\n'''\n\n_ext = load_inline(name='fused_linear_scale_bn_ext_v3', cpp_sources=cpp_source, cuda_sources=cuda_source,\n functions=['scale_params','forward_cached'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):\n super(ModelNew, self).__init__()\n self.gemm = nn.Linear(in_features, out_features)\n self.scale = nn.Parameter(torch.randn(scale_shape))\n self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum)\n self._sw = None\n self._sb = None\n self._cache_versions = None\n def _ensure_cache(self):\n vers = (self.gemm.weight._version, self.gemm.bias._version, self.scale._version,\n self.gemm.weight.data_ptr(), self.gemm.bias.data_ptr(), self.scale.data_ptr())\n if self._sw is None or self._cache_versions != vers:\n sw, sb = _ext.scale_params(self.gemm.weight, self.gemm.bias, self.scale)\n self._sw = sw\n self._sb = sb\n self._cache_versions = vers\n def forward(self, x):\n self._ensure_cache()\n return _ext.forward_cached(x, self._sw, self._sb,\n self.bn.weight, self.bn.bias, self.bn.running_mean, self.bn.running_var,\n self.training, float(self.bn.eps), float(self.bn.momentum))\n"
}