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>\n#include <ATen/cuda/CUDAContext.h>\nvoid launch_conv1x1_64_128(const float* x, const float* w, const float* b, float* y, int N, int H, int Wd, int has_bias, cudaStream_t stream);\ntorch::Tensor conv1x1_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias_opt) {\n int N=x.size(0), H=x.size(2), Wd=x.size(3);\n auto y=torch::empty({N,(int64_t)128,H,Wd}, x.options());\n const float* bp=nullptr; int hb=0;\n if (bias_opt.has_value() && bias_opt.value().defined()) { bp=bias_opt.value().data_ptr<float>(); hb=1; }\n launch_conv1x1_64_128(x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), N,H,Wd,hb, at::cuda::getCurrentCUDAStream());\n return y;\n}\n'''\n\ncuda_src = r'''\n#include <cuda_runtime.h>\n#define BM 128\n#define BN 32\n#define CI 64\n#define CO 128\n__global__ __launch_bounds__(256,3) void conv1x1_64_128_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ y,int N,int H,int Wd,int has_bias){\n int HW=H*Wd, M=N*HW, m0=blockIdx.x*BM, o0=blockIdx.y*BN, tid=threadIdx.y*16+threadIdx.x;\n __shared__ float sx[CI*BM]; __shared__ float sw[CI*BN];\n for(int idx=tid; idx<CI*BM; idx+=256){int k=idx/BM,p=idx-k*BM,m=m0+p; float v=0.f; if(m<M){int n=m/HW, hw=m-n*HW; v=x[((n*CI+k)*HW)+hw];} sx[idx]=v;}\n for(int idx=tid; idx<CI*BN; idx+=256){int k=idx/BN,o=idx-k*BN; sw[idx]=w[((o0+o)*CI)+k];}\n __syncthreads();\n int tx=threadIdx.x, ty=threadIdx.y, pbase=ty*8, obase=tx*2;\n float a00=0,a01=0,a10=0,a11=0,a20=0,a21=0,a30=0,a31=0,a40=0,a41=0,a50=0,a51=0,a60=0,a61=0,a70=0,a71=0;\n #pragma unroll\n for(int k=0;k<CI;k++){float w0=sw[k*BN+obase], w1=sw[k*BN+obase+1]; const float* xs=&sx[k*BM+pbase]; float x0=xs[0],x1=xs[1],x2=xs[2],x3=xs[3],x4=xs[4],x5=xs[5],x6=xs[6],x7=xs[7];\n a00=fmaf(x0,w0,a00);a01=fmaf(x0,w1,a01);a10=fmaf(x1,w0,a10);a11=fmaf(x1,w1,a11);a20=fmaf(x2,w0,a20);a21=fmaf(x2,w1,a21);a30=fmaf(x3,w0,a30);a31=fmaf(x3,w1,a31);\n a40=fmaf(x4,w0,a40);a41=fmaf(x4,w1,a41);a50=fmaf(x5,w0,a50);a51=fmaf(x5,w1,a51);a60=fmaf(x6,w0,a60);a61=fmaf(x6,w1,a61);a70=fmaf(x7,w0,a70);a71=fmaf(x7,w1,a71);}\n int oA=o0+obase,oB=oA+1; if(has_bias){float b0=b[oA],b1=b[oB]; a00+=b0;a10+=b0;a20+=b0;a30+=b0;a40+=b0;a50+=b0;a60+=b0;a70+=b0; a01+=b1;a11+=b1;a21+=b1;a31+=b1;a41+=b1;a51+=b1;a61+=b1;a71+=b1;}\n int mb=m0+pbase; float v0[8]={a00,a10,a20,a30,a40,a50,a60,a70}, v1[8]={a01,a11,a21,a31,a41,a51,a61,a71};\n #pragma unroll\n for(int i=0;i<8;i++){int m=mb+i; if(m<M){int nn=m/HW,hwi=m-nn*HW; y[((nn*CO+oA)*HW)+hwi]=v0[i]; y[((nn*CO+oB)*HW)+hwi]=v1[i];}}\n}\nextern \"C\" void launch_conv1x1_64_128(const float* x,const float* w,const float* b,float* y,int N,int H,int Wd,int has_bias,cudaStream_t stream){dim3 block(16,16); dim3 grid((N*H*Wd+BM-1)/BM,CO/BN); conv1x1_64_128_kernel<<<grid,block,0,stream>>>(x,w,b,y,N,H,Wd,has_bias);} \n'''\n\n_ext = load_inline(name='conv1x1_ext_v4', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['conv1x1_forward'], extra_cuda_cflags=['-O3'], verbose=False, no_implicit_headers=True)\nclass ModelNew(nn.Module):\n def __init__(self, in_channels:int, out_channels:int, bias:bool=False):\n super(ModelNew,self).__init__(); self.conv1d=nn.Conv2d(in_channels,out_channels,kernel_size=1,stride=1,padding=0,bias=bias)\n def forward(self,x:torch.Tensor)->torch.Tensor:\n return _ext.conv1x1_forward(x,self.conv1d.weight,self.conv1d.bias)\n"
}