call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\nvoid stats_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor var, torch::Tensor running_mean, torch::Tensor running_var, double momentum);\nvoid bn_relu_eval_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, torch::Tensor y, double eps);\nvoid conv3_cuda(torch::Tensor y, torch::Tensor residual, torch::Tensor w, torch::Tensor b, torch::Tensor out, int64_t dilation);\n\ntorch::Tensor forward_ext(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var,\n torch::Tensor conv_w, torch::Tensor conv_b, int64_t dilation, double eps, bool training, double momentum) {\n auto y = torch::empty_like(x);\n if (training) {\n auto mean = torch::empty({x.size(1)}, x.options());\n auto var = torch::empty({x.size(1)}, x.options());\n stats_cuda(x, mean, var, running_mean, running_var, momentum);\n bn_relu_eval_cuda(x, gamma, beta, mean, var, y, eps);\n } else {\n bn_relu_eval_cuda(x, gamma, beta, running_mean, running_var, y, eps);\n }\n auto out = torch::empty_like(x);\n conv3_cuda(y, x, conv_w, conv_b, out, dilation);\n return out;\n}\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void stats_kernel(const float* __restrict__ x, float* __restrict__ mean, float* __restrict__ var,\n float* __restrict__ running_mean, float* __restrict__ running_var,\n int N, int C, int L, float momentum) {\n int c = blockIdx.x, tid = threadIdx.x, M = N * L;\n float sum=0.f, sum2=0.f;\n for (int i=tid;i<M;i+=blockDim.x){ int n=i/L; int l=i-n*L; float v=x[(n*C+c)*L+l]; sum+=v; sum2+=v*v; }\n __shared__ float s1[256], s2[256]; s1[tid]=sum; s2[tid]=sum2; __syncthreads();\n for(int s=128;s>0;s>>=1){ if(tid<s){s1[tid]+=s1[tid+s]; s2[tid]+=s2[tid+s];} __syncthreads(); }\n if(tid==0){ float m=s1[0]/(float)M; float v=s2[0]/(float)M-m*m; v=v<0.f?0.f:v; mean[c]=m; var[c]=v; running_mean[c]=running_mean[c]*(1.f-momentum)+momentum*m; running_var[c]=running_var[c]*(1.f-momentum)+momentum*v*((float)M/(float)(M-1)); }\n}\n\n__global__ void bn_relu_kernel(const float* __restrict__ x,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ var,float* __restrict__ y,int total,int L,int C,float eps){\n int idx=blockIdx.x*blockDim.x+threadIdx.x;\n if(idx<total){ int c=(idx/L)%C; float v=(x[idx]-mean[c])*(gamma[c]*rsqrtf(var[c]+eps))+beta[c]; y[idx]=v>0.f?v:0.f; }\n}\n\n__global__ void conv3_kernel(const float* __restrict__ y, const float* __restrict__ residual, const float* __restrict__ w, const float* __restrict__ b, float* __restrict__ out, int N, int C, int L, int d){\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = N*C*L;\n if (idx >= total) return;\n int l = idx % L;\n int t = idx / L;\n int co = t % C;\n int n = t / C;\n const int base_n = n * C * L;\n const float* __restrict__ wp = w + co * C * 3;\n float acc = b[co];\n int lm = l - d;\n int lp = l + d;\n if (lm >= 0 && lp < L) {\n #pragma unroll 4\n for (int ci=0; ci<C; ++ci) {\n const float* yp = y + base_n + ci*L;\n const float* wci = wp + ci*3;\n acc = fmaf(yp[lm], wci[0], acc);\n acc = fmaf(yp[l], wci[1], acc);\n acc = fmaf(yp[lp], wci[2], acc);\n }\n } else if (lm >= 0) {\n #pragma unroll 4\n for (int ci=0; ci<C; ++ci) {\n const float* yp = y + base_n + ci*L;\n const float* wci = wp + ci*3;\n acc = fmaf(yp[lm], wci[0], acc);\n acc = fmaf(yp[l], wci[1], acc);\n }\n } else if (lp < L) {\n #pragma unroll 4\n for (int ci=0; ci<C; ++ci) {\n const float* yp = y + base_n + ci*L;\n const float* wci = wp + ci*3;\n acc = fmaf(yp[l], wci[1], acc);\n acc = fmaf(yp[lp], wci[2], acc);\n }\n } else {\n #pragma unroll 4\n for (int ci=0; ci<C; ++ci) {\n acc = fmaf(y[base_n + ci*L + l], wp[ci*3+1], acc);\n }\n }\n out[idx] = acc + residual[idx];\n}\n\nvoid stats_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor var, torch::Tensor running_mean, torch::Tensor running_var, double momentum){ int N=x.size(0),C=x.size(1),L=x.size(2); stats_kernel<<<C,256>>>(x.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),N,C,L,(float)momentum); }\nvoid bn_relu_eval_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, torch::Tensor y, double eps){ int C=x.size(1),L=x.size(2),total=x.numel(); int threads=256,blocks=(total+threads-1)/threads; bn_relu_kernel<<<blocks,threads>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),y.data_ptr<float>(),total,L,C,(float)eps); }\nvoid conv3_cuda(torch::Tensor y, torch::Tensor residual, torch::Tensor w, torch::Tensor b, torch::Tensor out, int64_t dilation){ int N=y.size(0),C=y.size(1),L=y.size(2),total=y.numel(); int threads=256,blocks=(total+threads-1)/threads; conv3_kernel<<<blocks,threads>>>(y.data_ptr<float>(),residual.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),out.data_ptr<float>(),N,C,L,(int)dilation); }\n'''\n_ext = load_inline(name='dilated_bn_conv_direct_v1', cpp_sources=[cpp_src], cuda_sources=[cuda_src], functions=['forward_ext'], verbose=False, extra_cflags=['-O3'], 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.forward_ext(x.contiguous(), self.norm.weight, self.norm.bias, self.norm.running_mean, self.norm.running_var, self.conv.weight, self.conv.bias, self.dilation, self.norm.eps, self.training, self.norm.momentum)\n"
}