call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n\ntorch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\ntorch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {\n return layernorm_cuda(x, weight, bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <stdint.h>\n\n#ifndef CHECK_CUDA\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#endif\n\nstatic constexpr int THREADS = 256;\nstatic constexpr int BLOCKS_PER_ROW = 256;\n\n__inline__ __device__ double warp_reduce_sum_double(double v) {\n unsigned mask = 0xffffffffu;\n v += __shfl_down_sync(mask, v, 16);\n v += __shfl_down_sync(mask, v, 8);\n v += __shfl_down_sync(mask, v, 4);\n v += __shfl_down_sync(mask, v, 2);\n v += __shfl_down_sync(mask, v, 1);\n return v;\n}\n\n__inline__ __device__ double block_reduce_sum_double(double v) {\n __shared__ double shared[32];\n int lane = threadIdx.x & 31;\n int wid = threadIdx.x >> 5;\n v = warp_reduce_sum_double(v);\n if (lane == 0) shared[wid] = v;\n __syncthreads();\n v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0;\n if (wid == 0) v = warp_reduce_sum_double(v);\n return v;\n}\n\n__global__ void reduce_stage1_kernel(const float* __restrict__ x,\n double* __restrict__ psum,\n double* __restrict__ psq,\n int64_t N4) {\n int row = blockIdx.x / BLOCKS_PER_ROW;\n int part = blockIdx.x - row * BLOCKS_PER_ROW;\n const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x + (int64_t)row * (N4 * 4));\n double s = 0.0;\n double q = 0.0;\n for (int64_t i = (int64_t)part * blockDim.x + threadIdx.x; i < N4; i += (int64_t)BLOCKS_PER_ROW * blockDim.x) {\n float4 v = x4[i];\n s += (double)v.x + (double)v.y + (double)v.z + (double)v.w;\n q += (double)v.x * (double)v.x + (double)v.y * (double)v.y + (double)v.z * (double)v.z + (double)v.w * (double)v.w;\n }\n double bs = block_reduce_sum_double(s);\n double bq = block_reduce_sum_double(q);\n if (threadIdx.x == 0) {\n int idx = row * BLOCKS_PER_ROW + part;\n psum[idx] = bs;\n psq[idx] = bq;\n }\n}\n\n__global__ void reduce_stage2_kernel(const double* __restrict__ psum,\n const double* __restrict__ psq,\n float* __restrict__ mean,\n float* __restrict__ rstd,\n int64_t N) {\n int row = blockIdx.x;\n double s = 0.0;\n double q = 0.0;\n for (int i = threadIdx.x; i < BLOCKS_PER_ROW; i += blockDim.x) {\n s += psum[row * BLOCKS_PER_ROW + i];\n q += psq[row * BLOCKS_PER_ROW + i];\n }\n s = block_reduce_sum_double(s);\n q = block_reduce_sum_double(q);\n if (threadIdx.x == 0) {\n double m = s / (double)N;\n double var = q / (double)N - m * m;\n if (var < 0.0) var = 0.0;\n mean[row] = (float)m;\n rstd[row] = (float)rsqrt(var + 1.0e-5);\n }\n}\n\n__global__ void norm_kernel(const float* __restrict__ x,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n const float* __restrict__ mean,\n const float* __restrict__ rstd,\n float* __restrict__ y,\n int64_t total4,\n int64_t N4) {\n int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;\n int64_t stride = (int64_t)blockDim.x * gridDim.x;\n const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);\n const float4* __restrict__ w4 = reinterpret_cast<const float4*>(weight);\n const float4* __restrict__ b4 = reinterpret_cast<const float4*>(bias);\n float4* __restrict__ y4 = reinterpret_cast<float4*>(y);\n for (int64_t idx = tid; idx < total4; idx += stride) {\n int row = (int)(idx / N4);\n int64_t j = idx - (int64_t)row * N4;\n float m = mean[row];\n float r = rstd[row];\n float4 xv = x4[idx];\n float4 wv = w4[j];\n float4 bv = b4[j];\n float4 out;\n out.x = (xv.x - m) * r * wv.x + bv.x;\n out.y = (xv.y - m) * r * wv.y + bv.y;\n out.z = (xv.z - m) * r * wv.z + bv.z;\n out.w = (xv.w - m) * r * wv.w + bv.w;\n y4[idx] = out;\n }\n}\n\ntorch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {\n CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias);\n TORCH_CHECK(x.scalar_type() == torch::kFloat32, \"x must be float32\");\n TORCH_CHECK(weight.scalar_type() == torch::kFloat32 && bias.scalar_type() == torch::kFloat32, \"params must be float32\");\n TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), \"tensors must be contiguous\");\n int64_t N = weight.numel();\n int64_t total = x.numel();\n TORCH_CHECK(total % N == 0, \"bad shape\");\n TORCH_CHECK((N % 4) == 0 && (total % 4) == 0, \"N must be divisible by 4\");\n int64_t B = total / N;\n auto y = torch::empty_like(x);\n auto optsD = x.options().dtype(torch::kFloat64);\n auto optsF = x.options().dtype(torch::kFloat32);\n auto psum = torch::empty({B, BLOCKS_PER_ROW}, optsD);\n auto psq = torch::empty({B, BLOCKS_PER_ROW}, optsD);\n auto mean = torch::empty({B}, optsF);\n auto rstd = torch::empty({B}, optsF);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n dim3 grid1((unsigned int)(B * BLOCKS_PER_ROW));\n reduce_stage1_kernel<<<grid1, THREADS, 0, stream>>>(x.data_ptr<float>(), psum.data_ptr<double>(), psq.data_ptr<double>(), N / 4);\n reduce_stage2_kernel<<<(unsigned int)B, THREADS, 0, stream>>>(psum.data_ptr<double>(), psq.data_ptr<double>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), N);\n int threads = 256;\n int64_t total4 = total / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n if (blocks > 65535) blocks = 65535;\n norm_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), y.data_ptr<float>(), total4, N / 4);\n return y;\n}\n'''\n\n_ext = load_inline(name='ln_ext_v1', cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=['layernorm_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])\n\nclass ModelNew(nn.Module):\n def __init__(self, normalized_shape: tuple):\n super(ModelNew, self).__init__()\n self.ln = nn.LayerNorm(normalized_shape=normalized_shape)\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _ext.layernorm_forward(x, self.ln.weight, self.ln.bias)\n"
}