callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_sources = r'''\n#include <torch/extension.h>\n\nvoid bn_eval_launcher(torch::Tensor z, torch::Tensor out, torch::Tensor scale, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, double eps);\nvoid bn_stats_launcher(torch::Tensor z, torch::Tensor scale, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor mean, torch::Tensor invstd, double momentum, double eps);\nvoid bn_train_norm_launcher(torch::Tensor z, torch::Tensor out, torch::Tensor scale, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor mean, torch::Tensor invstd);\n\ntorch::Tensor forward_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,\n torch::Tensor scale, torch::Tensor bn_w, torch::Tensor bn_b,\n torch::Tensor running_mean, torch::Tensor running_var,\n bool training, double momentum, double eps) {\n auto z = at::linear(input, weight, bias);\n auto out = torch::empty_like(z);\n if (training) {\n auto mean = torch::empty({z.size(1)}, z.options());\n auto invstd = torch::empty({z.size(1)}, z.options());\n bn_stats_launcher(z, scale, running_mean, running_var, mean, invstd, momentum, eps);\n bn_train_norm_launcher(z, out, scale, bn_w, bn_b, mean, invstd);\n } else {\n bn_eval_launcher(z, out, scale, bn_w, bn_b, running_mean, running_var, eps);\n }\n return out;\n}\n'''\n\ncuda_sources = r'''\n#include <torch/types.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void bn_eval_kernel(const float* __restrict__ z, float* __restrict__ out,\n const float* __restrict__ scale, const float* __restrict__ gamma,\n const float* __restrict__ beta, const float* __restrict__ rmean,\n const float* __restrict__ rvar, int total, int N, float eps) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (; idx < total; idx += stride) {\n int j = idx % N;\n float v = z[idx] * scale[j];\n float y = (v - rmean[j]) * rsqrtf(rvar[j] + eps);\n out[idx] = y * gamma[j] + beta[j];\n }\n}\n\n__global__ void bn_train_norm_kernel(const float* __restrict__ z, float* __restrict__ out,\n const float* __restrict__ scale, const float* __restrict__ gamma,\n const float* __restrict__ beta, const float* __restrict__ mean,\n const float* __restrict__ invstd, int total, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride = blockDim.x * gridDim.x;\n for (; idx < total; idx += stride) {\n int j = idx % N;\n float v = z[idx] * scale[j];\n out[idx] = (v - mean[j]) * invstd[j] * gamma[j] + beta[j];\n }\n}\n\n// block: (32,8), one block computes statistics for 32 contiguous columns.\n__global__ void bn_stats_kernel(const float* __restrict__ z, const float* __restrict__ scale,\n float* __restrict__ running_mean, float* __restrict__ running_var,\n float* __restrict__ mean, float* __restrict__ invstd,\n int M, int N, float momentum, float eps) {\n int tx = threadIdx.x; // feature within tile, 0..31\n int ty = threadIdx.y; // row-stripe, 0..7\n int j = blockIdx.x * 32 + tx;\n __shared__ float sh_sum[8][32];\n __shared__ float sh_sumsq[8][32];\n float s = 0.0f;\n float ss = 0.0f;\n if (j < N) {\n for (int r = ty; r < M; r += 8) {\n float v = z[r * N + j];\n s += v;\n ss += v * v;\n }\n }\n sh_sum[ty][tx] = s;\n sh_sumsq[ty][tx] = ss;\n __syncthreads();\n if (ty == 0 && j < N) {\n float sum = 0.0f, sumsq = 0.0f;\n #pragma unroll\n for (int k = 0; k < 8; ++k) {\n sum += sh_sum[k][tx];\n sumsq += sh_sumsq[k][tx];\n }\n float invM = 1.0f / (float)M;\n float mz = sum * invM;\n float varz = sumsq * invM - mz * mz;\n varz = fmaxf(varz, 0.0f);\n float sc = scale[j];\n float my = mz * sc;\n float vary = varz * sc * sc;\n mean[j] = my;\n invstd[j] = rsqrtf(vary + eps);\n float unbiased = (M > 1) ? (vary * ((float)M / (float)(M - 1))) : 0.0f;\n running_mean[j] = running_mean[j] * (1.0f - momentum) + momentum * my;\n running_var[j] = running_var[j] * (1.0f - momentum) + momentum * unbiased;\n }\n}\n\nvoid bn_eval_launcher(torch::Tensor z, torch::Tensor out, torch::Tensor scale, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, double eps) {\n int M = (int)z.size(0), N = (int)z.size(1);\n int total = M * N;\n int threads = 256;\n int blocks = min((total + threads - 1) / threads, 32768);\n bn_eval_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(\n z.data_ptr<float>(), out.data_ptr<float>(), scale.data_ptr<float>(), bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), total, N, (float)eps);\n}\n\nvoid bn_stats_launcher(torch::Tensor z, torch::Tensor scale, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor mean, torch::Tensor invstd, double momentum, double eps) {\n int M = (int)z.size(0), N = (int)z.size(1);\n dim3 threads(32, 8);\n dim3 blocks((N + 31) / 32);\n bn_stats_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(\n z.data_ptr<float>(), scale.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), M, N, (float)momentum, (float)eps);\n}\n\nvoid bn_train_norm_launcher(torch::Tensor z, torch::Tensor out, torch::Tensor scale, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor mean, torch::Tensor invstd) {\n int M = (int)z.size(0), N = (int)z.size(1);\n int total = M * N;\n int threads = 256;\n int blocks = min((total + threads - 1) / threads, 32768);\n bn_train_norm_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(\n z.data_ptr<float>(), out.data_ptr<float>(), scale.data_ptr<float>(), bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), total, N);\n}\n'''\n\n_ext = load_inline(name='gemm_scale_bn_ext_v1', cpp_sources=cpp_sources, cuda_sources=cuda_sources,\n functions=['forward_cuda'], 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\n def forward(self, x):\n return _ext.forward_cuda(x, self.gemm.weight, self.gemm.bias, self.scale,\n self.bn.weight, self.bn.bias, self.bn.running_mean, self.bn.running_var,\n bool(self.bn.training), float(self.bn.momentum), float(self.bn.eps))\n"
}