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>\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta);\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <c10/cuda/CUDAException.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\nstatic constexpr int IC=64,OC=64,IH=256,IW=256,OH=258,OW=258,KH=3,KW=3,NUM_GROUPS=8,CH_PER_GROUP=8,GROUP_ELEMS=8*258*258; static constexpr float INV_SQRT2=0.7071067811865475244f;\n__device__ __forceinline__ float gelu_exact(float x){return 0.5f*x*(1.0f+erff(x*INV_SQRT2));}\n__device__ __forceinline__ float madd_tf32(float a,float b,float c){return fmaf(__float_to_tf32(a),__float_to_tf32(b),c);} \n__global__ void conv_gelu_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ a,long total){long idx=(long)blockIdx.x*blockDim.x+threadIdx.x; long stride=(long)blockDim.x*gridDim.x; for(;idx<total;idx+=stride){int ow=idx%OW; long t=idx/OW; int oh=t%OH; t/=OH; int oc=t%OC; int n=t/OC; float acc=b?__ldg(b+oc):0.f; if(oh>=2&&oh<=255&&ow>=2&&ow<=255){#pragma unroll 4\nfor(int ic=0;ic<IC;++ic){const float* xp=x+(((long)n*IC+ic)*IH*IW); const float* wp=w+(((long)ic*OC+oc)*KH*KW); acc=madd_tf32(__ldg(xp+(oh)*IW+ow),__ldg(wp+0),acc); acc=madd_tf32(__ldg(xp+(oh)*IW+ow-1),__ldg(wp+1),acc); acc=madd_tf32(__ldg(xp+(oh)*IW+ow-2),__ldg(wp+2),acc); acc=madd_tf32(__ldg(xp+(oh-1)*IW+ow),__ldg(wp+3),acc); acc=madd_tf32(__ldg(xp+(oh-1)*IW+ow-1),__ldg(wp+4),acc); acc=madd_tf32(__ldg(xp+(oh-1)*IW+ow-2),__ldg(wp+5),acc); acc=madd_tf32(__ldg(xp+(oh-2)*IW+ow),__ldg(wp+6),acc); acc=madd_tf32(__ldg(xp+(oh-2)*IW+ow-1),__ldg(wp+7),acc); acc=madd_tf32(__ldg(xp+(oh-2)*IW+ow-2),__ldg(wp+8),acc);}} else {#pragma unroll\nfor(int ic=0;ic<IC;++ic){const float* xp=x+(((long)n*IC+ic)*IH*IW); const float* wp=w+(((long)ic*OC+oc)*KH*KW); #pragma unroll\nfor(int kh=0;kh<KH;++kh){int ih=oh-kh; if((unsigned)ih<IH){#pragma unroll\nfor(int kw=0;kw<KW;++kw){int iw=ow-kw; if((unsigned)iw<IW) acc=madd_tf32(__ldg(xp+ih*IW+iw),__ldg(wp+kh*KW+kw),acc);}}}}} a[idx]=gelu_exact(acc);}}\n__global__ void group_stats_kernel(const float* __restrict__ a,float* __restrict__ mean,float* __restrict__ invstd){int bg=blockIdx.x,n=bg/NUM_GROUPS,g=bg-n*NUM_GROUPS,c0=g*CH_PER_GROUP; float sum=0.f,sumsq=0.f; for(int i=threadIdx.x;i<GROUP_ELEMS;i+=blockDim.x){int spatial=i%(OH*OW); int c=c0+i/(OH*OW); float v=__ldg(a+(((long)n*OC+c)*OH*OW+spatial)); sum+=v; sumsq=fmaf(v,v,sumsq);} __shared__ float ssum[256],ssq[256]; ssum[threadIdx.x]=sum; ssq[threadIdx.x]=sumsq; __syncthreads(); for(int off=128;off>0;off>>=1){if(threadIdx.x<off){ssum[threadIdx.x]+=ssum[threadIdx.x+off]; ssq[threadIdx.x]+=ssq[threadIdx.x+off];} __syncthreads();} if(threadIdx.x==0){float m=ssum[0]/(float)GROUP_ELEMS; float var=fmaxf(ssq[0]/(float)GROUP_ELEMS-m*m,0.f); mean[bg]=m; invstd[bg]=rsqrtf(var+1.0e-5f);}}\n__global__ void norm_kernel(const float* __restrict__ a,const float* __restrict__ mean,const float* __restrict__ invstd,const float* __restrict__ gamma,const float* __restrict__ beta,float* __restrict__ out,long total){long idx=(long)blockIdx.x*blockDim.x+threadIdx.x; long stride=(long)blockDim.x*gridDim.x; for(;idx<total;idx+=stride){long t=idx/(OH*OW); int c=t%OC; int n=t/OC; int bg=n*NUM_GROUPS+c/CH_PER_GROUP; float y=(__ldg(a+idx)-__ldg(mean+bg))*__ldg(invstd+bg); out[idx]=fmaf(y,__ldg(gamma+c),__ldg(beta+c));}}\ntorch::Tensor fused_forward(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,torch::Tensor gamma,torch::Tensor beta){auto out=torch::empty({x.size(0),(long)OC,(long)OH,(long)OW},x.options()); auto a=torch::empty_like(out); auto mean=torch::empty({x.size(0)*NUM_GROUPS},x.options()); auto invstd=torch::empty({x.size(0)*NUM_GROUPS},x.options()); long total=x.size(0)*(long)OC*OH*OW; cudaStream_t stream=at::cuda::getCurrentCUDAStream(); int threads=256; int blocks=(int)min((long)65535,(total+threads-1)/threads); conv_gelu_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bias.defined()?bias.data_ptr<float>():nullptr,a.data_ptr<float>(),total); C10_CUDA_KERNEL_LAUNCH_CHECK(); group_stats_kernel<<<(int)(x.size(0)*NUM_GROUPS),256,0,stream>>>(a.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>()); C10_CUDA_KERNEL_LAUNCH_CHECK(); norm_kernel<<<blocks,threads,0,stream>>>(a.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),out.data_ptr<float>(),total); C10_CUDA_KERNEL_LAUNCH_CHECK(); return out;}\n'''\n_ext=load_inline(name='fused_deconv_gelu_gn_ext_tf32',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['fused_forward'],extra_cuda_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,in_channels,out_channels,kernel_size,stride,groups,num_groups):\n super(ModelNew,self).__init__(); self.conv_transpose=nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride=stride); self.group_norm=nn.GroupNorm(num_groups=num_groups,num_channels=out_channels)\n def forward(self,x):\n return _ext.fused_forward(x.contiguous(),self.conv_transpose.weight.contiguous(),self.conv_transpose.bias.contiguous(),self.group_norm.weight.contiguous(),self.group_norm.bias.contiguous())\n"
}