基於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

#

以上是基於Taichi的Python高效能運算入門指南的詳細內容。更多資訊請關注PHP中文網其他相關文章！