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#include <ATen/ATen.h>\n#include <ATen/ops/linear.h>\n\nvoid launch_swish_bias_gn(torch::Tensor y, torch::Tensor bias, torch::Tensor gn_weight, torch::Tensor gn_bias, int64_t num_groups);\n\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor lin_bias,\n torch::Tensor bias, torch::Tensor gn_weight, torch::Tensor gn_bias,\n int64_t num_groups) {\n auto y = at::linear(x, weight, lin_bias);\n launch_swish_bias_gn(y, bias, gn_weight, gn_bias, num_groups);\n return y;\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\nstatic __inline__ __device__ float warp_sum(float 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__global__ __launch_bounds__(256, 2)\nvoid swish_bias_gn64_kernel(float* __restrict__ y,\n const float* __restrict__ bias,\n const float* __restrict__ gamma,\n const float* __restrict__ beta,\n int N, int C, int G) {\n int row = blockIdx.x;\n int tid = threadIdx.x;\n int lane = tid & 31;\n int warp = tid >> 5; // 0..7\n const int warps = 8;\n // Optimized path for the benchmark: C=4096, G=64, group size=64.\n for (int g = warp; g < G; g += warps) {\n int c0 = g * 64 + lane;\n int idx0 = row * C + c0;\n float a0 = y[idx0];\n float z0 = a0 / (1.0f + expf(-a0)) + bias[c0];\n int c1 = c0 + 32;\n int idx1 = idx0 + 32;\n float a1 = y[idx1];\n float z1 = a1 / (1.0f + expf(-a1)) + bias[c1];\n\n float s = z0 + z1;\n float ss = z0 * z0 + z1 * z1;\n s = warp_sum(s);\n ss = warp_sum(ss);\n float mean = __shfl_sync(0xffffffffu, s, 0) * (1.0f / 64.0f);\n float var = __shfl_sync(0xffffffffu, ss, 0) * (1.0f / 64.0f) - mean * mean;\n float inv = rsqrtf(var + 1.0e-5f);\n y[idx0] = (z0 - mean) * inv * gamma[c0] + beta[c0];\n y[idx1] = (z1 - mean) * inv * gamma[c1] + beta[c1];\n }\n}\n\n__global__ __launch_bounds__(256, 2)\nvoid swish_bias_gn_generic_kernel(float* __restrict__ y,\n const float* __restrict__ bias,\n const float* __restrict__ gamma,\n const float* __restrict__ beta,\n int N, int C, int G, int group_size) {\n int row = blockIdx.x;\n int g = blockIdx.y;\n int tid = threadIdx.x;\n extern __shared__ float smem[];\n float sum = 0.0f, sq = 0.0f;\n int base = row * C + g * group_size;\n for (int i = tid; i < group_size; i += blockDim.x) {\n int c = g * group_size + i;\n float a = y[base + i];\n float z = a / (1.0f + expf(-a)) + bias[c];\n y[base + i] = z;\n sum += z;\n sq += z * z;\n }\n smem[tid] = sum;\n smem[tid + blockDim.x] = sq;\n __syncthreads();\n for (int stride = blockDim.x/2; stride > 0; stride >>= 1) {\n if (tid < stride) {\n smem[tid] += smem[tid + stride];\n smem[tid + blockDim.x] += smem[tid + blockDim.x + stride];\n }\n __syncthreads();\n }\n float mean = smem[0] / (float)group_size;\n float var = smem[blockDim.x] / (float)group_size - mean * mean;\n float inv = rsqrtf(var + 1.0e-5f);\n for (int i = tid; i < group_size; i += blockDim.x) {\n int c = g * group_size + i;\n float z = y[base + i];\n y[base + i] = (z - mean) * inv * gamma[c] + beta[c];\n }\n}\n\nvoid launch_swish_bias_gn(torch::Tensor y, torch::Tensor bias, torch::Tensor gn_weight, torch::Tensor gn_bias, int64_t num_groups) {\n int N = (int)y.size(0);\n int C = (int)y.size(1);\n int G = (int)num_groups;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (C == 4096 && G == 64) {\n swish_bias_gn64_kernel<<<N, 256, 0, stream>>>(\n y.data_ptr<float>(), bias.data_ptr<float>(), gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(), N, C, G);\n } else {\n int group_size = C / G;\n dim3 grid(N, G);\n int threads = 256;\n swish_bias_gn_generic_kernel<<<grid, threads, 2*threads*sizeof(float), stream>>>(\n y.data_ptr<float>(), bias.data_ptr<float>(), gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(), N, C, G, group_size);\n }\n}\n'''\n\n_ext = load_inline(\n name='fused_linear_swish_bias_gn_ext_v1',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cflags=['-O3'],\n extra_cuda_cflags=['-O3'],\n with_cuda=True,\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_features, out_features, num_groups, bias_shape):\n super().__init__()\n self.matmul = nn.Linear(in_features, out_features)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n self.group_norm = nn.GroupNorm(num_groups, out_features)\n self.num_groups = num_groups\n\n def forward(self, x):\n return _ext.fused_forward(x, self.matmul.weight, self.matmul.bias, self.bias,\n self.group_norm.weight, self.group_norm.bias, self.num_groups)\n"
}