pytorch中使用cuda扩展的实现示例
以下面这个例子作为教程,实现功能是element-wiseadd;
(pytorch中想调用cuda模块,还是用另外使用C编写接口脚本)
第一步:cuda编程的源文件和头文件
//mathutil_cuda_kernel.cu //头文件,最后一个是cuda特有的 #include#include #include #include #include"mathutil_cuda_kernel.h" //获取GPU线程通道信息 dim3cuda_gridsize(intn) { intk=(n-1)/BLOCK+1; intx=k; inty=1; if(x>65535){ x=ceil(sqrt(k)); y=(n-1)/(x*BLOCK)+1; } dim3d(x,y,1); returnd; } //这个函数是cuda执行函数,可以看到细化到了每一个元素 __global__voidbroadcast_sum_kernel(float*a,float*b,intx,inty,intsize) { inti=(blockIdx.x+blockIdx.y*gridDim.x)*blockDim.x+threadIdx.x; if(i>=size)return; intj=i%x;i=i/x; intk=i%y; a[IDX2D(j,k,y)]+=b[k]; } //这个函数是与c语言函数链接的接口函数 voidbroadcast_sum_cuda(float*a,float*b,intx,inty,cudaStream_tstream) { intsize=x*y; cudaError_terr; //上面定义的函数 broadcast_sum_kernel<< >>(a,b,x,y,size); err=cudaGetLastError(); if(cudaSuccess!=err) { fprintf(stderr,"CUDAkernelfailed:%s\n",cudaGetErrorString(err)); exit(-1); } }
#ifndef_MATHUTIL_CUDA_KERNEL #define_MATHUTIL_CUDA_KERNEL #defineIDX2D(i,j,dj)(dj*i+j) #defineIDX3D(i,j,k,dj,dk)(IDX2D(IDX2D(i,j,dj),k,dk)) #defineBLOCK512 #defineMAX_STREAMS512 #ifdef__cplusplus extern"C"{ #endif voidbroadcast_sum_cuda(float*a,float*b,intx,inty,cudaStream_tstream); #ifdef__cplusplus } #endif #endif
第二步:C编程的源文件和头文件(接口函数)
//mathutil_cuda.c //THC是pytorch底层GPU库 #include#include"mathutil_cuda_kernel.h" externTHCState*state; intbroadcast_sum(THCudaTensor*a_tensor,THCudaTensor*b_tensor,intx,inty) { float*a=THCudaTensor_data(state,a_tensor); float*b=THCudaTensor_data(state,b_tensor); cudaStream_tstream=THCState_getCurrentStream(state); //这里调用之前在cuda中编写的接口函数 broadcast_sum_cuda(a,b,x,y,stream); return1; }
intbroadcast_sum(THCudaTensor*a_tensor,THCudaTensor*b_tensor,intx,inty);
第三步:编译,先编译cuda模块,再编译接口函数模块(不能放在一起同时编译)
nvcc-c-omathutil_cuda_kernel.cu.omathutil_cuda_kernel.cu-xcu-Xcompiler-fPIC-arch=sm_52
importos importtorch fromtorch.utils.ffiimportcreate_extension this_file=os.path.dirname(__file__) sources=[] headers=[] defines=[] with_cuda=False iftorch.cuda.is_available(): print('IncludingCUDAcode.') sources+=['src/mathutil_cuda.c'] headers+=['src/mathutil_cuda.h'] defines+=[('WITH_CUDA',None)] with_cuda=True this_file=os.path.dirname(os.path.realpath(__file__)) extra_objects=['src/mathutil_cuda_kernel.cu.o']#这里是编译好后的.o文件位置 extra_objects=[os.path.join(this_file,fname)forfnameinextra_objects] ffi=create_extension( '_ext.cuda_util', headers=headers, sources=sources, define_macros=defines, relative_to=__file__, with_cuda=with_cuda, extra_objects=extra_objects ) if__name__=='__main__': ffi.build()
第四步:调用cuda模块
from_extimportcuda_util#从对应路径中调用编译好的模块 a=torch.randn(3,5).cuda() b=torch.randn(3,1).cuda() mathutil.broadcast_sum(a,b,*map(int,a.size())) #上面等价于下面的效果: a=torch.randn(3,5) b=torch.randn(3,1) a+=b
以上就是本文的全部内容,希望对大家的学习有所帮助,也希望大家多多支持毛票票。
声明:本文内容来源于网络,版权归原作者所有,内容由互联网用户自发贡献自行上传,本网站不拥有所有权,未作人工编辑处理,也不承担相关法律责任。如果您发现有涉嫌版权的内容,欢迎发送邮件至:czq8825#qq.com(发邮件时,请将#更换为@)进行举报,并提供相关证据,一经查实,本站将立刻删除涉嫌侵权内容。