call → run_correctness
callrun_correctness
{
"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 = r'''\n#include <torch/extension.h>\n#include <c10/cuda/CUDAStream.h>\n#include <cuda_runtime_api.h>\n\nvoid eval_kernel_launch(const float*, const float*, const float*, const float*, const float*, const float*, float*, long long, int, float, cudaStream_t);\nvoid partial_stats_launch(const float*, float*, float*, int, int, int, cudaStream_t);\nvoid finalize_stats_launch(const float*, const float*, const float*, float*, float*, float*, float*, int, int, float, float, cudaStream_t);\nvoid train_norm_launch(const float*, const float*, const float*, const float*, const float*, const float*, float*, long long, int, cudaStream_t);\n\ntorch::Tensor eval_forward(torch::Tensor lin, torch::Tensor scale, torch::Tensor running_mean, torch::Tensor running_var,\n torch::Tensor gamma, torch::Tensor beta, double eps) {\n auto out = torch::empty_like(lin);\n eval_kernel_launch(lin.data_ptr<float>(), scale.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),\n gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), lin.numel(), (int)lin.size(1), (float)eps, c10::cuda::getCurrentCUDAStream());\n return out;\n}\n\ntorch::Tensor train_forward(torch::Tensor lin, torch::Tensor scale, torch::Tensor running_mean, torch::Tensor running_var,\n torch::Tensor gamma, torch::Tensor beta, double eps, double momentum) {\n int N=(int)lin.size(0), M=(int)lin.size(1); const int ROWS=256; int chunks=(N+ROWS-1)/ROWS;\n auto opts=lin.options();\n auto ps=torch::empty({M, chunks}, opts); auto psq=torch::empty({M, chunks}, opts);\n auto mean=torch::empty({M}, opts); auto invstd=torch::empty({M}, opts); auto out=torch::empty_like(lin);\n auto stream=c10::cuda::getCurrentCUDAStream();\n partial_stats_launch(lin.data_ptr<float>(), ps.data_ptr<float>(), psq.data_ptr<float>(), N, M, chunks, stream);\n finalize_stats_launch(ps.data_ptr<float>(), psq.data_ptr<float>(), scale.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), N, M, (float)eps, (float)momentum, stream);\n train_norm_launch(lin.data_ptr<float>(), scale.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), lin.numel(), M, stream);\n return out;\n}\n'''\n_cuda = r'''\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n\n__global__ void eval_kernel(const float* __restrict__ lin,const float* __restrict__ scale,const float* __restrict__ mean,const float* __restrict__ var,const float* __restrict__ gamma,const float* __restrict__ beta,float* __restrict__ out,long long n4,int M,float eps){\n long long i4=(long long)blockIdx.x*blockDim.x+threadIdx.x; if(i4<n4){ long long base=i4*4; int c=(int)(base%M); float4 v=reinterpret_cast<const float4*>(lin)[i4]; float a[4]={v.x,v.y,v.z,v.w}; float r[4];\n#pragma unroll\n for(int k=0;k<4;k++){int j=c+k; float t=a[k]*__ldg(scale+j); t=(t-__ldg(mean+j))*rsqrtf(__ldg(var+j)+eps); r[k]=t*__ldg(gamma+j)+__ldg(beta+j);} reinterpret_cast<float4*>(out)[i4]=make_float4(r[0],r[1],r[2],r[3]); }}\n\n__global__ void partial_stats_kernel(const float* __restrict__ lin, float* __restrict__ ps, float* __restrict__ psq, int N, int M, int chunks){\n const int tx=threadIdx.x, ty=threadIdx.y; int j=blockIdx.x*32+tx; int row0=blockIdx.y*256; float sum=0.f, sumq=0.f;\n if(j<M){ int end=min(row0+256,N); for(int r=row0+ty; r<end; r+=8){ float v=lin[(long long)r*M+j]; sum+=v; sumq+=v*v; }}\n __shared__ float ss[8][32]; __shared__ float sq[8][32]; ss[ty][tx]=sum; sq[ty][tx]=sumq; __syncthreads();\n if(ty==0 && j<M){ float s=0.f, q=0.f; #pragma unroll for(int k=0;k<8;k++){s+=ss[k][tx]; q+=sq[k][tx];} ps[(long long)j*chunks+blockIdx.y]=s; psq[(long long)j*chunks+blockIdx.y]=q; }\n}\n\n__global__ void finalize_stats_kernel(const float* __restrict__ ps,const float* __restrict__ psq,const float* __restrict__ scale,float* __restrict__ running_mean,float* __restrict__ running_var,float* __restrict__ mean,float* __restrict__ invstd,int N,int M,int chunks,float eps,float momentum){\n int j=blockIdx.x; float s=0.f,q=0.f; for(int c=threadIdx.x;c<chunks;c+=blockDim.x){ s+=ps[(long long)j*chunks+c]; q+=psq[(long long)j*chunks+c]; }\n __shared__ float ss[128], sq[128]; ss[threadIdx.x]=s; sq[threadIdx.x]=q; __syncthreads(); for(int stride=64; stride>0; stride>>=1){ if(threadIdx.x<stride){ss[threadIdx.x]+=ss[threadIdx.x+stride]; sq[threadIdx.x]+=sq[threadIdx.x+stride];} __syncthreads(); }\n if(threadIdx.x==0){ float sc=scale[j]; float m=(ss[0]/(float)N)*sc; float ex2=(sq[0]/(float)N)*sc*sc; float varb=fmaxf(ex2-m*m,0.f); mean[j]=m; invstd[j]=rsqrtf(varb+eps); running_mean[j]=running_mean[j]*(1.f-momentum)+momentum*m; running_var[j]=running_var[j]*(1.f-momentum)+momentum*(varb*((float)N/(float)(N-1))); }}\n\n__global__ void train_norm_kernel(const float* __restrict__ lin,const float* __restrict__ scale,const float* __restrict__ mean,const float* __restrict__ invstd,const float* __restrict__ gamma,const float* __restrict__ beta,float* __restrict__ out,long long n4,int M){\n long long i4=(long long)blockIdx.x*blockDim.x+threadIdx.x; if(i4<n4){ long long base=i4*4; int c=(int)(base%M); float4 v=reinterpret_cast<const float4*>(lin)[i4]; float a[4]={v.x,v.y,v.z,v.w}; float r[4];\n#pragma unroll\n for(int k=0;k<4;k++){int j=c+k; float t=a[k]*__ldg(scale+j); t=(t-__ldg(mean+j))*__ldg(invstd+j); r[k]=t*__ldg(gamma+j)+__ldg(beta+j);} reinterpret_cast<float4*>(out)[i4]=make_float4(r[0],r[1],r[2],r[3]); }}\n\nvoid eval_kernel_launch(const float* lin,const float* scale,const float* mean,const float* var,const float* gamma,const float* beta,float* out,long long total,int M,float eps,cudaStream_t stream){int th=256; long long n4=total/4; int bl=(int)((n4+th-1)/th); eval_kernel<<<bl,th,0,stream>>>(lin,scale,mean,var,gamma,beta,out,n4,M,eps);} \nvoid partial_stats_launch(const float* lin,float* ps,float* psq,int N,int M,int chunks,cudaStream_t stream){dim3 block(32,8); dim3 grid((M+31)/32,chunks); partial_stats_kernel<<<grid,block,0,stream>>>(lin,ps,psq,N,M,chunks);} \nvoid finalize_stats_launch(const float* ps,const float* psq,const float* scale,float* running_mean,float* running_var,float* mean,float* invstd,int N,int M,float eps,float momentum,cudaStream_t stream){int chunks=(N+255)/256; finalize_stats_kernel<<<M,128,0,stream>>>(ps,psq,scale,running_mean,running_var,mean,invstd,N,M,chunks,eps,momentum);} \nvoid train_norm_launch(const float* lin,const float* scale,const float* mean,const float* invstd,const float* gamma,const float* beta,float* out,long long total,int M,cudaStream_t stream){int th=256; long long n4=total/4; int bl=(int)((n4+th-1)/th); train_norm_kernel<<<bl,th,0,stream>>>(lin,scale,mean,invstd,gamma,beta,out,n4,M);} \n'''\n_ext=load_inline(name='fused_linear_scale_bn_ext3',cpp_sources=[_cpp],cuda_sources=[_cuda],functions=['eval_forward','train_forward'],extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'],verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):\n super(ModelNew,self).__init__(); self.gemm=nn.Linear(in_features,out_features); self.scale=nn.Parameter(torch.randn(scale_shape)); self.bn=nn.BatchNorm1d(out_features,eps=eps,momentum=momentum)\n def forward(self,x):\n lin=F.linear(x,self.gemm.weight,self.gemm.bias)\n if self.bn.training:\n if self.bn.track_running_stats: self.bn.num_batches_tracked.add_(1)\n mom=self.bn.momentum if self.bn.momentum is not None else 0.1\n return _ext.train_forward(lin,self.scale,self.bn.running_mean,self.bn.running_var,self.bn.weight,self.bn.bias,float(self.bn.eps),float(mom))\n return _ext.eval_forward(lin,self.scale,self.bn.running_mean,self.bn.running_var,self.bn.weight,self.bn.bias,float(self.bn.eps))\n"
}