callrun_correctness
{
"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>\ntorch::Tensor fused_bn_relu_dilated_conv_forward(torch::Tensor x, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor conv_weight, torch::Tensor conv_bias, bool training, double eps, int64_t dilation);\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 CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be CUDA\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n__global__ void bn_stats_kernel(const float* __restrict__ x, float* __restrict__ mean, float* __restrict__ invstd, int N, int C, int L, float eps){\n int c=blockIdx.x, tid=threadIdx.x, total=N*L; float sum=0.f,sumsq=0.f;\n for(int idx=tid;idx<total;idx+=blockDim.x){int n=idx/L;int t=idx-n*L;float v=x[(n*C+c)*L+t];sum+=v;sumsq+=v*v;}\n __shared__ float ssum[256], ssq[256]; ssum[tid]=sum; ssq[tid]=sumsq; __syncthreads();\n for(int s=128;s>0;s>>=1){if(tid<s){ssum[tid]+=ssum[tid+s];ssq[tid]+=ssq[tid+s];}__syncthreads();}\n if(tid==0){float m=ssum[0]/(float)total;float var=ssq[0]/(float)total-m*m;if(var<0.f)var=0.f;mean[c]=m;invstd[c]=rsqrtf(var+eps);} }\n__device__ __forceinline__ float tf32_rna(float x){ float y; asm(\"cvt.rna.tf32.f32 %0, %1;\" : \"=f\"(y) : \"f\"(x)); return y; }\n__global__ void fused_conv_kernel(const float* __restrict__ x,const float* __restrict__ bn_weight,const float* __restrict__ bn_bias,const float* __restrict__ running_mean,const float* __restrict__ running_var,const float* __restrict__ mean,const float* __restrict__ invstd,const float* __restrict__ conv_weight,const float* __restrict__ conv_bias,float* __restrict__ out,int N,int C,int L,int K,int dilation,float eps,int training){\n int t=blockIdx.x*blockDim.x+threadIdx.x; int oc0=blockIdx.y*4; int n=blockIdx.z; if(t>=L)return;\n float acc0=0.f,acc1=0.f,acc2=0.f,acc3=0.f;\n for(int ic=0;ic<C;++ic){\n float sc,bs; if(training){sc=bn_weight[ic]*invstd[ic];bs=bn_bias[ic]-mean[ic]*sc;}else{sc=bn_weight[ic]*rsqrtf(running_var[ic]+eps);bs=bn_bias[ic]-running_mean[ic]*sc;}\n #pragma unroll\n for(int kk=0;kk<3;++kk){\n int tt=t+(kk-1)*dilation; if((unsigned)tt<(unsigned)L){\n float v=x[(n*C+ic)*L+tt]*sc+bs; v=v>0.f?v:0.f; float vt=tf32_rna(v);\n int wbase=(oc0*C+ic)*K+kk;\n acc0=fmaf(vt,tf32_rna(conv_weight[wbase]),acc0);\n acc1=fmaf(vt,tf32_rna(conv_weight[wbase+C*K]),acc1);\n acc2=fmaf(vt,tf32_rna(conv_weight[wbase+2*C*K]),acc2);\n acc3=fmaf(vt,tf32_rna(conv_weight[wbase+3*C*K]),acc3);\n }} }\n int base=(n*C+oc0)*L+t; out[base]=acc0+conv_bias[oc0]+x[base]; out[base+L]=acc1+conv_bias[oc0+1]+x[base+L]; out[base+2*L]=acc2+conv_bias[oc0+2]+x[base+2*L]; out[base+3*L]=acc3+conv_bias[oc0+3]+x[base+3*L]; }\ntorch::Tensor fused_bn_relu_dilated_conv_forward(torch::Tensor x, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor conv_weight, torch::Tensor conv_bias, bool training, double eps, int64_t dilation){\n CHECK_CUDA(x); CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(conv_weight); int N=x.size(0),C=x.size(1),L=x.size(2),K=conv_weight.size(2); auto out=torch::empty_like(x); torch::Tensor mean,invstd; cudaStream_t stream=at::cuda::getCurrentCUDAStream();\n if(training){mean=torch::empty({C},x.options());invstd=torch::empty({C},x.options());bn_stats_kernel<<<C,256,0,stream>>>(x.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),N,C,L,(float)eps);}else{mean=torch::empty({1},x.options());invstd=torch::empty({1},x.options());}\n dim3 block(128); dim3 grid((L+127)/128,C/4,N); fused_conv_kernel<<<grid,block,0,stream>>>(x.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),training?mean.data_ptr<float>():nullptr,training?invstd.data_ptr<float>():nullptr,conv_weight.data_ptr<float>(),conv_bias.data_ptr<float>(),out.data_ptr<float>(),N,C,L,K,(int)dilation,(float)eps,training?1:0); return out; }\n'''\n_ext=load_inline(name='fused_bn_relu_dilated_conv_ext_v4',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['fused_bn_relu_dilated_conv_forward'],verbose=False,extra_cuda_cflags=['-O3'])\nclass ModelNew(nn.Module):\n def __init__(self, channels, kernel_size, dilation):\n super(ModelNew,self).__init__(); self.norm=nn.BatchNorm1d(channels); self.activation=nn.ReLU(); self.conv=nn.Conv1d(channels,channels,kernel_size,padding=dilation,dilation=dilation); self.dilation=dilation\n def forward(self,x):\n return _ext.fused_bn_relu_dilated_conv_forward(x.contiguous(),self.norm.weight,self.norm.bias,self.norm.running_mean,self.norm.running_var,self.conv.weight,self.conv.bias,self.training,self.norm.eps,self.dilation)\n"
}