怎么在pytorch中使用cuda扩展

怎么在pytorch中使用cuda扩展？针对这个问题，这篇文章详细介绍了相对应的分析和解答，希望可以帮助更多想解决这个问题的小伙伴找到更简单易行的方法。

第一步：cuda编程的源文件和头文件

//mathutil_cuda_kernel.cu
//头文件，最后一个是cuda特有的
#include<curand.h>
#include<stdio.h>
#include<math.h>
#include<float.h>
#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<<<cuda_gridsize(size),BLOCK,0,stream>>>(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<THC/THC.h>
#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

关于怎么在pytorch中使用cuda扩展问题的解答就分享到这里了，希望以上内容可以对大家有一定的帮助，如果你还有很多疑惑没有解开，可以关注恰卡编程网行业资讯频道了解更多相关知识。