callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <ATen/ATen.h>\n#include <ATen/ops/addmm.h>\n#include <ATen/ops/t.h>\nvoid bn_swish_train_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor extra_bias, double momentum, double eps, double divide_value);\nvoid bn_swish_eval_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor extra_bias, double eps, double divide_value);\ntorch::Tensor linear_bn_swish(torch::Tensor x, torch::Tensor weight, torch::Tensor lin_bias, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor extra_bias, bool training, double momentum, double eps, double divide_value) {\n auto y = at::addmm(lin_bias, x, weight.t());\n if (training) bn_swish_train_cuda(y,bn_weight,bn_bias,running_mean,running_var,extra_bias,momentum,eps,divide_value);\n else bn_swish_eval_cuda(y,bn_weight,bn_bias,running_mean,running_var,extra_bias,eps,divide_value);\n return y;\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#define THREADS 256\n__global__ void reduce_cols_coal_kernel(const float* __restrict__ y, float* __restrict__ running_mean, float* __restrict__ running_var, float* __restrict__ mean, float* __restrict__ invstd, int N, int M, float momentum, float eps){\n int cx=threadIdx.x; int lane=threadIdx.y; int col=blockIdx.x*32+cx;\n float sum=0.f, sumsq=0.f;\n if(col<M){\n for(int row=lane; row<N; row+=8){ float v=y[row*M+col]; sum+=v; sumsq=fmaf(v,v,sumsq); }\n }\n __shared__ float ss[8][32]; __shared__ float ssq[8][32];\n ss[lane][cx]=sum; ssq[lane][cx]=sumsq; __syncthreads();\n if(lane==0 && col<M){\n float s=0.f, q=0.f;\n #pragma unroll\n for(int i=0;i<8;i++){ s+=ss[i][cx]; q+=ssq[i][cx]; }\n float m=s/(float)N; float var=q/(float)N - m*m; var=fmaxf(var,0.f);\n mean[col]=m; invstd[col]=rsqrtf(var+eps);\n running_mean[col]=fmaf(momentum,m,running_mean[col]*(1.f-momentum));\n running_var[col]=fmaf(momentum,(var*((float)N/(float)(N-1))),running_var[col]*(1.f-momentum));\n }\n}\n__global__ void apply_train_kernel(float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,const float bias0,int total){\n int idx=blockIdx.x*blockDim.x+threadIdx.x;\n for(;idx<total;idx+=blockDim.x*gridDim.x){ int col=idx & 8191; float v=(y[idx]-mean[col])*invstd[col]; v=fmaf(v,gamma[col],beta[col]+bias0); y[idx]=v/(1.f+__expf(-v)); }\n}\n__global__ void apply_train_general_kernel(float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,const float* __restrict__ extra_bias,int total,int M,int bias_numel,float divv){\n int idx=blockIdx.x*blockDim.x+threadIdx.x;\n for(;idx<total;idx+=blockDim.x*gridDim.x){ int col=idx%M; float v=(y[idx]-mean[col])*invstd[col]; v=v*gamma[col]+beta[col]; v=(v+extra_bias[(bias_numel==1)?0:col])/divv; y[idx]=v/(1.f+__expf(-v)); }\n}\n__global__ void apply_eval_kernel(float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ running_mean,const float* __restrict__ running_var,const float bias0,int total,float eps){\n int idx=blockIdx.x*blockDim.x+threadIdx.x;\n for(;idx<total;idx+=blockDim.x*gridDim.x){ int col=idx & 8191; float v=(y[idx]-running_mean[col])*rsqrtf(running_var[col]+eps); v=fmaf(v,gamma[col],beta[col]+bias0); y[idx]=v/(1.f+__expf(-v)); }\n}\n__global__ void apply_eval_general_kernel(float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ running_mean,const float* __restrict__ running_var,const float* __restrict__ extra_bias,int total,int M,int bias_numel,float eps,float divv){\n int idx=blockIdx.x*blockDim.x+threadIdx.x;\n for(;idx<total;idx+=blockDim.x*gridDim.x){ int col=idx%M; float v=(y[idx]-running_mean[col])*rsqrtf(running_var[col]+eps); v=v*gamma[col]+beta[col]; v=(v+extra_bias[(bias_numel==1)?0:col])/divv; y[idx]=v/(1.f+__expf(-v)); }\n}\nvoid bn_swish_train_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor extra_bias, double momentum, double eps, double divide_value){\n int N=y.size(0), M=y.size(1); auto mean=torch::empty({M},y.options()); auto inv=torch::empty({M},y.options()); cudaStream_t st=at::cuda::getCurrentCUDAStream();\n dim3 block(32,8); dim3 grid((M+31)/32); reduce_cols_coal_kernel<<<grid,block,0,st>>>(y.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), inv.data_ptr<float>(), N,M,(float)momentum,(float)eps);\n int total=N*M; int blocks=(total+THREADS-1)/THREADS; if(blocks>65535) blocks=65535;\n if (M==8192 && extra_bias.numel()==1 && (float)divide_value==1.0f) apply_train_kernel<<<blocks,THREADS,0,st>>>(y.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),extra_bias.data_ptr<float>()[0],total);\n else apply_train_general_kernel<<<blocks,THREADS,0,st>>>(y.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),extra_bias.data_ptr<float>(),total,M,extra_bias.numel(),(float)divide_value);\n}\nvoid bn_swish_eval_cuda(torch::Tensor y, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor extra_bias, double eps, double divide_value){\n int N=y.size(0), M=y.size(1), total=N*M; int blocks=(total+THREADS-1)/THREADS; if(blocks>65535) blocks=65535; cudaStream_t st=at::cuda::getCurrentCUDAStream();\n if (M==8192 && extra_bias.numel()==1 && (float)divide_value==1.0f) apply_eval_kernel<<<blocks,THREADS,0,st>>>(y.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),extra_bias.data_ptr<float>()[0],total,(float)eps);\n else apply_eval_general_kernel<<<blocks,THREADS,0,st>>>(y.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),extra_bias.data_ptr<float>(),total,M,extra_bias.numel(),(float)eps,(float)divide_value);\n}\n'''\n_ext=load_inline(name='lin_bn_swish_ext_v4',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['linear_bn_swish'],verbose=False,extra_cflags=['-O3'],extra_cuda_cflags=['-O3','--use_fast_math'])\nclass ModelNew(nn.Module):\n def __init__(self,in_features,out_features,bn_eps=1e-5,bn_momentum=0.1,bias_shape=(1,),divide_value=1.0):\n super(ModelNew,self).__init__(); self.matmul=nn.Linear(in_features,out_features); self.bn=nn.BatchNorm1d(out_features,eps=bn_eps,momentum=bn_momentum); self.bias=nn.Parameter(torch.randn(bias_shape)); self.divide_value=divide_value\n def forward(self,x):\n return _ext.linear_bn_swish(x,self.matmul.weight,self.matmul.bias,self.bn.weight,self.bn.bias,self.bn.running_mean,self.bn.running_var,self.bias,self.training,float(self.bn.momentum),float(self.bn.eps),float(self.divide_value))\n"
}