callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r\"\"\"\n#include <torch/extension.h>\n\ntorch::Tensor fused_gn_pool_clamp(torch::Tensor y,\n torch::Tensor gn_weight,\n torch::Tensor gn_bias,\n torch::Tensor scale,\n int64_t num_groups,\n double clamp_min,\n double clamp_max);\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\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n\n__inline__ __device__ float warp_reduce_sum(float v) {\n for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);\n return v;\n}\n\n__inline__ __device__ float block_reduce_sum(float v) {\n static __shared__ float shared[32];\n int lane = threadIdx.x & 31;\n int wid = threadIdx.x >> 5;\n v = warp_reduce_sum(v);\n if (lane == 0) shared[wid] = v;\n __syncthreads();\n v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;\n if (wid == 0) v = warp_reduce_sum(v);\n return v;\n}\n\n__global__ void gn_stats_kernel(const float* __restrict__ y, float* __restrict__ mean, float* __restrict__ invstd,\n int N, int C, int H, int W, int G) {\n int ng = blockIdx.x;\n int n = ng / G;\n int g = ng - n * G;\n int cpg = C / G;\n int HW = H * W;\n int elems = cpg * HW;\n int c0 = g * cpg;\n const float* base = y + ((n * C + c0) * HW);\n float sum = 0.0f;\n float sumsq = 0.0f;\n for (int i = threadIdx.x; i < elems; i += blockDim.x) {\n float v = base[(i / HW) * HW + (i - (i / HW) * HW)];\n sum += v;\n sumsq = fmaf(v, v, sumsq);\n }\n sum = block_reduce_sum(sum);\n sumsq = block_reduce_sum(sumsq);\n if (threadIdx.x == 0) {\n float m = sum / (float)elems;\n float var = sumsq / (float)elems - m * m;\n var = fmaxf(var, 0.0f);\n mean[ng] = m;\n invstd[ng] = rsqrtf(var + 1.0e-5f);\n }\n}\n\n__global__ void gn_pool_clamp_kernel(const float* __restrict__ y,\n const float* __restrict__ mean,\n const float* __restrict__ invstd,\n const float* __restrict__ gamma,\n const float* __restrict__ beta,\n const float* __restrict__ scale,\n float* __restrict__ out,\n int total, int C, int H, int W, int G, int OH, int OW,\n float clamp_min, float clamp_max) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int HW = H * W;\n int cpg = C / G;\n for (; idx < total; idx += blockDim.x * gridDim.x) {\n int t = idx;\n int ow = t % OW; t /= OW;\n int oh = t % OH; t /= OH;\n int c = t % C;\n int n = t / C;\n int g = c / cpg;\n int ng = n * G + g;\n float m = mean[ng];\n float is = invstd[ng];\n float a = gamma[c] * scale[c];\n float b = beta[c];\n int h0 = oh * 4;\n int w0 = ow * 4;\n const float* ptr = y + ((n * C + c) * HW + h0 * W + w0);\n float mx = -CUDART_INF_F;\n #pragma unroll\n for (int dh = 0; dh < 4; ++dh) {\n #pragma unroll\n for (int dw = 0; dw < 4; ++dw) {\n float v = (ptr[dh * W + dw] - m) * is;\n v = v * a + b;\n mx = fmaxf(mx, v);\n }\n }\n mx = fminf(fmaxf(mx, clamp_min), clamp_max);\n out[idx] = mx;\n }\n}\n\ntorch::Tensor fused_gn_pool_clamp(torch::Tensor y,\n torch::Tensor gn_weight,\n torch::Tensor gn_bias,\n torch::Tensor scale,\n int64_t num_groups,\n double clamp_min,\n double clamp_max) {\n CHECK_CUDA(y); CHECK_CUDA(gn_weight); CHECK_CUDA(gn_bias); CHECK_CUDA(scale);\n CHECK_CONTIGUOUS(y); CHECK_CONTIGUOUS(gn_weight); CHECK_CONTIGUOUS(gn_bias); CHECK_CONTIGUOUS(scale);\n int N = (int)y.size(0);\n int C = (int)y.size(1);\n int H = (int)y.size(2);\n int W = (int)y.size(3);\n int G = (int)num_groups;\n int OH = H / 4;\n int OW = W / 4;\n auto opts = y.options();\n auto out = torch::empty({N, C, OH, OW}, opts);\n auto mean = torch::empty({N, G}, opts);\n auto invstd = torch::empty({N, G}, opts);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n gn_stats_kernel<<<N * G, 256, 0, stream>>>(y.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), N, C, H, W, G);\n int total = N * C * OH * OW;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n if (blocks > 4096) blocks = 4096;\n gn_pool_clamp_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),\n gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(), scale.data_ptr<float>(), out.data_ptr<float>(),\n total, C, H, W, G, OH, OW, (float)clamp_min, (float)clamp_max);\n return out;\n}\n\"\"\"\n\n_ext = load_inline(name=\"fused_gn_pool_clamp_ext_v1\",\n cpp_sources=_cpp_src,\n cuda_sources=_cuda_src,\n functions=[\"fused_gn_pool_clamp\"],\n extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"],\n extra_cflags=[\"-O3\"],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, num_groups, scale_shape, maxpool_kernel_size, clamp_min, clamp_max):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.group_norm = nn.GroupNorm(num_groups, out_channels)\n self.scale = nn.Parameter(torch.ones(scale_shape))\n self.num_groups = num_groups\n self.maxpool_kernel_size = maxpool_kernel_size\n self.clamp_min = float(clamp_min)\n self.clamp_max = float(clamp_max)\n def forward(self, x):\n y = F.conv2d(x, self.conv.weight, self.conv.bias)\n return _ext.fused_gn_pool_clamp(y, self.group_norm.weight, self.group_norm.bias, self.scale, self.num_groups, self.clamp_min, self.clamp_max)\n"
}