From CUDA to Taichi: 如安在图形项目中使用GPU计算

sec9281

我从去年开始接触CUDA编程，写了一年之后感觉项目迭代周期有点慢，本年终于用上了taichi。此次因为要讲组会趁便总结一下，如何从CUDA编程的角度理解GPU计算，以及taichi与CUDA对比提供了哪些便当。
CUDA Minimal Example

1. // compile with: nvcc add.cu -o add_cuda
2. #include <iostream>
4. __global__ void addKernel(int *c, const int *a, const int *b, int n)
5. {
6.     int i = blockIdx.x * blockDim.x + threadIdx.x;
7.     if(i >= n) return;
8.     c[i] = a[i] + b[i];
9. }
11. constexpr int N = 100;
13. int main() {
14.   // host memory
15.   int* a = new int[N];
16.   int* b = new int[N];
17.   int* c = new int[N];
18.   for(int i = 0; i < N; ++i) {
19.     a[i] = i;
20.     b[i] = 2 * i;
21.   }
22.   // device memory
23.   int* da;
24.   int* db;
25.   int* dc;
26.   cudaMalloc(&da, N * sizeof(int));
27.   cudaMemcpy(da, a, N * sizeof(int), cudaMemcpyHostToDevice);
28.   cudaMalloc(&db, N * sizeof(int));
29.   cudaMemcpy(db, b, N * sizeof(int), cudaMemcpyHostToDevice);
30.   cudaMalloc(&dc, N * sizeof(int));
32.   int n_thread = 32;
33.   int n_block = (N + n_thread - 1) / n_thread;
34.   addKernel<<<n_block, n_thread>>>(dc, da, db, N);
35.   cudaMemcpy(c, dc, N * sizeof(int), cudaMemcpyDeviceToHost);
36.   for(int i = 0; i < N; ++i) {
37.     std::cout << c[i] << ” ”;
38.   }
39.   std::cout << std::endl;
41.   delete[] a;
42.   delete[] b;
43.   delete[] c;
44.   cudaFree(da);
45.   cudaFree(db);
46.   cudaFree(dc);
47. }

复制代码

这是一份使用cuda进行并行加法的最基础的例子，实现了两个数组的并行加法。接下来就从这份代码出发，总结一下我感觉CUDA编程斗劲重要的三个方面：编程模型和调剂机制，内存模型以及同步机制。

CUDA的编程模型和调剂机制

CUDA与一般C++法式最明显的区别，就是用__global__标识符定义的kernel函数。__global__标识表记标帜了这个函数是一个CUDA的kernel，是跑在GPU上的代码，也就决定了这个函数的调用方式是这样的：addKernel<<<n_block, n_thread>>>() 。此中的两个参数分袂暗示使用多少个block，每个block里有多少thread来跑这个kernel。深入理解这两个参数需要了解CUDA的编程模型和调剂机制。
在抽象层面，每一个thread是执行任务的最小单元，会按挨次执行kernel中的代码。一些数量的thread会组织成一个block，数量由调用参数中的n_thread决定。 所有的block会组织成一个grid，数量由调用参数中的n_block决定。


基于这样的模型，每个thread在具体执行kernel的时候会同时被分配两个id：blockIdx和 threadIdx，来标识当前thread对应的是grid中的哪一个。 事实上，CUDA默认使用dim3类型的向量来定义n_thread和n_block，也就是grid内block的布局和block内部thread的布局是一个三维空间的网格。这主要是为了法式员理解的便当， 比如当我们写三维流体模拟的时候就可以用n_thread和n_block的三维布局对应真实的流体网格。相应的，blockIdx和 threadIdx也都默认是三维向量，所以在上面例子中我们只取此中的x分量。
为什么CUDA要设计出一个中间的层级block而不是一个grid下都是thread呢？这主要是从硬件实现上的考虑。


GPU上具体执行计算的单元是CUDA Core，有独立的寄存器。逻辑层面的一个thread最终在硬件上在某一个CUDA Core上执行。而一些数量的CUDA Core会组织成一个Streaming Multiprocessor (SM)。一个SM中的CUDA Core可以共享一个L1缓存，而从局部性角度考虑一个SM中的CUDA Core访谒的内存应该是接近的，所以这样设计能够大大加快访存速度。一块GPU的通用计算部门可以简单当作由很多的SM和global memory组成。
CUDA在调剂时候会把一个逻辑的block分配给独一的一个SM，不会呈现一个block跑在多个SM上的情况。但反过来有可能一个SM被分配了多个block，这取决于一个block需要的资源和SM能提供的资源。在执行运算的时候，SM采用的架构称为Single-Instruction, Multiple-Thread (SIMT)，每个时刻一个SM上的所有CUDA Core执行的都是不异的命令，只不外因为每个CUDA Core上跑的thread的threadIdx和blockIdx分歧， 所以运行成果分歧。当呈现if分支的时候，跳出if的thread也只能等待还在分支里的thread跑完之后再一起执行下面的命令。 当几个block被分配给了一个SM时，GPU会将block划分为更细粒度的组进行调剂，称为warp。在此刻所有的硬件上，一个warp都是32个线程。我们可以列举几种SM会面对的情况：

• SM上有一个block，这个block有64个线程，那么GPU就会将其划分为2个warp
  
• SM上有一个block，这个block有30个线程，那么GPU就会补上2个空线程变成一个warp
  
• SM上有两个block，一个block有12个线程，一个线程有20个线程，那么GPU就会将其补为两个warp
  

使用cudaGetDeviceProperties()可以得到GPU硬件的相关参数。SM存在最大的warp数量，我的机器上是64个warp，也就是2048个线程。同时一个block的线程数也有上限，我的机器上是1024个线程。也就是说，如果我创建的block有1024个线程，那么GPU可以将两个block分配一个SM，一共64个warp。
尽管如此，warp中的线程在物理上可能也不是同时执行。因为早期的显卡可能存在SM中的CUDA Core少于32个的情况。这时GPU就会把warp分时执行，但是我们其实不用担忧这一点，这个分时操作对法式员是完全不成见的。GPU可能为了填补流水线空泡对来自分歧block的warp一起进行调剂，来使得执行效率最高。
以上就是CUDA的编程模型和调剂模型。回到代码中来，我们在选择n_thread的时候应该尽量选择32的倍数，但是不能超过硬件上限，一般情况下选择32、64、128都可以，性能分歧不是很大。

除了__global__关键字，还有__device__和__host__，分袂暗示：
__global__：CUDA kernel，只能通过<<<,>>> 的方式调用，返回值必需为void
__device__：CUDA kernel中调用的函数，不能从CPU调用，可以有返回值
__host__：一般CPU上的函数，可以不写
GPU内存模型

GPU上有一个大的global memory，也就是GPU上的显存，我的机器上是4GB。然后到SM中有各级缓存。每个block可以使用缓存中的一部门，称为shared memory，我的机器上是48KB。shared memory是每个block独立的，可以被block内的线程共用，对比global memory仿写速度很快，充实操作能大大提高计算效率。因为CUDA的内存跟CPU内存是分隔的，因此需要使用CUDA提供的函数进行操作。而CUDA暗示显存的方式也是指针，因此要注意不能跟CPU上的内存混淆，不能把CPU指针传给kernel。实例中就展示了一种常见的使用情况，在GPU上新建空间，然后把CPU上的数据拷贝到GPU上，得到成果之后拷贝回CPU内存，最后释放GPU空间：

1. cudaMalloc(&da, N * sizeof(int));
2. cudaMemcpy(da, a, N * sizeof(int), cudaMemcpyHostToDevice);
3. cudaMalloc(&db, N * sizeof(int));
4. cudaMemcpy(db, b, N * sizeof(int), cudaMemcpyHostToDevice);
5. cudaMalloc(&dc, N * sizeof(int));
6. addKernel<<<n_block, n_thread>>>(dc, da, db, N);
7. cudaMemcpy(c, dc, N * sizeof(int), cudaMemcpyDeviceToHost);
8. cudaFree(da);
9. cudaFree(db);
10. cudaFree(dc);

复制代码

这里使用cudaMalloc()创建的空间在global memory上。那如何使用shared memory呢？
如果要使用的shared memory大小已知，可以直接在kernel中使用__shared__关键字暗示：

1. __global__ void someKernel() {
2. // ...
3. int i = threadIdx.x;
4. __shared__ volatile float tmp[16];
5. tmp[i] = r[i] * r[i];
6. // ...
7. }

复制代码

__shared__标识表记标帜了tmp是shared memory，volatile感化跟在CPU上一样，避免对tmp的操作被寄存器优化。
此外，CUDA还撑持shared memory的大小是动态的情况。在调用CUDA kernel的时候我们可以在n_block和n_thread后加第三个参数s_mem，暗示每个block使用的shared memory大小。然后在kernel里使用extern __shared__:

1. someKernel<<<n_thread, n_block, s_mem>>>(a, b);
2. // ...
3. __global__ void someKernel(float* a, float* b) {
4. extern __shared__ float tmp[];
5. // the maximum size of tmp is s_mem
6. // ...
7. }

复制代码

CUDA的同步问题

CUDA的同步问题其实包含两个方面：GPU与CPU之间的同步以及GPU内部的同步。
GPU与CPU的同步
在之前的样例中，GPU与CPU之间的同步是隐式发生的。GPU接受并执行CPU的命令是一个队列模型，CPU不竭向队尾提交任务，GPU不竭从队首取任务执行。默认情况下，所有kernel调用和cudaMemcpy这样的函数调用城市被push到默认队列stream 0里。kernel调用在push到队列之后会立刻返回，CPU可以继续执行剩下的代码。而从GPU到CPU的cudaMemcpy会阻塞CPU，直到拷贝完成之后才会返回。通过这样的方式，我们能够保证cudaMemcpy返回的时候我们可以拿到正确成果。
CUDA撑持向多个stream提交任务。这些stream可以是跟stream 0同步的，也就是可以保证stream 0里的先提交的任务也会在当前stream任务执行前完成；也可以跟stream 0是异步的，这样可以在stream 0被cudaMemcpy阻塞的时候进行其他的任务。 向其他stream提交任务首先需要创建stream，然后在调用kernel的时候在<<<>>>的最后一个参数里填上对应的stream编号n_stream。这里就不赘述具体代码。总的来说，完整的kernel调用需要四个参数：

1. someKernel<<<n_block, n_thread, s_mem, n_stream>>>(a, b);

复制代码

CUDA同样提供了异步的cudaMemcpy方式:cudaMemcpyAsync,参数里可以填执行的stream编号。

GPU内部同步

1. // part of conjugate gradiant of 16x16 matrix in one block
2. // ...
3. int i = threadIdx.x;
4. __shared__ volatile float tmp[16 + 8];
5. tmp[i] = r[i] * r[i];
6. tmp[i] += tmp[i + 8];
7. tmp[i] += tmp[i + 4];
8. tmp[i] += tmp[i + 2];
9. tmp[i] += tmp[i + 1];
10. if (i == 0) {
11. // tmp[0] is the inner product of r
12. }

复制代码

我们还是以代码为例。这个代码片段是用Conjugate Gradient (CG)方式并行解很多16x16矩阵方程的一部门。假设我们有很多的16x16的小矩阵要解Ax=b，于是我们可以将一个矩阵assign给一个block，每个block16个线程对应矩阵的16维，在每个block里跑CG。

按照之前的分析其实最好应该每个block 32个线程解两个矩阵，但为了说明便利就牺牲一下效率

这个代码片段是用来计算长度为16的向量r的内积，关键部门是使用reduction的方式计算r*r的求和，算法流程可以用下面的图暗示：


这里我们其实使用了warp的隐式同步机制。上面提到一个warp内的线程是同步执行的，因此tmp+=执行完之后tmp[0]里就是求和成果。 而当block内的线程数量超过一个warp时，线程之间就不能完全保证同步，因此每一步tmp+=前后都需要加上__syncthreads()来同步block内的线程。__syncthreads()可以保证一个block内的线程执行到同一位置之后再继续执行剩下的代码。
解决了block内部线程的同步之后，还有block之间的同步问题。最简单的情况下，我们可以使用atomicAdd来确保对同一块global memory的改削是原子的。除此之外还有一些更复杂的同步机制，需要按照具体情况进行选择，这里就不赘述了。
我感觉理解了CUDA的编程模型和调剂机制，内存模型以及同步机制这三个方面，就基本能开始使用CUDA写并行法式了。然而在一个图形项目中，往往有1w+行的C++代码，需要的也不止于两个数组并行相加这样的简单功能，所以想在图形项目中使用CUDA还需要注意一些此外事情。
文件组织

所有的CUDA kernel和device function都必需实此刻.cu文件中， kernel的调用也必需在.cu中，但是CUDA并没有特殊的头文件。NVCC需要将.cu文件里的kernel和device function编译成二进制文件，其他部门可以交由C++ compiler解决。 我的一般组织方式是：
a.h：声明__global___，__device__或者__host__的函数
a.cpp：实现a.h中定义的CPU代码
a.cu：实现a.h中定义的GPU代码，或者需要调用kernel的函数
一般C++项目会使用Visual Studio，CMake等东西进行配置。与这些传统东西对比，xmake 是一个更现代更好用的项目配置东西。如果使用的是xmake，那么加上CUDA的撑持非常便利：

1. set_xmakever(”2.5.9”)
3. add_requires(”cuda”, {configs={utils={”cusparse”, ”cusolver”, ”cublas”}}})
4. add_rules(”mode.release”, ”mode.debug”)
5. set_languages(”cxx17”)
7. target(”hello”)
8.     set_kind(”binary”)
9.     add_includedirs(”src/inc”)
10.     add_files(”src/main.cpp”, ”src/*.cpp”, ”src/*.cu”)
11.     add_cuflags(”-rdc=true”)
12.     add_cugencodes(”native”)
13.     set_targetdir(”bin”)

复制代码

使用Class

CUDA中可以使用class。代码中定义的struct/class复合数据类型可以直接在kernel/device function中作为参数类型使用。但是如果在class/struct中定义了成员函数，必需要加上__device__修饰符才能在kernel中调用，比如下面的例子：

1. class Point2D {
2.   public:
3.   __host__ __device__ Point2D(float xx=0, float yy=0):x(xx),y(yy){};
4.   float x, y;
5. };
7. __global__ void kernel(Point2D* arr) {
8.   Point2D a(1., 2.);
9.   // ...
10. }

复制代码

因为在图形中经常需要措置向量和几何计算，在CPU上Eigen有这样的数学库可以用，但GPU上没有一个主流的数学库。Eigen在文档中说明了撑持CUDA，但我测试的成果是<Eigen/Core>中定义的类型可以使用，但是引入其他部门比如JacobiSVD会呈现问题。在我的项目中我就只使用了<Eigen/Core>中的功能，进行简单的向量矩阵运算，然后本身实现斗劲复杂的功能，比如polar decomposition和SVD。
cuSparse, cuBlas, cuSolver

CUDA提供了很多开箱即用的数学库，比如cuBlas用来措置dense的向量矩阵运算，cuSparse用来措置稀疏矩阵，cuSolver用来求解大规模稀疏矩阵方程等等。这些库只能从host端调用，比起本身手写的算法往往效率更高。这些库对于单精度和双精度的浮点数定义了分歧的函数，比如cublasDaxpy暗示double浮点数的a*x+y计算， cublasSaxpy暗示single浮点数的a*x+y运算。为了防止措置高精度模型单精度浮点数不够用的情况，可以使用宏命令或者if constexpr切换分歧的函数：

1. #ifdef REAL_AS_DOUBLE
2.         cublasDaxpy(cublas_handle_, dim_[l], &one, P, 1, lhs, 1);
3. #else
4.         cublasSaxpy(cublas_handle_, dim_[l], &one, P, 1, lhs, 1);
5. #endif

复制代码

taichi

在学习CUDA之前我其实测验考试过taichi，但是当时我对GPU编程不甚了解，加上当时taichi功能不是很完善，所以我的体验是taichi并没有给我带来很多便利，反而因为各种基础的向量拷贝，稀疏矩阵求解都需要本身实现，而且文档不是很齐全，api不是很不变，让我感觉taichi只能做一些toy example，不太适合做项目。还有一部门原因是因为组里项目祖传代码都是C++的，而且很多任务（比如做remeshing，构造流体稀疏矩阵）不是太适合并行，就没有继续深入了解taichi。但是当我做了一段时间的弹性体模拟，开始测验考试到CUDA的酸爽之后，回到taichi的感觉非常好。1.3版本的taichi撑持GPU上的稀疏矩阵求解，有3D可视化东西，api斗劲不变，文档也规范了很多，而且可以操作到python原生的便当特性，因此对比CUDA体验有了很大的提升。
最开始的CUDA代码，核心部门换成taichi只需要写成下面这个样子：

1. a = ti.field(ti.i32, N)
2. b = ti.field(ti.i32, N)
3. c = ti.field(ti.i32, N)
4. @ti.kernel
5. def add():
6.   for i in c:
7.     c[i] = a[i] + b[i]

复制代码

ti.kernel

ti.kernel是类似于CUDA kernel的存在，taichi还有ti.func对应CUDA中的__device__。对比CUDA，taichi一个很大的区别在于，taichi并行的语法是自动并行最外层的for循环，不用像CUDA一样声明一次再调用一次。这个语法就更容易理解，而且简化了很大一部门代码。比如在CUDA中我们经常要写定名和参数很复杂的kernel：

1. __global__ void EnergyGradientWithCtrlReordered(
2.     const real* tet_grad, const int32_t* v2t_ids, const int32_t* v2t_off,
3.     const bool* fixed, const Vec3* fixed_X, const Vec3* X, const int32_t* m2h,
4.     real* grad, const real control_mag, int ctrl_vert, Vec3 ctrl_pos,
5.     const int32_t n_vert) {
6. // ...
7. }

复制代码

这是因为我们需要把kernel涉及到的所有参数列举出来，而且每一个kernel可能只干一个很小的事情，所以定名也斗劲麻烦。但是在taichi中，我们可以直接写：

1. @ti.kernel
2. def func():
3.   for i, j in A:
4.     # ...
5.   for i in B:
6.     # ...

复制代码

taichi可以找到对应的全局变量，而不需要我们本身写到参数里。同时我们也可以在一个ti.kernel里写很多个分歧的循环，最外层的循环城市自动并行，在循环之间也可以进行一些其他的计算。同时taichi的kernel还允许有最多一个返回值，而CUDA kernel只能返回void。
不外taichi的kernel有一个处所需要额外注意。taichi的kernel是采用的是just in-time compilation的模式，每个kernel会在第一次调用的时候被编译成后端代码，如果kernel中使用了全局变量，那么这个全局变量会直接取值放进被编译的代码里，之后改这个全局变量也不会改变kernel中的值。taichi官方给出了一个样例：

1. import taichi as ti
2. ti.init()
4. a = 1
6. @ti.kernel
7. def kernel_1():
8.     print(a)
10. @ti.kernel
11. def kernel_2():
12.     print(a)
14. kernel_1()  # Prints 1
15. a = 2
16. kernel_1()  # Prints 1
17. kernel_2()  # Prints 2

复制代码

但是如果是向kernel传递参数，或者引用ti.field定义的全局变量，kernel编译的时候会把这些值作为引用传入，成果就会纷歧样：

1. import taichi as ti
2. ti.init()
4. a = ti.field(ti.i32, shape=(1,))
5. a[0] = 1
7. @ti.kernel
8. def kernel_1():
9.     print(a[0])
11. @ti.kernel
12. def kernel_2():
13.     print(a[0])
15. kernel_1()  # Prints 1
16. a[0] = 2
17. kernel_1()  # Prints 2
18. kernel_2()  # Prints 2

复制代码

这是因为如果是传入的参数在每次调用的时候城市复制一次，而ti.field定义的数据会保留在GPU上，默认按引用访谒。
ti.field

taichi的field是一大块比CUDA更好用的功能。概况上ti.field就是一个高维的数组，撑持scalar、vector、matrix等基础数据布局：

1. a = ti.field(ti.f32, shape=(N, M)) # 2 dim scalar field
2. b = ti.Vector.field(3, ti.i32, shape=(N, )) # 1 dim field of vector3
3. c = ti.Matrix.field(2, 2, ti.i32, shape=(N, M)) # 2 dim field of 2x2 matrix

复制代码

但是taichi在ti.field背后提供了非常强大的自定义功能，比如自定义数据的储存挨次，AOS和SOA，稀疏数据布局等。刘天添老师在22年的SIGGRPH Asia上有一个非常好的教程，把taichi的数据布局讲得非常清楚：
举例说明，如果我想在C++中实现一个row major的二维数据布局，我需要手动计算下标或者直接定义成C++中的二维数组，这样当我想换成column major的时候就需要把所有用到这个数据布局的处所都改削一遍。但是taichi把数据布局的储存挨次 (layout) 跟数据布局本身解耦开了。数据布局的储存挨次通过SNode-tree实现：


上图中的x和y都是3x2的整数矩阵，但是x是row major，y是column major。ti.root暗示根节点，ti.root.dense(ti.i, 3)暗示root的第i维是以dense的方式存储了3个子节点，然后每个子节点通过.dense(ti.j, 2)暗示在第j维按dense的方式存储了2个叶子节点，而最后.place(x)则暗示每个叶子节点上是一个int32。通过这样的树形布局我们能够便利定义row/column major的储存挨次。而到实际访谒的时候，我们都可以使用for i, j in x这样的方式遍历数据布局，taichi就会自动按照定义的row/column major的挨次访谒：

1. @ti.kernel
2. def func():
3.   for i, j in x:
4.     # ...

复制代码

SNode的定义非常巧妙，比如下面展示的实现Z-order的存储，切换AOS/SOA，稀疏网格的功能都可以在一两行内实现：






taichi总结

taichi还有其他很好用的功能。比如taichi中的+=会默认视为原子操作，所以基本不用担忧racing的问题。meshtaichi提供了针对mesh数据布局的优化。quantaichi可以压缩数值精度减少内存占用。 稀疏矩阵库可以实现基本的稀疏系统求解。此外还有difftaichi实现自动微分，开箱即用的可视化东西等等。从我的使用体验上来看，因为python本身语法简洁、生态丰硕，使用taichi+python进行原型开发和迭代是比CUDA+C++快得多的选择。尽管taichi对比CUDA在计算效率上必定是有损掉，但是总体来说不会有数量级上的差距，大大都情况下基本可以忽略不记。