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>\ntorch::Tensor fused_bn_post_forward(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor gn_w, torch::Tensor gn_b, int64_t num_groups, double bn_eps, double gn_eps);\n'''\n\ncuda_sources = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n\n__global__ void bn_stats_kernel(const float* __restrict__ x, float* __restrict__ mean, float* __restrict__ invstd,\n int N, int C, int H, int W, float eps) {\n int c = blockIdx.x;\n int tid = threadIdx.x;\n int HW = H * W;\n int M = N * HW;\n extern __shared__ float smem[];\n float* ss = smem;\n float* ss2 = smem + blockDim.x;\n float sum = 0.f, sq = 0.f;\n for (int i = tid; i < M; i += blockDim.x) {\n int n = i / HW;\n int hw = i - n * HW;\n float v = x[((long)n * C + c) * HW + hw];\n sum += v;\n sq += v * v;\n }\n ss[tid] = sum; ss2[tid] = sq;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (tid < s) { ss[tid] += ss[tid+s]; ss2[tid] += ss2[tid+s]; }\n __syncthreads();\n }\n if (tid == 0) {\n float m = ss[0] / (float)M;\n float var = ss2[0] / (float)M - m*m;\n if (var < 0.f) var = 0.f;\n mean[c] = m;\n invstd[c] = rsqrtf(var + eps);\n }\n}\n\n__device__ __forceinline__ float bnval(float z, int c, const float* __restrict__ mean, const float* __restrict__ invstd,\n const float* __restrict__ bw, const float* __restrict__ bb) {\n return (z - mean[c]) * invstd[c] * bw[c] + bb[c];\n}\n\n__global__ void fused_pool_gn_from_conv_kernel(const float* __restrict__ x,\n const float* __restrict__ mean,\n const float* __restrict__ invstd,\n const float* __restrict__ bn_w,\n const float* __restrict__ bn_b,\n const float* __restrict__ gamma,\n const float* __restrict__ beta,\n float* __restrict__ out,\n int N, int C, int H, int W,\n int Ho, int Wo,\n int G, float eps) {\n int bg = blockIdx.x;\n int n = bg / G;\n int g = bg - n * G;\n int tid = threadIdx.x;\n int group_c = C / G;\n int elems = group_c * Ho * Wo;\n extern __shared__ float smem[];\n float* ssum = smem;\n float* ssq = smem + blockDim.x;\n float sum = 0.0f;\n float sq = 0.0f;\n int c0 = g * group_c;\n for (int idx = tid; idx < elems; idx += blockDim.x) {\n int hw = idx % (Ho * Wo);\n int lc = idx / (Ho * Wo);\n int oh = hw / Wo;\n int ow = hw - oh * Wo;\n int c = c0 + lc;\n long base = ((long)n * C + c) * H * W + (oh * 2) * W + (ow * 2);\n float m = bnval(x[base], c, mean, invstd, bn_w, bn_b);\n float v = bnval(x[base + 1], c, mean, invstd, bn_w, bn_b); if (v > m) m = v;\n v = bnval(x[base + W], c, mean, invstd, bn_w, bn_b); if (v > m) m = v;\n v = bnval(x[base + W + 1], c, mean, invstd, bn_w, bn_b); if (v > m) m = v;\n float t = tanhf(m);\n sum += t;\n sq += t * t;\n }\n ssum[tid] = sum; ssq[tid] = sq;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (tid < s) { ssum[tid] += ssum[tid+s]; ssq[tid] += ssq[tid+s]; }\n __syncthreads();\n }\n float gm = ssum[0] / (float)elems;\n float var = ssq[0] / (float)elems - gm * gm;\n if (var < 0.f) var = 0.f;\n float ginv = rsqrtf(var + eps);\n for (int idx = tid; idx < elems; idx += blockDim.x) {\n int hw = idx % (Ho * Wo);\n int lc = idx / (Ho * Wo);\n int oh = hw / Wo;\n int ow = hw - oh * Wo;\n int c = c0 + lc;\n long base = ((long)n * C + c) * H * W + (oh * 2) * W + (ow * 2);\n float m = bnval(x[base], c, mean, invstd, bn_w, bn_b);\n float v = bnval(x[base + 1], c, mean, invstd, bn_w, bn_b); if (v > m) m = v;\n v = bnval(x[base + W], c, mean, invstd, bn_w, bn_b); if (v > m) m = v;\n v = bnval(x[base + W + 1], c, mean, invstd, bn_w, bn_b); if (v > m) m = v;\n float t = tanhf(m);\n long obase = ((long)n * C + c) * Ho * Wo + oh * Wo + ow;\n out[obase] = (t - gm) * ginv * gamma[c] + beta[c];\n }\n}\n\ntorch::Tensor fused_bn_post_forward(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor gn_w, torch::Tensor gn_b, int64_t num_groups, double bn_eps, double gn_eps) {\n int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);\n int Ho = H/2, Wo = W/2;\n auto mean = torch::empty({C}, x.options());\n auto invstd = torch::empty({C}, x.options());\n int threads = 256;\n bn_stats_kernel<<<C, threads, threads*2*sizeof(float), at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), N, C, H, W, (float)bn_eps);\n auto out = torch::empty({N, C, Ho, Wo}, x.options());\n fused_pool_gn_from_conv_kernel<<<N*(int)num_groups, threads, threads*2*sizeof(float), at::cuda::getCurrentCUDAStream()>>>(\n x.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), gn_w.data_ptr<float>(), gn_b.data_ptr<float>(), out.data_ptr<float>(),\n N, C, H, W, Ho, Wo, (int)num_groups, (float)gn_eps);\n return out;\n}\n'''\n\n_ext = load_inline(name='fused_bn_post_ext_v1', cpp_sources=cpp_sources, cuda_sources=cuda_sources, functions=['fused_bn_post_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups, num_groups):\n super(ModelNew,self).__init__()\n self.conv_transpose=nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding)\n self.batch_norm=nn.BatchNorm2d(out_channels)\n self.tanh=nn.Tanh(); self.max_pool=nn.MaxPool2d(2,2)\n self.group_norm=nn.GroupNorm(num_groups=num_groups,num_channels=out_channels)\n self.num_groups=num_groups\n def forward(self,x):\n x=self.conv_transpose(x)\n return _ext.fused_bn_post_forward(x.contiguous(), self.batch_norm.weight, self.batch_norm.bias, self.group_norm.weight, self.group_norm.bias, self.num_groups, self.batch_norm.eps, self.group_norm.eps)\n"
}