很多时候，我们希望对tensor有一些自定义的操作，一种实现方式就是使用pycuda。本文以实现两个tensor的加法为例，讲解如何实现pycuda与pytorch交互。
1. pycuda的使用方式
首先看下pycuda文档对pycuda的定义：
PyCUDA gives you easy, Pythonic access to Nvidia’s CUDA parallel computation API.
即pycuda是在python中调用CUDA进行并行计算的接口。
使用这个接口的流程为：
1.初始化PyCUDA模块。
2.导入一个用于编译的类。
3.定义一个存储着CUDA代码的字符串。
4.构造一个类对象（编译CUDA代码，使之成为可以被执行的GPU程序）。
5.从类对象中获得一个函数的执行入口。
6.通过这个入口，执行GPU程序。
原理和CUDA c++一致，不太了解的同学可以先去看一下。不同的是kenel函数要定义在一个字符串中，通过构造类对象编译CUDA代码，然后统过类对象中的函数接口进行调用。
下面举一个简单的例子，实现了两个二维数组的相加：
import numpy as np
import pycuda.autoinit #以自动的方式对pycuda进行初始
from pycuda.compiler import SourceModule # 编译kernel函数的类
import pycuda.gpuarray as gpuarray
# 通过字符串定义kernel函数
kernel_code = r"""
void __global__ add(const float *x, const float *y, float *z)
{
const int n = threadIdx.y*blockDim.x+threadIdx.x;
z[n] = x[n] + y[n];
}
"""
# 编译kernel函数
mod = SourceModule(kernel_code)
# 获取函数接口
add= mod.get_function("add")
# 定义三个二维数组并转换成pycuda中的gpuarray
A = np.zeros((10, 10))
B = np.ones((10, 10))
C = np.zeros((10, 10))
A_GPU = gpuarray.to_gpu(A.astype(np.float32))
B_GPU = gpuarray.to_gpu(B.astype(np.float32))
C_GPU = gpuarray.to_gpu(B.astype(np.float32))
# 执行kernel函数，和cuda c++一样定义griddim和blockdim;
add(A_GPU, B_GPU, C_GPU, grid=(1,1,1), block=(10,10,1))
C = C_GPU.get()
print(C)
输出：
[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]]
编译kernel代码的时间比较耗时，编译一次之后后面就快了。
2. pycuda与pytorch交互处理tensor
要想完成两者的交互，需要定义以下这个类，并利用tensor初始化一个对象。
class Holder(pycuda.driver.PointerHolderBase):
def __init__(self, t):
super(Holder, self).__init__()
self.t = t
self.gpudata = t.data_ptr()
def get_pointer():
return self.t.data_ptr()
还是上面的二维矩阵相加的例子，但A,B,C换成pytorch的tensor:
import torch
import pycuda.autoinit
from pycuda.compiler import SourceModule
x = torch.cuda.FloatTensor(8)
class Holder(pycuda.driver.PointerHolderBase):
def __init__(self, t):
super(Holder, self).__init__()
self.t = t
self.gpudata = t.data_ptr()
def get_pointer():
return self.t.data_ptr()
kernel_code = r"""
void __global__ add(const float *x, const float *y, float *z)
{
const int n = threadIdx.y*blockDim.x+threadIdx.x;
z[n] = x[n] + y[n];
}
"""
mod = SourceModule(kernel_code)
add= mod.get_function("add")
A = torch.zeros((10, 10), device="cuda:0")
B = torch.ones((10, 10), device="cuda:0")
C = torch.rand((10, 10), device="cuda:0")
add(Holder(A), Holder(B), Holder(C), grid=(1,1,1), block=(10,10,1))
print(C)
输出：
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]], device='cuda:0')
注意第四行这句话如果没有的话会报错：
Traceback (most recent call last):
File "/data/zhuyinghao/jarvis-dain/frame_rate_up_conversion/pycuda_ex2.py", line 27, in <module>
add(Holder(A), Holder(B), Holder(C), grid=(1,1,1), block=(10,10,1))
File "/root/anaconda3/lib/python3.7/site-packages/pycuda/driver.py", line 436, in function_call
func._set_block_shape(*block)
pycuda._driver.LogicError: cuFuncSetBlockShape failed: invalid resource handle
3.参考文献
[1] interop-between-pycuda-and-pytorch
[2] pyCUDA加速你的python代码