在Python里寻求Thrust是否搞错了什么

本文主要讨论：

解释型语言中的GPU编程
ThrustRTC 项目, 以及如何在Python中使用CUDA模板库

解释型语言与GPU编程

众所周知，C++ 是CUDA的默认语言。

在解释型语言（在以人工智能为代表的领域）如此流行的今天，我们发现，以CUDA为代表的GPU并行编程依然主要依靠C/C++等编译型语言。这也不难理解。因为同样的代码如果让解释器为每个线程单独解释一遍，属于没有必要的重复计算。为了追求高性能，解释执行GPU代码是不考虑的。不管如何设计，GPU编程一个基本的原则是：一次编译 ->成千上万次的执行。

于是，作为习惯Python的用户，如果想使用CUDA加速的话，无非以下几个选择：

使用 Numba，host 代码和 device 代码都用 Python 编写

Numba这个技术，本质上是把一部分Python代码编译执行。由于语法还是Python，它无法使用完整的C++语法特性，也和C++编写的库不兼容。同时，由于是编译，而非解释执行，它也无法完整地发挥解释型语言的灵活性。结果就是，能做的事情非常的简单和局限。

当然对于Numba我们不能苛求。这样一个大胆的设计，能够工作已经是奇迹了。

预编译全部的 device 代码，通过 host API 与 Python交互

这个是大多数GPU算法库的选择。这种方法存在以下几个缺点：

代码膨胀。这种库一般体积比较庞大。因为往往要为每一代需要支持的GPU单独生成一份二进制代码。如果是模板库的话，这个份数还要乘上需要支持的数据类型的数量
无法从外部扩展。一个CUDA模块一旦编译好了，想从外部加入自定义的代码将非常困难，也不支持回调。

使用运行时编译(RTC)技术

CUDA 从 7.x 开始增加了一个新模块叫NVRTC。有了它我们可以在运行时把字符串形式的CUDA 代码编译成 PTX 代码，然后通过CUDA Driver API来调用其中的Kernel。这个技术为我们提供了一种新的可能性，那就是，我们依然用C++来编写device 代码，但是将编译工作留到主程序的运行时来完成。在此之前，device 代码只是程序中内嵌的一段字符串。显然，这段字符串是可以跨越主程序的编程语言的，它既可以来自C++代码，也可以来自Python代码，更可以由这两部分系统，结合运行时的数据类型等信息综合而成。

由于device代码就是CUDA C++代码，在其中我们可以使用全部的CUDA C++特性，以及引用已有的头文件库。由于是运行时编译，我们只需要根据调用的环境为特定的硬件和特定的数据类型编译特定的二进制代码。同时，由于我们完全掌握源代码，可以在运行时随意向其中加入自定义的代码。

扫描二维码关注公众号，回复： 5884789 查看本文章

使用这个技术的一个代价是会引入编译时间。为了减少对同样代码的重复编译，可以在内存设置Kernel缓存，在磁盘中设置PTX代码的缓存。（PTX到cubin的缓存已经由driver提供。）正如Python使用了__pycache__来加快host代码的解释速度，使用RTC技术时，我们同样可以通过缓存来加快device代码的分析和编译速度。

让 Python 也可以使用 CUDA 的模板库

一些优秀的CUDA算法库，如Thrust，使用了模板技术。这些库以纯头文件形式提供。因为是源代码，本身不存在膨胀问题。类型多态性、扩展性均在编译期实现，这种设计使得这个库只适合于C++。Numba 里显然无法使用Thrust。而如果采用预编译，则会失去可扩展性，导致代码膨胀，并在一定程度上失去类型多态（只能预编译有限的类型）。

Thrust 提供的功能非常的基础和常用，包括查找、规约、排序等等。ThrustRTC项目采用运行时编译的技术路线，目的是让Python用户也能使用类似Thrust的这些功能。

ThrustRTC 的设计包含以下几个核心概念：

Context: 每一个Context管理着一组需要包含的头文件，全局常数，并提供一个Kernel缓存，遇到已经编译过的代码可以直接找到对应的Kernel而无需再次编译。Context 为各种GPU算法提供了运行时编译和运行的环境。
DeviceViewable: device 可见对象。这些对象本身是host端的对象，但是预先规定了向device 端传递数据的方法。每一个DeviceViewable类对应于一个device端的数据类型作为它的 view 类型。这样一来，我们就可以在 host 端管理这些对象，同时随时可以把数据传递给device端的程序来处理。
DVVector: 最重要的一个DeviceViewable类，相当于Thrust里的device_vector，是ThrustRTC 中各种算法主要操作的数据对象。它对应的 view 类型是一个名为VectorView的结构体。

通过Context提供的接口，只要以字符串形式提供代码，随时可以编译、执行一个CUDA Kernel。

import ThrustRTC as trtc

trtc.set_ptx_cache('__ptx_cache__')
ctx = trtc.Context()

darr = trtc.device_vector_from_list(ctx, [3.0, 5.0, 7.0, 9.0 , 11.0], 'float')
ctx.launch_kernel(1,128, {'arr': darr}, 
	'''
	size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
	if (idx >= arr.size) return;
	arr[idx]*=100.0f;
	''')

其中 {'arr': darr} 是参数映射表，代表将 darr 这个对象映射到 device 端名为 arr 的 view。于是在 device 代码里，我们用 arr 来操作这个数组。

CUDA Kernel 也可以先定义，后执行。注意 Kernel 对象起到的作用只是暂存一下代码，并不进行编译。

import ThrustRTC as trtc
import numpy as np

trtc.set_ptx_cache('__ptx_cache__')
ctx = trtc.Context()

kernel = trtc.Kernel(['arr_in', 'arr_out', 'k'],
	'''
	size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
	if (idx >= arr_in.size) return;
	arr_out[idx] = arr_in[idx]*k;
	''')

harr = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype='float32')
darr = trtc.device_vector_from_numpy(ctx, harr)

darr_out = trtc.device_vector(ctx, 'float', 5)
kernel.launch(ctx, 1,128, [darr, darr_out, trtc.DVFloat(10.0)])

这里 Kernel 实际已经具有了模板的性质。所有参数 arr_in, arr_out, k 的类型都是在调用Kernel的时候才决定的。

Kernel 还有一个简化版本，叫做 For，专门用于一维的Kernel:

import ThrustRTC as trtc
import numpy as np

trtc.set_ptx_cache('__ptx_cache__')
ctx = trtc.Context()

forLoop = trtc.For(['arr_in','arr_out','k'], "idx",
	'''
	arr_out[idx] = arr_in[idx]*k;
	''')

harr = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype='float32')
darr = trtc.device_vector_from_numpy(ctx, harr)

darr_out = trtc.device_vector(ctx, 'float', 5)
forLoop.launch(ctx, 0, 5, [darr, darr_out, trtc.DVFloat(10.0)])

以上可以看作 ThrustRTC 的“引擎”部分。有了这个引擎，就可以着手实现 Thrust 中的各种算法了。尽管目前只实现了少量的简单算法，但已足以证明这种思路是可行的。比如下面是调用 Replace_If 算法的例子：

import ThrustRTC as trtc

trtc.set_ptx_cache('__ptx_cache__')
ctx = trtc.Context()

is_less_than_zero =  trtc.Functor( {}, ['x'], 'ret',
'''
         ret = x<0;
''')

darr = trtc.device_vector_from_list(ctx, [1, -2, 3, -4, 5 ], 'int32_t')
trtc.Replace_If(ctx, darr, is_less_than_zero , trtc.DVInt32(0))

这里用到了 Functor 功能，这也是 Thrust 中常用的技术。这里 is_less_than_zero 这个Functor 用来作为 Replace_If 判断是否替换某个元素的条件。

从这里给出的例子我们不难看出 ThrustRTC 和 Thrust C++ 之间存在的区别：

ThrustRTC 没有采用 Iterator 的设计，所有的算法就直接在 DVVector 上执行。这主要是出于简化设计的目的，同时作者认为 Iterator 在其他语言中并不具备其在 C++ 语言中的那种优越性。
Functor 不支持用 return 返回，需要单独定义一个返回变量并通过对它赋值来实现返回。在Verbose 模式下，我们可以看到 Functor 实际上是被转换成了一个 do{...} while(false); 形式的程序块：

// CUDA Kernel 的完整代码
#define DEVICE_ONLY
#include "DVVector.h"
#include "cstdint"
extern "C" __global__
void saxpy(VectorView<int32_t> _view_vec, int32_t _new_value, size_t _begin, size_t _end)
{
      size_t _idx = threadIdx.x + blockIdx.x*blockDim.x + _begin;
      if (_idx>=_end) return;
      bool ret;
      do{
          auto x = _view_vec[_idx];
          ret = x<0;
      } while(false);
      if (ret) _view_vec[_idx] = _new_value;
}

因此，虽然return不支持，但是break是支持的。

目前只移植了 Thrust 中少量的算法，也许最终也不能移植全部的算法，但是我会努力的。

关于这个项目的版权许可

值得一提的是，这个项目我决定采用目前大热的 Anti 996 许可证以示对这个运动的全力支持。为了和加班文化这个社会毒瘤进行斗争，我们程序员自己创造了 Anti 996 许可证。要让这一举措真正发挥作用，Anti 996 这个阵营中就必须诞生出伟大的项目。对于这件事情，个人认为我们这些为955良心公司工作的程序员责无旁贷，因为我们有更多的时间投身于自己感兴趣的项目。让我们把炉火烧得通红吧！