手把手教你用 Taichi 做高性能并行运算（0）

作者：幼儿园理化笙
链接：https://zhuanlan.zhihu.com/p/145222094
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
 

最近 GAMES 201 开课了，许多同学对其中 Taichi 编程语言的使用产生了很多疑问。我作为 Taichi 开发者之一，就试着水了一下这篇教程。

性能

众所周知，圆周率是一个常数，约为3.1415926535897，可以用下面这个公式计算：

（对，我故意挑了一个收敛很慢的级数，方便下面进行比较）

写一个Python程序来计算：

import time

def calc_pi():
  sum = 0.0
  for i in range(66666666):
    n = 2 * i + 1
    sum += pow(-1.0, i) / n
  return sum * 4


t0 = time.time()
print('PI =', calc_pi())
t1 = time.time()
print(f'{t1 - t0:.3} sec')

运行结果：

PI = 3.1415926385900446
14.5 sec

 

同样的代码，再试试C语言：

#include <stdio.h>
#include <math.h>
#include <time.h>

double calc_pi(void)
{
  int i, n;
  double sum = 0.0;
  for (i = 0; i < 66666666; i++) {
    n = 2 * i + 1;
    sum += pow(-1.0, i) / n;
  }
  return sum * 4;
}

int main(void)
{
  clock_t t0, t1;
  double pi;
  t0 = clock();
  printf("PI = %.15lf\n", calc_pi());
  t1 = clock();
  printf("%.3f sec\n", (double)(t1 - t0) / CLOCKS_PER_SEC);
  return 0;
}

运行结果：

PI = 3.141592638590045
1.348 sec

可见，解释型的Python语言，性能上比编译型的C语言差了许多，这是很自然的。

生产力

可是很多时候，同学们更喜欢Python，因为它的拥有面向对象，语法更简洁，还有很多现成的工具包可以使用（比如tensorflow，numpy），大大提升了我们的开发效率。

那么有没有办法，能够将写出的Python代码，高性能地运行呢？

答案是有的，Taichi 就是这样一个包，它会帮你把你的自定义函数编译成机器指令码，在CPU或GPU上并行执行，从而既保证了性能，又保证了生产力。

Taichi 可以通过 pip 安装（需要Python 3.6/3.7/3.8 64位)：

pip install taichi

还是圆周率的例子，要用 Taichi 编译一个函数，以便高性能执行，只需在前面加一个@ti.kernel的装饰器：

import time
import taichi as ti

@ti.kernel
def calc_pi() -> ti.f32:
  sum = 0.0
  for i in range(66666666):
    n = 2 * i + 1
    sum += pow(-1.0, i) / n
  return sum * 4


t0 = time.time()
print('PI =', calc_pi())
t1 = time.time()
print(f'{t1 - t0:.3} sec')

运行结果：

[Taichi] mode=release
[Taichi] version 0.6.6, supported archs: [cpu, cuda, opengl], commit 7d76c01c, python 3.8.2
PI = 3.141596794128418
3.62 sec

可见虽然性能不如C语言，还是有很大进步的，毕竟需要动态编译。

Taichi 会默认以 CPU 模式运行，如果指定为 GPU 的话，结果会更快：

import time
import taichi as ti

ti.init(arch=ti.gpu)  # 添加了这行

@ti.kernel
def calc_pi() -> ti.f32:
  sum = 0.0
  for i in range(66666666):
    n = 2 * i + 1
    sum += pow(-1.0, i) / n
  return sum * 4


t0 = time.time()
print('PI =', calc_pi())
t1 = time.time()
print(f'{t1 - t0:.3} sec')

运行结果：

[Taichi] mode=release
[Taichi] version 0.6.6, supported archs: [cpu, cuda, opengl], commit 7d76c01c, python 3.8.2
PI = 3.1415982246398926
1.56 sec

从以上的例子可以看到，通过 Taichi 的使用，我们可以在 Python 代码里写出接近编译型语言的性能。大大提升了 Python 作为胶水语言的运算能力，特别是大规模并行时的运算能力。

代码分析

@ti.kernel

一切以 @ti.kernel 修饰的函数都会被 Taichi 编译。

def calc_pi() -> ti.f32:

Taichi 是静态类型的，因此需要指定函数的返回类型。

for i in range(66666666):

编译器会自动把 for 循环分配到多个线程，并行执行，线程数量取决于你CPU的核数。而用户完全不必担心 for 循环的展开，编译器会处理一切优化问题。

sum += pow(-1.0, i) / n

Taichi 可以保证+=是原子操作，不会出现多线程数据竞争的问题。

sum = 0.0

这里如果是sum = 0的话，就不行，因为 Taichi 是静态类型的，如果 sum 初始为一个 int，则无法保存接下来数据类型是 float 的 pow(-1.0, i) / n 。编译器将会弹出警告来提醒你这一错误。

可见，虽然代码被传递给了 Taichi 编译器来编译，但大体上依然遵循 Python 的语法。

在下一篇文章中，我将继续介绍 Taichi 的更多特性和应用。

鸣谢

最后，十分感谢

@胡渊鸣

同学和他的开发团队为我们带来了 Taichi 这门高性能并行语言！