pycuda pytorch交互——使用pycuda处理pytorch tensor

很多时候，我们希望对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代码

今天的文章pycuda pytorch交互——使用pycuda处理pytorch tensor分享到此就结束了，感谢您的阅读，如果确实帮到您，您可以动动手指转发给其他人。