callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\nvoid em_forward_out_cuda(torch::Tensor x, torch::Tensor mus, torch::Tensor covs, torch::Tensor log_pi,\n torch::Tensor new_mus, torch::Tensor new_covs, torch::Tensor new_log_pi,\n torch::Tensor log_lik, torch::Tensor resp, torch::Tensor L, torch::Tensor logc, torch::Tensor partial);\nvoid em_forward_out(torch::Tensor x, torch::Tensor mus, torch::Tensor covs, torch::Tensor log_pi,\n torch::Tensor new_mus, torch::Tensor new_covs, torch::Tensor new_log_pi,\n torch::Tensor log_lik, torch::Tensor resp, torch::Tensor L, torch::Tensor logc, torch::Tensor partial) {\n em_forward_out_cuda(x, mus, covs, log_pi, new_mus, new_covs, new_log_pi, log_lik, resp, L, logc, partial);\n}\n'''\n# same cuda kernels but out wrapper no allocation\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 DCONST 8\n#define KCONST 5\n__global__ void prep_kernel(const float* __restrict__ covs,const float* __restrict__ log_pi,float* __restrict__ Lout,float* __restrict__ logc){int k=blockIdx.x;if(threadIdx.x||k>=KCONST)return;float m=log_pi[0];#pragma unroll\nfor(int i=1;i<KCONST;i++)m=fmaxf(m,log_pi[i]);float s=0.f;#pragma unroll\nfor(int i=0;i<KCONST;i++)s+=expf(log_pi[i]-m);float lse_pi=m+logf(s);float L[DCONST][DCONST];#pragma unroll\nfor(int i=0;i<DCONST;i++){for(int j=0;j<DCONST;j++)L[i][j]=0.f;}float logdet=0.f;#pragma unroll\nfor(int i=0;i<DCONST;i++){#pragma unroll\nfor(int j=0;j<=i;j++){float sum=covs[(k*DCONST+i)*DCONST+j];#pragma unroll\nfor(int p=0;p<j;p++)sum-=L[i][p]*L[j][p];if(i==j){sum=fmaxf(sum,1.e-20f);float v=sqrtf(sum);L[i][j]=v;logdet+=2.f*logf(v);}else L[i][j]=sum/L[j][j];}}#pragma unroll\nfor(int i=0;i<DCONST;i++){#pragma unroll\nfor(int j=0;j<DCONST;j++)Lout[(k*DCONST+i)*DCONST+j]=L[i][j];}logc[k]=(log_pi[k]-lse_pi)-0.5f*((float)DCONST*1.8378770664093453f+logdet);} \n__global__ void estep_kernel(const float* __restrict__ x,const float* __restrict__ mus,const float* __restrict__ L,const float* __restrict__ logc,float* __restrict__ resp,float* __restrict__ partial_ll,int N){extern __shared__ float sh[];int tid=threadIdx.x;int global=blockIdx.x*blockDim.x+tid;int stride=blockDim.x*gridDim.x;float ll=0.f;for(int n=global;n<N;n+=stride){float xv[DCONST];#pragma unroll\nfor(int d=0;d<DCONST;d++)xv[d]=x[n*DCONST+d];float lp[KCONST];#pragma unroll\nfor(int k=0;k<KCONST;k++){float y[DCONST];float maha=0.f;#pragma unroll\nfor(int i=0;i<DCONST;i++){float sum=xv[i]-mus[k*DCONST+i];#pragma unroll\nfor(int j=0;j<i;j++)sum-=L[(k*DCONST+i)*DCONST+j]*y[j];float yi=sum/L[(k*DCONST+i)*DCONST+i];y[i]=yi;maha+=yi*yi;}lp[k]=logc[k]-0.5f*maha;}float mx=lp[0];#pragma unroll\nfor(int k=1;k<KCONST;k++)mx=fmaxf(mx,lp[k]);float se=0.f;float e[KCONST];#pragma unroll\nfor(int k=0;k<KCONST;k++){e[k]=expf(lp[k]-mx);se+=e[k];}ll+=mx+logf(se);float inv=1.f/se;#pragma unroll\nfor(int k=0;k<KCONST;k++)resp[n*KCONST+k]=e[k]*inv;}sh[tid]=ll;__syncthreads();for(int off=blockDim.x/2;off>0;off>>=1){if(tid<off)sh[tid]+=sh[tid+off];__syncthreads();}if(tid==0)partial_ll[blockIdx.x]=sh[0];}\n__global__ void mstep_kernel(const float* __restrict__ x,const float* __restrict__ resp,float* __restrict__ new_mus,float* __restrict__ new_covs,float* __restrict__ new_log_pi,int N){int k=blockIdx.x,tid=threadIdx.x;const int QN=36,STR=45;extern __shared__ float sh[];float nk=0.f;float ss[DCONST];float qq[QN];#pragma unroll\nfor(int d=0;d<DCONST;d++)ss[d]=0.f;#pragma unroll\nfor(int q=0;q<QN;q++)qq[q]=0.f;for(int n=tid;n<N;n+=blockDim.x){float r=resp[n*KCONST+k];nk+=r;float xv[DCONST];#pragma unroll\nfor(int d=0;d<DCONST;d++){float v=x[n*DCONST+d];xv[d]=v;ss[d]+=r*v;}int qi=0;#pragma unroll\nfor(int i=0;i<DCONST;i++){float xi=xv[i];#pragma unroll\nfor(int j=i;j<DCONST;j++)qq[qi++]+=r*xi*xv[j];}}float* base=sh+tid*STR;base[0]=nk;#pragma unroll\nfor(int d=0;d<DCONST;d++)base[1+d]=ss[d];#pragma unroll\nfor(int q=0;q<QN;q++)base[1+DCONST+q]=qq[q];__syncthreads();for(int off=blockDim.x/2;off>0;off>>=1){if(tid<off){float* a=sh+tid*STR;float* b=sh+(tid+off)*STR;#pragma unroll\nfor(int t=0;t<STR;t++)a[t]+=b[t];}__syncthreads();}if(tid==0){float denom=fmaxf(sh[0],1.e-8f);float invn=1.f/denom;float mu[DCONST];#pragma unroll\nfor(int d=0;d<DCONST;d++){mu[d]=sh[1+d]*invn;new_mus[k*DCONST+d]=mu[d];}new_log_pi[k]=logf(denom)-logf((float)N);int qi=0;#pragma unroll\nfor(int i=0;i<DCONST;i++){#pragma unroll\nfor(int j=i;j<DCONST;j++){float c=sh[1+DCONST+qi++]*invn-mu[i]*mu[j];if(i==j)c+=1.e-4f;new_covs[(k*DCONST+i)*DCONST+j]=c;new_covs[(k*DCONST+j)*DCONST+i]=c;}}}}\n__global__ void reduce_ll_kernel(const float* __restrict__ partial,float* __restrict__ out,int M){extern __shared__ float sh[];int tid=threadIdx.x;float s=0.f;for(int i=tid;i<M;i+=blockDim.x)s+=partial[i];sh[tid]=s;__syncthreads();for(int off=blockDim.x/2;off>0;off>>=1){if(tid<off)sh[tid]+=sh[tid+off];__syncthreads();}if(tid==0)*out=sh[0];}\nvoid em_forward_out_cuda(torch::Tensor x,torch::Tensor mus,torch::Tensor covs,torch::Tensor log_pi,torch::Tensor new_mus,torch::Tensor new_covs,torch::Tensor new_log_pi,torch::Tensor log_lik,torch::Tensor resp,torch::Tensor L,torch::Tensor logc,torch::Tensor partial){int N=(int)x.size(0);int threads_e=128;int blocks_e=(N+threads_e-1)/threads_e;if(blocks_e<1)blocks_e=1;if(blocks_e>32)blocks_e=32;cudaStream_t stream=at::cuda::getCurrentCUDAStream();prep_kernel<<<KCONST,1,0,stream>>>(covs.data_ptr<float>(),log_pi.data_ptr<float>(),L.data_ptr<float>(),logc.data_ptr<float>());estep_kernel<<<blocks_e,threads_e,threads_e*sizeof(float),stream>>>(x.data_ptr<float>(),mus.data_ptr<float>(),L.data_ptr<float>(),logc.data_ptr<float>(),resp.data_ptr<float>(),partial.data_ptr<float>(),N);int threads_m=128;size_t smem_m=(size_t)threads_m*45*sizeof(float);mstep_kernel<<<KCONST,threads_m,smem_m,stream>>>(x.data_ptr<float>(),resp.data_ptr<float>(),new_mus.data_ptr<float>(),new_covs.data_ptr<float>(),new_log_pi.data_ptr<float>(),N);reduce_ll_kernel<<<1,128,128*sizeof(float),stream>>>(partial.data_ptr<float>(),log_lik.data_ptr<float>(),blocks_e);}\n'''\n\n_ext = load_inline(name='em_gmm_ext_out_v1', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['em_forward_out'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self, dim, num_components):\n super().__init__(); self.dim=dim; self.K=num_components\n self.register_buffer('mus', torch.randn(num_components, dim))\n self.register_buffer('covs', torch.eye(dim).unsqueeze(0).expand(num_components,-1,-1).clone())\n self.register_buffer('log_pi', torch.zeros(num_components))\n self._cache = None\n def _ensure_cache(self, x):\n N=x.shape[0]; dev=x.device\n if self._cache is None or self._cache[0].shape[0] != N or self._cache[0].device != dev:\n self._cache = (torch.empty((N,self.K),device=dev,dtype=x.dtype), torch.empty((self.K,self.dim,self.dim),device=dev,dtype=x.dtype), torch.empty((self.K,),device=dev,dtype=x.dtype), torch.empty((32,),device=dev,dtype=x.dtype), torch.empty((self.K,self.dim),device=dev,dtype=x.dtype), torch.empty((self.K,self.dim,self.dim),device=dev,dtype=x.dtype), torch.empty((self.K,),device=dev,dtype=x.dtype), torch.empty((),device=dev,dtype=x.dtype))\n return self._cache\n def forward(self,x:torch.Tensor)->tuple:\n x=x.contiguous(); resp,L,logc,partial,nm,nc,nlp,ll=self._ensure_cache(x)\n _ext.em_forward_out(x,self.mus.contiguous(),self.covs.contiguous(),self.log_pi.contiguous(),nm,nc,nlp,ll,resp,L,logc,partial)\n return nm,nc,nlp,ll\n"
}