pytorch中使用cuda扩展的实现示例

2023-08-01 17:09:04 344

以下面这个例子作为教程，实现功能是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（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。

pytorch中使用cuda扩展的实现示例

热门推荐

随机推荐