用 Taichi 加速 Python：提速 100+ 倍！

Python 已經(jīng)成為世界上最流行的編程語言，尤其在深度學(xué)習(xí)、數(shù)據(jù)科學(xué)等領(lǐng)域占據(jù)主導(dǎo)地位。但是由于其解釋執(zhí)行的屬性，Python 較低的性能很影響它在計(jì)算密集（比如多重 for 循環(huán)）的場景下發(fā)揮作用，實(shí)在讓人又愛又恨。如果你是一名經(jīng)常需要使用 Python 進(jìn)行密集計(jì)算的開發(fā)者，我相信你肯定會有下面的類似經(jīng)歷：

• 我的 Python 程序里面有個很大的 for 循環(huán)，循環(huán)體里面全是密集的計(jì)算，跑起來好慢啊...
• 我的程序里面只有一小部分計(jì)算是性能瓶頸，雖然可以用 C++ 改寫然后用 ctypes 綁定一下，但是那樣會很麻煩，還會有在別的機(jī)器上編譯不了的風(fēng)險。我希望所有的工作都能在一個 Python 腳本中完成！
• 我之前是忠實(shí)的 C++/Fortran 用戶，但是最近周圍的同學(xué)用 Python 的越來越多，我也想試試 Python，但是無奈很多祖?zhèn)鞔a用 Python 改寫以后就會慢 100 多倍，我接受不了...
• 我的工作中需要處理大量圖片數(shù)據(jù)，而需要的圖像處理功能 OpenCV 又不提供，只能自己手寫兩重 for 循環(huán)，在 Python 里面這么搞真是太痛苦了 ...

如果你有類似的煩惱，那真的值得了解一下 Taichi。我來簡單介紹一下：Taichi 是一個嵌入在 Python 中的領(lǐng)域特定語言，其一大功能就是加速 Python。Taichi 通過自己的編譯器將被 @ti.kernel 修飾的函數(shù)編譯到各種硬件上，包括 CPU 和 GPU，然后高性能執(zhí)行。

由于 Taichi 開發(fā)者社區(qū)花了大量的精力優(yōu)化 Taichi 在 Python 中的使用體驗(yàn)，所有的 Taichi 功能都可以在 import taichi as ti 以后使用，Taichi 本身也可以使用 pip 進(jìn)行安裝。當(dāng)然，Taichi 也可以與常用的 Python 包（numpy、matplotlib、PyTorch 等）進(jìn)行交互。

在這篇文章中，我們將通過三個計(jì)算例子來演示如何使用 Taichi 讓你的 Python 輕松加速 > 50 倍。這三個例子是：1. 計(jì)算質(zhì)數(shù)數(shù)目；2. 動態(tài)規(guī)劃求解最長公共子序列；3. 求解反應(yīng)-擴(kuò)散方程。

?? https://github.com/taichi-dev/faster-python-with-taichi

計(jì)算素?cái)?shù)個數(shù)

作為開胃小菜，我們先做一個小實(shí)驗(yàn)：計(jì)算小于給定正整數(shù) 的素?cái)?shù)的個數(shù)。相信任何對 Python 有基礎(chǔ)了解的人都不難寫出類似下面這樣的解法：

"""Count the number of primes in range [1, n].
"""

def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1

    return count

print(count_primes(1000000))

這個方法的思路簡單且粗暴：我們用一個函數(shù) is_prime 來判斷某個正整數(shù) 是不是素?cái)?shù)，是素?cái)?shù)則返回 1，不是則返回 0。這只要遍歷檢查從 2 到 之間是否有整數(shù)能夠整除 即可。然后將小于 的全部整數(shù)依次代入此函數(shù)并統(tǒng)計(jì)結(jié)果。將上面的代碼保存為 count_primes.py，在命令行運(yùn)行：

time python count_primes.py

在我的電腦上輸出的運(yùn)行結(jié)果是：

78498

real        0m2.235s
user        0m2.235s
sys        0m0.000s

耗時 2.235 秒。也許代碼中 設(shè)置成一百萬對你的電腦來說太輕松了，要不要把 改成一千萬試試？我打賭不管你的電腦多么高端，你起碼都要等個半分鐘才能看到結(jié)果。

好了下面是魔法時刻：我們不修改上面的函數(shù)體，只 import 一個“庫”，然后給兩個函數(shù)分別加一個裝飾器：

"""Count the number of primes below a given bound.
"""
import taichi as ti
ti.init()

@ti.func
def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1

    return count

print(count_primes(1000000))

仍然運(yùn)行 time python count_primes.py 命令，輸出的結(jié)果是：

78498

real        0m0.363s
user        0m0.546s
sys        0m0.179s

速度直接 x6! 而將 改成一千萬的話，Taichi 的耗時只會增加到 0.8s 左右，而 Python 則需要大約 55 秒，Taichi 直接加速了 70 倍！不僅如此，我們還可以在 ti.init 中加上 ti.init(arch=ti.gpu) 參數(shù)，指定 Taichi 使用 GPU 來進(jìn)行計(jì)算。在 GPU 上同樣的計(jì)算 Taichi 只花了不到 0.45 秒，比 Python 足足快了 120 倍！你可以運(yùn)行這里的代碼親身體會一下。

上面這個計(jì)算素?cái)?shù)的例子使用的方法有點(diǎn)土，作為習(xí)題還可以，但在實(shí)際生產(chǎn)中就顯得不那么實(shí)用了。我們接下來看一個實(shí)際中普遍使用的算法。

動態(tài)規(guī)劃

動態(tài)規(guī)劃（Dynamic Programming）是一類特別實(shí)用的算法，這類算法的哲學(xué)是以空間換時間，通過存儲中間計(jì)算結(jié)果來減少重復(fù)計(jì)算量。我們這里選擇一個求解最長公共子序列（Longest common subsequence, LCS）的例子 (算法導(dǎo)論的讀者有木有)。

插播兩個來自淵鳴的《算法導(dǎo)論》小故事：

1. 筆者小時候買過一本《算法導(dǎo)論》，書中提到四位作者都來自“麻雀理工學(xué)院”。當(dāng)時還很好奇：怎么會有學(xué)校叫這么奇怪的名字... 過了一陣才意識到自己可能成了盜版書籍的受害者。
2. 10 年后，我還真的來到了麻省理工學(xué)院（MIT）讀博士，一年后進(jìn)行碩士論文答辯（MIT 叫做 Research Qualification Exam），我自然就帶著 Taichi 的論文去了。答辯委員會里面有一位慈祥的教授，Charles E. Leiserson，嗯，就是《算法導(dǎo)論》的作者之一，“CLRS” 之中的 L。

言歸正傳。所謂子序列，就是一個序列的子集，但是保持它們在原序列中的順序。比如說 [1, 2, 1] 是 [1, 2, 3, 1] 的子序列，而 [3, 2] 則不是。我們這里考慮對兩條給定的序列，求出它們最長公共子序列的長度。最長公共子序列就是兩個序列的所有公共子序列中最長的一條 (這個最長子序列未必唯一，但它的長度是唯一確定的)。

舉個例子：

a = [0, 1, 0, 2, 4, 3, 1, 2, 1]

和

b = [4, 0, 1, 4, 5, 3, 1, 2]

的最長公共子序列是

LCS(a, b) = [0, 1, 4, 3, 1, 2]

最長公共子序列有很多應(yīng)用。比如大家日常使用的 Linux diff 命令和 git 工具（比較兩個文件之間的相似度），還有生物信息學(xué)中判斷兩段基因的相似度（把數(shù)字換成 ACGT 就行），其中的實(shí)現(xiàn)都用到了 LCS。

動態(tài)規(guī)劃計(jì)算 LCS 的想法是我們依次求解序列 a 的前 i 個元素和序列 b 的前 j 個元素的最長公共子序列的長度，通過讓 i 和 j 逐漸增加我們就逐步得出了最終的結(jié)果。我們用 f[i, j] 表示 LCS((prefix(a, i), prefix(b, j)，其中 prefix(a, i) 表示序列 a 的前 i 個元素，即 a[0], a[1], ..., a[i - 1]。這樣我們就得到遞推式：

f[i, j] = max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
              max(f[i - 1, j], f[i, j - 1]))

于是，一個 LCS 算法用 Python 可以很自然地書寫為：

for i in range(1, len_a + 1):
    for j in range(1, len_b + 1):
        f[i, j] = max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
                      max(f[i - 1, j], f[i, j - 1]))

這里我們給出一個 Taichi 的加速實(shí)現(xiàn)：

import taichi as ti
import numpy as np

ti.init(arch=ti.cpu)

benchmark = True

N = 15000

f = ti.field(dtype=ti.i32, shape=(N + 1, N + 1))

if benchmark:
    a_numpy = np.random.randint(0, 100, N, dtype=np.int32)
    b_numpy = np.random.randint(0, 100, N, dtype=np.int32)
else:
    a_numpy = np.array([0, 1, 0, 2, 4, 3, 1, 2, 1], dtype=np.int32)
    b_numpy = np.array([4, 0, 1, 4, 5, 3, 1, 2], dtype=np.int32)


@ti.kernel
def compute_lcs(a: ti.types.ndarray(), b: ti.types.ndarray()) -> ti.i32:
    len_a, len_b = a.shape[0], b.shape[0]

    ti.loop_config(serialize=True) # 避免 Taichi 自動并行
    for i in range(1, len_a + 1):
        for j in range(1, len_b + 1):
            f[i, j] = max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
                          max(f[i - 1, j], f[i, j - 1]))

    return f[len_a, len_b]


print(compute_lcs(a_numpy, b_numpy))

將上面的代碼保存為 lcs.py，然后在終端運(yùn)行：

time python lcs.py

得到的結(jié)果為（具體結(jié)果每次未必一致）：

2721

real        0m1.409s
user        0m1.112s
sys        0m0.549s

我們在代碼中同時提供了分別使用 Taichi 和 Numpy 計(jì)算的版本，在我的電腦上對兩個長度是 N=15000 的隨機(jī)序列進(jìn)行計(jì)算 Taichi 版本大約需要 0.9 秒，而 Python 則需要 476s，足足差了 500 多倍！大家可以運(yùn)行一下體會 Taichi 相對 Numpy 那種飛一樣的感覺。

當(dāng)然，Numpy 主要針對的場景是以數(shù)組為基本單位的運(yùn)算，遇到這種需要在數(shù)組內(nèi)更細(xì)粒度進(jìn)行計(jì)算的情況就比較無力了。而這正是 Taichi 能夠發(fā)揮作用的地方。

反應(yīng) - 擴(kuò)散方程

在大自然中我們常常會在動植物的表面見到一些有趣的圖案，比如斑馬身上的條紋，獵豹身上的斑點(diǎn)，河豚表面的花紋等等。

這些圖案看起來是不規(guī)則的，但是又有一定的規(guī)律，并不完全隨機(jī)。從進(jìn)化的觀點(diǎn)，這些圖案是生物在長期演進(jìn)和自然選擇中逐漸形成的，但到底是什么規(guī)則決定了它們的形狀一直是個有趣的問題。阿蘭 . 圖靈 (正是圖靈機(jī)的發(fā)明人) 是最早注意到這一現(xiàn)象并嘗試給出模型描述的人。他在論文 "The Chemical Basis of Morphogenesis" 中提出可以用兩種化學(xué)物質(zhì) U, V 之間的相互作用來模擬圖案的形成過程，其中物質(zhì) U 的角色類似被捕食者 (prey)，物質(zhì) V 的角色類似捕食者 (predator)。它們之間的作用服從如下規(guī)則：

1. 初始時空間中隨機(jī)地分布了一些 U, V。
2. 在每個時刻 1, 2, 3, ..., U, V 兩種物質(zhì)都向其鄰域擴(kuò)散。
3. 當(dāng) U, V 相遇時，一定比例的 U, V 會合并轉(zhuǎn)化為更多的 V (捕食者在捕食后數(shù)量會增加)
4. 為了避免捕食者 V 的數(shù)量過多導(dǎo)致 U 的數(shù)量被消耗光，我們在每個時刻按照一定的比例 f 添加 U，同時按照一定的比例 k 移走 V。

于是整個過程可以用下面的反應(yīng) - 擴(kuò)散方程描述：

這里關(guān)鍵的控制參數(shù)有四個，分別是 ， 分別控制 U, V 的擴(kuò)散速度， 代表 feed，控制 U 的添加量，而 代表 kill，控制移走 V 的比例。

為了在 Taichi 中模擬這一過程，我們將空間劃分為網(wǎng)格，每個網(wǎng)格中 U, V 的濃度值用一個 vec2 來表示。注意拉普拉斯算子 的數(shù)值計(jì)算是需要訪問當(dāng)前網(wǎng)格周圍的網(wǎng)格的，為了避免一邊修改一邊讀取這種操作的發(fā)生，我們需要開辟兩個形狀為 的網(wǎng)格，每次用其中一個網(wǎng)格的值作為舊值，將更新后的濃度值寫入另一個網(wǎng)格中，然后交換兩個網(wǎng)格的角色。所以我們需要的數(shù)據(jù)結(jié)構(gòu)應(yīng)該是：

W, H = 800, 600
uv = ti.Vector.field(2, float, shape=(2, W, H)) 

初始時，我們假定網(wǎng)格中 U 的濃度處處是 1，然后隨機(jī)選擇 50 個點(diǎn)撒上 V：

import numpy as np

uv_grid = np.zeros((2, W, H, 2), dtype=np.float32)
uv_grid[0, :, :, 0] = 1.0
rand_rows = np.random.choice(range(W), 50)
rand_cols = np.random.choice(range(H), 50)
uv_grid[0, rand_rows, rand_cols, 1] = 1.0
uv.from_numpy(uv_grid)

實(shí)際的計(jì)算代碼非常之簡短：

@ti.kernel
def compute(phase: int):
    for i, j in ti.ndrange(W, H):
        cen = uv[phase, i, j]
        lapl = uv[phase, i + 1, j] + uv[phase, i, j + 1] + uv[phase, i - 1, j] + uv[phase, i, j - 1] - 4.0 * cen
        du = Du * lapl[0] - cen[0] * cen[1] * cen[1] + feed * (1 - cen[0])
        dv = Dv * lapl[1] + cen[0] * cen[1] * cen[1] - (feed + kill) * cen[1]
        val = cen + 0.5 * tm.vec2(du, dv)
        uv[1 - phase, i, j] = val

這里我們使用了取值為 0 或 1 的整數(shù) phase 來控制使用 uv 的哪一層來作為舊的網(wǎng)格，并將更新的值寫入 1-phase 對應(yīng)的層中。

根據(jù) V 的濃度進(jìn)行染色，我們得到了如下的動畫效果:

非常有趣的是，雖然 V 的初始濃度是隨機(jī)設(shè)置的，但是最終得到的圖案卻具有相似性。

我們在代碼中提供了基于 Taichi 和 Numba 的兩份不同的實(shí)現(xiàn)，Taichi 的版本由于使用了 GPU 進(jìn)行計(jì)算，計(jì)算的部分可以輕松達(dá)到 300+ fps，而 Numba 的版本計(jì)算部分雖然也是編譯執(zhí)行的，但由于是在 CPU 上計(jì)算的，只有大約 30fps 左右。大家可以親自運(yùn)行代碼體會一下 Taichi 使用 GPU 加速的巨大優(yōu)勢。

總結(jié)

在這三個例子上 Taichi 都讓程序有了大幅加速。主要的性能來自三點(diǎn)：

1. Taichi 是編譯性的，而 Python 是解釋性的
2. Taichi 能自動并行，而 Python 通常是單線程的
3. Taichi 能在 GPU 上運(yùn)行，而 Python 本身是在 CPU 上運(yùn)行的

當(dāng)然，加速 Python 還有很多其他工具，這里我們分析一下他們和 Taichi 的優(yōu)劣。

與 Numpy/JAX/PyTorch/TensorFlow 比較：這幾類工具都高度基于數(shù)組運(yùn)算。計(jì)算的最小單位是數(shù)組，在 Data Science、Deep Learning 等領(lǐng)域是有明顯的優(yōu)勢的。但是在科學(xué)計(jì)算領(lǐng)域，這樣做導(dǎo)致靈活性缺失：比如說前面那個計(jì)算質(zhì)數(shù)的程序，就比較難使用數(shù)組運(yùn)算表示出來。Taichi 的優(yōu)勢就在于其靈活性，能夠直接操縱循環(huán)的每一次迭代，以一種更細(xì)顆粒度進(jìn)行對于計(jì)算的描述，類似 C++ 和 CUDA。

與 Cython 比較：使用 Cython 編寫程序?qū)崿F(xiàn)加速也是一種常見的選擇。在 Numpy 和 Scipy 的官方代碼中有不少模塊都是使用 Cython 編寫然后編譯的。但按照 Cython 的要求書寫代碼會比較麻煩，會犧牲一些可讀性。Cython 支持一定程度的并行計(jì)算，但不支持直接調(diào)用 GPU 進(jìn)行計(jì)算。

與 Numba 比較：Numba 顧名思義，是非常適合針對 Numpy 進(jìn)行加速的方案。當(dāng)你的函數(shù)是針對 Numpy 的數(shù)組向量化的操作時，使用 Numba 將其編譯以后執(zhí)行可以大大加速。Taichi 相比 Numba 的優(yōu)勢還有：1. Taichi 支持各種靈活的數(shù)據(jù)類型，比如 struct, dataclass, quant, sparse 等等，你可以任意指定它們的內(nèi)存排布，當(dāng)數(shù)據(jù)量龐大時這個優(yōu)勢會非常明顯。而 Numba 只有在針對 Numpy 的稠密數(shù)組時效果最佳。2. Taichi 可以調(diào)用不同的 GPU 后端進(jìn)行計(jì)算，所以寫大規(guī)模并行程序（如粒子仿真、渲染器等）這種操作對 Taichi 來說是小菜一碟。但你很難想象可以用 Numba 寫一個還過得去的 (哪怕離線) 渲染器。

與 Pypy 比較：Pypy 是一個 Python 的 JIT 編譯器，這個工具 2007 年就有了，和 Taichi 的解決方案有些類似，都是通過編譯的方式加速 Python。Pypy 最大優(yōu)勢在于 Python 代碼完全不用改變，就能通過 Pypy 加速。但是這也是 Pypy 加速比率比 Taichi 低的原因：因?yàn)?Pypy 需要在編譯的同時保持 Python 所有的語言特性，所以能夠進(jìn)行的優(yōu)化比較有限。而 Taichi 有一套自己的語法，雖然和 Python 很像但是也有自己的一些假設(shè)，這使得 Taichi 能夠?qū)崿F(xiàn)更大的加速。

與 ctypes 比較：ctypes 可以讓用戶在 Python 中調(diào)用 C 函數(shù)。C++、CUDA 編寫的程序也可以用過 C 接口暴露給 Python 使用。但是，ctypes 會讓工具鏈復(fù)雜化：為了寫一段比較快的程序，用戶需要同時掌握 C、Python、CMake、CUDA 等等語言，和本文描述的完全在 Python 中解決問題的方案比起來還是麻煩了一些。

總而言之，在科學(xué)計(jì)算任務(wù)上，Taichi 還是有自己獨(dú)特的優(yōu)勢的，大家可以根據(jù)自己的需求選擇最合適的工具。如果你需要在 Python 中實(shí)現(xiàn)類似 C/C++ 語言的性能，Taichi 不失為一個理想的選擇！

最后，我們希望 Taichi 能夠?yàn)槟銕韮r值，也希望能夠聽到你對 Taichi 的反饋，歡迎給我們提交 issues。如果想一鍵體驗(yàn) Taichi，只需要執(zhí)行????

pip install -U taichi

并執(zhí)行????

ti gallery

就可以體驗(yàn)各種基于 Taichi 的高性能可視化 Demos，期待與大家相遇！