基于Taichi的Python高性能計算入門指南
譯文自從Python編程語言誕生以來，它的核心理念一直是最大限度地提高代碼的可讀性和簡單性。Python對可讀性和簡單性的追求簡直達到了如癡如狂的境地。一個事實即可證實這一點：只要你在Python系統的根目錄中輸入命令“import this”后按下回車鍵，竟然馬上打印出一首英文小詩，翻譯成中文大致意思是：
“美麗勝過丑陋，顯式優于隱式。
簡單比復雜好，復雜比繁雜好。
扁平優于嵌套，稀疏勝過密集。
可讀性很重要……”
簡單總比復雜好，可讀性很重要。毫無疑問，Python確實在實現這些目標方面非常成功：它是迄今為止最友好的學習語言，并且一個普通的Python程序通常比等效的C++代碼短5到10倍。不幸的是，這里有一個陷阱：Python的簡單性是以降低性能為代價的!事實上，Python程序比C++對應的速度慢10到100倍。因此，似乎在速度和簡單性之間存在著一種永久的權衡，任何編程語言都不可能同時擁有這兩者。
但是，別擔心，所有的希望都沒有失去。
Taichi可實現兩全其美
Taichi編程語言是對Python編程語言進行擴展的一種嘗試，其結構支持通用、高性能的計算。它支持無縫地嵌入到Python中，而同時可以發揮計算機中所有的計算能力——包括多核CPU功能以及更為重要的GPU性能。
我們在本文中將展示一個使用Taichi編寫的示例程序。該程序使用GPU對落在球體上的一塊布進行實時物理模擬，同時渲染結果。
編寫實時GPU物理模擬器絕非易事，但是實現本例程的Taichi源代碼卻異常簡單。本文的其余部分將引導你完成整個實現，這樣你就可以領略Taichi提供的功能，以及它們的強大和友好程度。
在我們開始之前，你不妨先猜猜這個程序大概由多少行代碼組成。當然，你會在文章末尾找到這一答案。
算法概述
我們的程序將把一塊布料建模為一個質量彈簧系統。更具體地說，我們將此布料表示為點質量的N×N網格，其中相鄰點由彈簧連接。下面的圖形是由斯坦福大學的Matthew Fisher提供的，正展示了這種結構。
該質量彈簧系統的運動受4個因素影響：
- 重力
- 彈簧的內力
- 阻尼
- 與夾在中間的紅球碰撞
為了簡單起見，我們忽略了布料的自碰撞。我們的程序在t=0時開始。然后，在模擬的每一步，它將時間提前一個小常數dt。該程序通過評估上述4個因素中的每一個因素的影響來估計系統在這一小段時間內會發生什么，并在時間步結束時更新每個質量點的位置和速度。然后，更新的質點位置用于更新屏幕上渲染的圖像。
程序開始
盡管Taichi本身就是一種編程語言，但它以Python包的形式存在，只需運行pip install Taichi即可安裝。
要在Python程序中使用Taichi，首先需要使用別名ti導入Taichi：
import taichi as ti
如果您的機器具有支持CUDA的Nvidia GPU，Taichi程序的性能將會得到最大程度發揮。如果是這種情況，請在上述導入語句后添加以下代碼行：
ti.init(arch=ti.cuda)
如果你沒有CUDA GPU，Taichi仍然可以通過其他圖形API(如ti.metal，ti.vulkan和ti.opengl)與你的GPU交互。然而，Taichi對這些API的支持不如其對CUDA的支持那么全面。因此，目前情況下，我們使用CPU作為計算后端：
ti.init(arch=ti.cpu)
別擔心，Taichi即使只在CPU上運行，也會運行得很快。初始化Taichi之后，我們可以開始聲明用于描述質量彈簧布料的數據結構。為此，我們添加以下代碼行：
N = 128
x = ti.Vector.field(3, float, (N, N))
v = ti.Vector.field(3, float, (N, N))
這三行將x和v聲明為大小為N×N的二維數組，其中數組的每個元素都是一個浮點數的三維向量。在Taichi中，數組被稱為“場”，這兩個場分別記錄點質量的位置和速度。請注意，如果您將Taichi初始化為在CUDA GPU上運行，這些場/數組將自動存儲在GPU內存中。除了布料，我們還需要定義中間的球：
ball_radius = 0.2
ball_center = ti.Vector.field(3, float, (1,))
在這里，球中心是一個大小為1的一維場，其單個分量是一個三維浮點向量。聲明了所需的場之后，讓我們用t=0處的相應數據初始化這些場。我們希望確保，對于同一行或同一列上的任何一對相鄰點，它們之間的距離等于cell_size=1.0/N。這是通過以下初始化例程來實現的：
def init_scene():
for i, j in ti.ndrange(N, N):
x[i, j] = ti.Vector([i * cell_size,
j * cell_size / ti.sqrt(2),
(N - j) * cell_size / ti.sqrt(2)])
ball_center[0] = ti.Vector([0.5, -0.5, 0.0])
此處，你無需擔心每個x[i，j]值背后的含義——它的選擇只是為了使布料以45度角落下，參考下圖。
模擬
在每個時間步中，我們的程序都會模擬影響布料運動的4個因素：重力、彈簧內力、阻尼和與紅球的碰撞。其中，重力是最容易處理的。
下面是實現這一點的代碼：
@ti.kernel
def step():
for i in ti.grouped(v):
v[i].y -= gravity * dt
這里有兩點需要注意。首先，語句for i in ti.grouped(x)意味著將循環迭代x的所有元素，而不管x中有多少維度。其次，也是最重要的是：注解@ti.kernel意味著Taichi將自動并行運行函數中的任何頂級for循環。在本例中，Taichi將并行更新v中每個N*N向量的y分量。
接下來，我們來處理弦線的內力計算問題。首先，請注意前面圖形中的每個質點最多連接到八個鄰接質點。這些連接在我們的程序中表示如下：
links = [[-1, 0], [1, 0], [0, -1], [0, 1], [-1, -1], [1, -1], [-1, 1], [1, 1]
links = [ti.Vector(v) for v in links]
從物理角度來看，系統中的每個彈簧s都用固定長度l(s,0)初始化。在任何時間t，如果s的當前長度l(s,t)超過l(s,0)，則彈簧將在其端點上施加力，將它們拉在一起。相反，如果l(s,t)小于l(s,0)，則彈簧會將端點彼此推開。這些力的大小始終與l(s,0)-l(s,0)的絕對值成正比。此交互由以下代碼段捕獲：
for i in ti.grouped(x):
force = ti.Vector([0.0,0.0,0.0])
for d in ti.static(links):
j = min(max(i + d, 0), [N-1,N-1])
relative_pos = x[j] - x[i]
current_length = relative_pos.norm()
original_length = cell_size * float(i-j).norm()
if original_length != 0:
force += stiffness * relative_pos.normalized() *
(current_length - original_length) /
original_length
v[i] += force * dt
請注意，這個for循環仍應作為substep函數中的頂級for循環，該函數用@ti.kernel注解。這樣可以確保并行計算施加到每個質點的彈簧力。stiffness在此是一個常數，用于控制彈簧長度變化的程度。在上述程序中，我們使用stiffness =1600指定它的值。在現實世界中，當彈簧振動時，彈簧中儲存的能量會消散到周圍環境中，其振動最終停止。為了捕捉這種效應，在每個時間步，我們稍微降低每個點的速度大小：
for i in ti.grouped(x):
v[i] *= ti.exp(-damping * dt)
在此，damping取固定值2。
我們還需要處理布料和紅球之間的碰撞。要做到這一點，我們只需將質點與球接觸時的速度降低到0。這樣可以確保布料“掛”在球上，而不是穿透球或向下滑動：
if (x[i]-ball_center[0]).norm() <= ball_radius:
v[i] = ti.Vector([0.0, 0.0, 0.0])
最后，我們用每個質點的速度更新其自身的位置：
x[i] += dt * v[i]
這就是我們對一塊質量彈簧布料進行并行模擬所需的全部代碼。
渲染
我們將使用Taichi內置的基于GPU的GUI系統(昵稱是“GGUI”)渲染布料。GGUI使用Vulkan圖形API進行渲染，因此請確保您的計算機上安裝了Vulkan(https://docs.taichi.graphics/lang/articles/misc/ggui)。GGUI支持渲染兩種類型的3D對象：三角形網格和粒子。在我們的示例中，將把布料渲染為三角形網格，把紅色球渲染為單個粒子。
GGUI表示一個三角形網格，包含兩個Taichi場：一個頂點(vertices)場和一個索引(indices)場。頂點場是一個一維場，其中每個元素提取是一個表示頂點位置的三維向量，可能由多個三角形共享。在我們的應用程序中，每個點質量都是一個三角形頂點，因此我們可以簡單地將數據從x復制到vertices：
vertices = ti.Vector.field(3, float, N * N)
@ti.kernel
def set_vertices():
for i, j in ti.ndrange(N, N):
vertices[i * N + j] = x[i, j]
請注意，每一幀都需要調用set_vertices，因為頂點位置不斷被模擬更新。
我們的布料是用一個質點的N×N網格表示，也可以被看作一個由(N-1)×(N-1)小正方形組成的網格。每個正方形都將渲染為兩個三角形。因此，總共有(N-1)×(N-1)×2個三角形。每個三角形將在頂點場中表示為3個整數，該場記錄頂點場中三角形頂點的索引。以下代碼片段捕獲了這一結構：
num_triangles = (N - 1) * (N - 1) * 2
indices = ti.field(int, num_triangles * 3)
@ti.kernel
def set_indices():
for i, j in ti.ndrange(N, N):
if i < N - 1 and j < N - 1:
square_id = (i * (N - 1)) + j
#正方形的第一個小三角形
indices[square_id * 6 + 0] = i * N + j
indices[square_id * 6 + 1] = (i + 1) * N + j
indices[square_id * 6 + 2] = i * N + (j + 1)
#正方形的第二個小三角形
indices[square_id * 6 + 3] = (i + 1) * N + j + 1
indices[square_id * 6 + 4] = i * N + (j + 1)
indices[square_id * 6 + 5] = (i + 1) * N + j
請注意，與函數set_vertices不同，函數set_indices只需要調用一次。這是因為三角形頂點的索引實際上并沒有改變——只是位置在改變。
為了將紅球渲染為粒子，我們實際上不需要準備任何數據，我們之前定義的ball_center和ball_radius變量就是GGUI所需要的全部內容。
完整代碼
至此，我們已經介紹完本文示例程序的所有核心函數!下面代碼展示了我們如何調用這些函數：
init()
set_indices()
window = ti.ui.Window("Cloth", (800, 800), vsync=True)
canvas = window.get_canvas()
scene = ti.ui.Scene()
camera = ti.ui.make_camera()
while window.running:
for i in range(30):
step()
set_vertices()
camera.position(0.5, -0.5, 2)
camera.lookat(0.5, -0.5, 0)
scene.set_camera(camera)
scene.point_light(pos=(0.5, 1, 2), color=(1, 1, 1))
scene.mesh(vertices, indices=indices, color=(0.5, 0.5, 0.5), two_sided = True)
scene.particles(ball_center, radius=ball_radius, color=(0.5, 0, 0))
canvas.scene(scene)
window.show()
需要注意的一個小細節是，我們將在主程序循環中的每一幀調用函數step()30次，而不是調用一次。這樣做的目的就是讓動畫不會運行得太慢。把上述所有代碼放在一起，整個程序應該是這樣的：
import taichi as ti
ti.init(arch=ti.cuda) # 另一種可選擇方案： ti.init(arch=ti.cpu)
N = 128
cell_size = 1.0 / N
gravity = 0.5
stiffness = 1600
damping = 2
dt = 5e-4
ball_radius = 0.2
ball_center = ti.Vector.field(3, float, (1,))
x = ti.Vector.field(3, float, (N, N))
v = ti.Vector.field(3, float, (N, N))
num_triangles = (N - 1) * (N - 1) * 2
indices = ti.field(int, num_triangles * 3)
vertices = ti.Vector.field(3, float, N * N)
def init_scene():
for i, j in ti.ndrange(N, N):
x[i, j] = ti.Vector([i * cell_size ,
j * cell_size / ti.sqrt(2),
(N - j) * cell_size / ti.sqrt(2)])
ball_center[0] = ti.Vector([0.5, -0.5, -0.0])
@ti.kernel
def set_indices():
for i, j in ti.ndrange(N, N):
if i < N - 1 and j < N - 1:
square_id = (i * (N - 1)) + j
# 1st triangle of the square
indices[square_id * 6 + 0] = i * N + j
indices[square_id * 6 + 1] = (i + 1) * N + j
indices[square_id * 6 + 2] = i * N + (j + 1)
# 2nd triangle of the square
indices[square_id * 6 + 3] = (i + 1) * N + j + 1
indices[square_id * 6 + 4] = i * N + (j + 1)
indices[square_id * 6 + 5] = (i + 1) * N + j
links = [[-1, 0], [1, 0], [0, -1], [0, 1], [-1, -1], [1, -1], [-1, 1], [1, 1]]
links = [ti.Vector(v) for v in links]
@ti.kernel
def step():
for i in ti.grouped(x):
v[i].y -= gravity * dt
for i in ti.grouped(x):
force = ti.Vector([0.0,0.0,0.0])
for d in ti.static(links):
j = min(max(i + d, 0), [N-1,N-1])
relative_pos = x[j] - x[i]
current_length = relative_pos.norm()
original_length = cell_size * float(i-j).norm()
if original_length != 0:
force += stiffness * relative_pos.normalized() * (current_length - original_length) / original_length
v[i] += force * dt
for i in ti.grouped(x):
v[i] *= ti.exp(-damping * dt)
if (x[i]-ball_center[0]).norm() <= ball_radius:
v[i] = ti.Vector([0.0, 0.0, 0.0])
x[i] += dt * v[i]
@ti.kernel
def set_vertices():
for i, j in ti.ndrange(N, N):
vertices[i * N + j] = x[i, j]
init_scene()
set_indices()
window = ti.ui.Window("Cloth", (800, 800), vsync=True)
canvas = window.get_canvas()
scene = ti.ui.Scene()
camera = ti.ui.make_camera()
while window.running:
for i in range(30):
step()
set_vertices()
camera.position(0.5, -0.5, 2)
camera.lookat(0.5, -0.5, 0)
scene.set_camera(camera)
scene.point_light(pos=(0.5, 1, 2), color=(1, 1, 1))
scene.mesh(vertices, indices=indices, color=(0.5, 0.5, 0.5), two_sided = True)
scene.particles(ball_center, radius=ball_radius, color=(0.5, 0, 0))
canvas.scene(scene)
window.show()
注意到，上述代碼總行數僅有91行!
挑戰任務
我希望你喜歡本文中提供的上述示例程序!如果的確如此，下面幾個不同挑戰等級的任務留給你：
- 【容易】隨便地調整參數：觀察stiffness，damping和dt參數的修改如何改變程序的行為。
- 【容易】把程序中的 vsync=True修改為vsync=False。這將取消程序的每秒60幀限制，進而觀察程序在你的機器上運行的速度變化情況。
- 【中等難度】在布料和球之間實現稍微復雜的交互：使其在不穿透球的情況下滑下球。
- 【中等難度】添加更多球：使布料與多個球相互作用。
- 【高級難度】完成第二個挑戰后，嘗試用另一種編程語言或Python實現同一個程序，但不使用Taichi。觀察你能獲得的最大FPS(幀/秒)是多少，以及為了獲得類似的性能你需要寫多少代碼。
小結
最后，讓我們回顧一下Taichi讓我們在上面91行Python代碼中實現的功能：
- 模擬了具有超過一萬個質量點和大約十萬個彈簧的質量彈簧系統。
- 使用@ti.kernel注釋，通過CUDA GPU或CPU上的多線程自動并行化模擬
- 通過GPU渲染器實時渲染結果
Taichi不僅讓我們能夠用少量代碼實現所有這些復雜的功能，還省去了我們學習CUDA、多線程編程或GPU渲染的麻煩。借助于Taichi，任何人都可以編寫高性能的程序。他們可以專注于代碼的算法方面，而將性能方面留給編程語言本身。這讓我們聯想到Taichi的座右銘：人人并行編程!
要了解更多關于Taichi的信息，請訪問它的要了解更多關于Taichi的信息，請訪問它的??Github頁面??，在那里你可以找到詳細的文檔以及許多Taichi項目的例子，所有這些內容都很有趣。最后，如果你也篤信為并行計算開發一種友好而強大的語言的使命，那么非常歡迎你作為開源貢獻者加入Taichi。
在我的下一篇文章中，我將討論Taichi的內部工作原理以及它如何在不同的平臺上與GPU交互從而實現計算和渲染。屆時，你將開始快樂的Taichi編程!
譯者介紹
朱先忠，51CTO社區編輯，51CTO專家博客、講師，濰坊一所高校計算機教師，自由編程界老兵一枚。早期專注各種微軟技術(編著成ASP.NET AJX、Cocos 2d-X相關三本技術圖書)，近十多年投身于開源世界(熟悉流行全棧Web開發技術)，了解基于OneNet/AliOS+Arduino/ESP32/樹莓派等物聯網開發技術與Scala+Hadoop+Spark+Flink等大數據開發技術。
原文標題：A Beginner's Guide to High-Performance Computing in Python，作者：Dunfan Lu
