手把手教你用 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 這門高性能並行語言！