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

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

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

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

由于 Taichi 開(kāi)發(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)行交互。

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

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

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

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

"""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))

這個(gè)方法的思路簡(jiǎn)單且粗暴：我們用一個(gè)函數(shù) is_prime 來(lái)判斷某個(gè)正整數(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

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

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

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

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

動(dòng)態(tài)規(guī)劃

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

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

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

言歸正傳。所謂子序列，就是一個(gè)序列的子集，但是保持它們?cè)谠蛄兄械捻樞?。比如說(shuō) [1, 2, 1] 是 [1, 2, 3, 1] 的子序列，而 [3, 2] 則不是。我們這里考慮對(duì)兩條給定的序列，求出它們最長(zhǎng)公共子序列的長(zhǎng)度。最長(zhǎng)公共子序列就是兩個(gè)序列的所有公共子序列中最長(zhǎng)的一條 (這個(gè)最長(zhǎng)子序列未必唯一，但它的長(zhǎng)度是唯一確定的)。

舉個(gè)例子：

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

和

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

的最長(zhǎng)公共子序列是

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

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

動(dòng)態(tài)規(guī)劃計(jì)算 LCS 的想法是我們依次求解序列 a 的前 i 個(gè)元素和序列 b 的前 j 個(gè)元素的最長(zhǎng)公共子序列的長(zhǎng)度，通過(guò)讓 i 和 j 逐漸增加我們就逐步得出了最終的結(jié)果。我們用 f[i, j] 表示 LCS((prefix(a, i), prefix(b, j)，其中 prefix(a, i) 表示序列 a 的前 i 個(gè)元素，即 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]))

于是，一個(gè) LCS 算法用 Python 可以很自然地書(shū)寫(xiě)為：

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]))

這里我們給出一個(gè) 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 自動(dòng)并行
    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

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

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

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

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

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

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

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

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

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

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

初始時(shí)，我們假定網(wǎng)格中 U 的濃度處處是 1，然后隨機(jī)選擇 50 個(gè)點(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ì)算代碼非常之簡(jiǎn)短：

@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 來(lái)控制使用 uv 的哪一層來(lái)作為舊的網(wǎng)格，并將更新的值寫(xiě)入 1-phase 對(duì)應(yīng)的層中。

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

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

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

總結(jié)

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

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

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

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

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

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

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

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

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

最后，我們希望 Taichi 能夠?yàn)槟銕?lái)價(jià)值，也希望能夠聽(tīng)到你對(duì) Taichi 的反饋，歡迎給我們提交 issues，加入 Taichi 開(kāi)發(fā)者社區(qū)。如果想一鍵體驗(yàn) Taichi，只需要執(zhí)行????

pip install -U taichi

并執(zhí)行????

ti gallery

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