自從Python編程語(yǔ)言誕生以來(lái)，它的核心理念一直是最大限度地提高代碼的可讀性和簡(jiǎn)單性。Python對(duì)可讀性和簡(jiǎn)單性的追求簡(jiǎn)直達(dá)到了如癡如狂的境地。一個(gè)事實(shí)即可證實(shí)這一點(diǎn)：只要你在Python系統(tǒng)的根目錄中輸入命令“import this”后按下回車(chē)鍵，竟然馬上打印出一首英文小詩(shī)，翻譯成中文大致意思是：

“美麗勝過(guò)丑陋，顯式優(yōu)于隱式。

簡(jiǎn)單比復(fù)雜好，復(fù)雜比繁雜好。

扁平優(yōu)于嵌套，稀疏勝過(guò)密集。

可讀性很重要……”

簡(jiǎn)單總比復(fù)雜好，可讀性很重要。毫無(wú)疑問(wèn)，Python確實(shí)在實(shí)現(xiàn)這些目標(biāo)方面非常成功：它是迄今為止最友好的學(xué)習(xí)語(yǔ)言，并且一個(gè)普通的Python程序通常比等效的C++代碼短5到10倍。不幸的是，這里有一個(gè)陷阱：Python的簡(jiǎn)單性是以降低性能為代價(jià)的!事實(shí)上，Python程序比C++對(duì)應(yīng)的速度慢10到100倍。因此，似乎在速度和簡(jiǎn)單性之間存在著一種永久的權(quán)衡，任何編程語(yǔ)言都不可能同時(shí)擁有這兩者。

但是，別擔(dān)心，所有的希望都沒(méi)有失去。

Taichi可實(shí)現(xiàn)兩全其美

Taichi編程語(yǔ)言是對(duì)Python編程語(yǔ)言進(jìn)行擴(kuò)展的一種嘗試，其結(jié)構(gòu)支持通用、高性能的計(jì)算。它支持無(wú)縫地嵌入到Python中，而同時(shí)可以發(fā)揮計(jì)算機(jī)中所有的計(jì)算能力——包括多核CPU功能以及更為重要的GPU性能。

我們?cè)诒疚闹袑⒄故疽粋€(gè)使用Taichi編寫(xiě)的示例程序。該程序使用GPU對(duì)落在球體上的一塊布進(jìn)行實(shí)時(shí)物理模擬，同時(shí)渲染結(jié)果。

編寫(xiě)實(shí)時(shí)GPU物理模擬器絕非易事，但是實(shí)現(xiàn)本例程的Taichi源代碼卻異常簡(jiǎn)單。本文的其余部分將引導(dǎo)你完成整個(gè)實(shí)現(xiàn)，這樣你就可以領(lǐng)略Taichi提供的功能，以及它們的強(qiáng)大和友好程度。

在我們開(kāi)始之前，你不妨先猜猜這個(gè)程序大概由多少行代碼組成。當(dāng)然，你會(huì)在文章末尾找到這一答案。

算法概述

我們的程序?qū)岩粔K布料建模為一個(gè)質(zhì)量彈簧系統(tǒng)。更具體地說(shuō)，我們將此布料表示為點(diǎn)質(zhì)量的N×N網(wǎng)格，其中相鄰點(diǎn)由彈簧連接。下面的圖形是由斯坦福大學(xué)的Matthew Fisher提供的，正展示了這種結(jié)構(gòu)。

該質(zhì)量彈簧系統(tǒng)的運(yùn)動(dòng)受4個(gè)因素影響：

• 重力
• 彈簧的內(nèi)力
• 阻尼
• 與夾在中間的紅球碰撞

為了簡(jiǎn)單起見(jiàn)，我們忽略了布料的自碰撞。我們的程序在t=0時(shí)開(kāi)始。然后，在模擬的每一步，它將時(shí)間提前一個(gè)小常數(shù)dt。該程序通過(guò)評(píng)估上述4個(gè)因素中的每一個(gè)因素的影響來(lái)估計(jì)系統(tǒng)在這一小段時(shí)間內(nèi)會(huì)發(fā)生什么，并在時(shí)間步結(jié)束時(shí)更新每個(gè)質(zhì)量點(diǎn)的位置和速度。然后，更新的質(zhì)點(diǎn)位置用于更新屏幕上渲染的圖像。

程序開(kāi)始

盡管Taichi本身就是一種編程語(yǔ)言，但它以Python包的形式存在，只需運(yùn)行pip install Taichi即可安裝。

要在Python程序中使用Taichi，首先需要使用別名ti導(dǎo)入Taichi：

import taichi as ti

如果您的機(jī)器具有支持CUDA的Nvidia GPU，Taichi程序的性能將會(huì)得到最大程度發(fā)揮。如果是這種情況，請(qǐng)?jiān)谏鲜鰧?dǎo)入語(yǔ)句后添加以下代碼行：

ti.init(arch=ti.cuda)

如果你沒(méi)有CUDA GPU，Taichi仍然可以通過(guò)其他圖形API(如ti.metal，ti.vulkan和ti.opengl)與你的GPU交互。然而，Taichi對(duì)這些API的支持不如其對(duì)CUDA的支持那么全面。因此，目前情況下，我們使用CPU作為計(jì)算后端：

ti.init(arch=ti.cpu)

別擔(dān)心，Taichi即使只在CPU上運(yùn)行，也會(huì)運(yùn)行得很快。初始化Taichi之后，我們可以開(kāi)始聲明用于描述質(zhì)量彈簧布料的數(shù)據(jù)結(jié)構(gòu)。為此，我們添加以下代碼行：

N = 128
    x = ti.Vector.field(3, float, (N, N)) 
    v = ti.Vector.field(3, float, (N, N))

這三行將x和v聲明為大小為N×N的二維數(shù)組，其中數(shù)組的每個(gè)元素都是一個(gè)浮點(diǎn)數(shù)的三維向量。在Taichi中，數(shù)組被稱為“場(chǎng)”，這兩個(gè)場(chǎng)分別記錄點(diǎn)質(zhì)量的位置和速度。請(qǐng)注意，如果您將Taichi初始化為在CUDA GPU上運(yùn)行，這些場(chǎng)/數(shù)組將自動(dòng)存儲(chǔ)在GPU內(nèi)存中。除了布料，我們還需要定義中間的球：

ball_radius = 0.2
    ball_center = ti.Vector.field(3, float, (1,))

在這里，球中心是一個(gè)大小為1的一維場(chǎng)，其單個(gè)分量是一個(gè)三維浮點(diǎn)向量。聲明了所需的場(chǎng)之后，讓我們用t=0處的相應(yīng)數(shù)據(jù)初始化這些場(chǎng)。我們希望確保，對(duì)于同一行或同一列上的任何一對(duì)相鄰點(diǎn)，它們之間的距離等于cell_size=1.0/N。這是通過(guò)以下初始化例程來(lái)實(shí)現(xià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])

此處，你無(wú)需擔(dān)心每個(gè)x[i，j]值背后的含義——它的選擇只是為了使布料以45度角落下，參考下圖。

模擬

在每個(gè)時(shí)間步中，我們的程序都會(huì)模擬影響布料運(yùn)動(dòng)的4個(gè)因素：重力、彈簧內(nèi)力、阻尼和與紅球的碰撞。其中，重力是最容易處理的。

下面是實(shí)現(xiàn)這一點(diǎn)的代碼：

@ti.kernel
    def step():
        for i in ti.grouped(v):
        v[i].y -= gravity * dt

這里有兩點(diǎn)需要注意。首先，語(yǔ)句for i in ti.grouped(x)意味著將循環(huán)迭代x的所有元素，而不管x中有多少維度。其次，也是最重要的是：注解@ti.kernel意味著Taichi將自動(dòng)并行運(yùn)行函數(shù)中的任何頂級(jí)for循環(huán)。在本例中，Taichi將并行更新v中每個(gè)N*N向量的y分量。

接下來(lái)，我們來(lái)處理弦線的內(nèi)力計(jì)算問(wèn)題。首先，請(qǐng)注意前面圖形中的每個(gè)質(zhì)點(diǎn)最多連接到八個(gè)鄰接質(zhì)點(diǎn)。這些連接在我們的程序中表示如下：

 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]

從物理角度來(lái)看，系統(tǒng)中的每個(gè)彈簧s都用固定長(zhǎng)度l(s,0)初始化。在任何時(shí)間t，如果s的當(dāng)前長(zhǎng)度l(s,t)超過(guò)l(s,0)，則彈簧將在其端點(diǎn)上施加力，將它們拉在一起。相反，如果l(s,t)小于l(s,0)，則彈簧會(huì)將端點(diǎn)彼此推開(kāi)。這些力的大小始終與l(s,0)-l(s,0)的絕對(duì)值成正比。此交互由以下代碼段捕獲：

 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

請(qǐng)注意，這個(gè)for循環(huán)仍應(yīng)作為substep函數(shù)中的頂級(jí)for循環(huán)，該函數(shù)用@ti.kernel注解。這樣可以確保并行計(jì)算施加到每個(gè)質(zhì)點(diǎn)的彈簧力。stiffness在此是一個(gè)常數(shù)，用于控制彈簧長(zhǎng)度變化的程度。在上述程序中，我們使用stiffness =1600指定它的值。在現(xiàn)實(shí)世界中，當(dāng)彈簧振動(dòng)時(shí)，彈簧中儲(chǔ)存的能量會(huì)消散到周?chē)h(huán)境中，其振動(dòng)最終停止。為了捕捉這種效應(yīng)，在每個(gè)時(shí)間步，我們稍微降低每個(gè)點(diǎn)的速度大小：

for i in ti.grouped(x):
  v[i] *= ti.exp(-damping * dt)

在此，damping取固定值2。

我們還需要處理布料和紅球之間的碰撞。要做到這一點(diǎn)，我們只需將質(zhì)點(diǎn)與球接觸時(shí)的速度降低到0。這樣可以確保布料“掛”在球上，而不是穿透球或向下滑動(dòng)：

if (x[i]-ball_center[0]).norm() <= ball_radius: 
                v[i] = ti.Vector([0.0, 0.0, 0.0])

最后，我們用每個(gè)質(zhì)點(diǎn)的速度更新其自身的位置：

x[i] += dt * v[i]

這就是我們對(duì)一塊質(zhì)量彈簧布料進(jìn)行并行模擬所需的全部代碼。

渲染

我們將使用Taichi內(nèi)置的基于GPU的GUI系統(tǒng)(昵稱是“GGUI”)渲染布料。GGUI使用Vulkan圖形API進(jìn)行渲染，因此請(qǐng)確保您的計(jì)算機(jī)上安裝了Vulkan(https://docs.taichi.graphics/lang/articles/misc/ggui)。GGUI支持渲染兩種類型的3D對(duì)象：三角形網(wǎng)格和粒子。在我們的示例中，將把布料渲染為三角形網(wǎng)格，把紅色球渲染為單個(gè)粒子。

GGUI表示一個(gè)三角形網(wǎng)格，包含兩個(gè)Taichi場(chǎng)：一個(gè)頂點(diǎn)(vertices)場(chǎng)和一個(gè)索引(indices)場(chǎng)。頂點(diǎn)場(chǎng)是一個(gè)一維場(chǎng)，其中每個(gè)元素提取是一個(gè)表示頂點(diǎn)位置的三維向量，可能由多個(gè)三角形共享。在我們的應(yīng)用程序中，每個(gè)點(diǎn)質(zhì)量都是一個(gè)三角形頂點(diǎn)，因此我們可以簡(jiǎn)單地將數(shù)據(jù)從x復(fù)制到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]

請(qǐng)注意，每一幀都需要調(diào)用set_vertices，因?yàn)轫旤c(diǎn)位置不斷被模擬更新。

我們的布料是用一個(gè)質(zhì)點(diǎn)的N×N網(wǎng)格表示，也可以被看作一個(gè)由(N-1)×(N-1)小正方形組成的網(wǎng)格。每個(gè)正方形都將渲染為兩個(gè)三角形。因此，總共有(N-1)×(N-1)×2個(gè)三角形。每個(gè)三角形將在頂點(diǎn)場(chǎng)中表示為3個(gè)整數(shù)，該場(chǎng)記錄頂點(diǎn)場(chǎng)中三角形頂點(diǎn)的索引。以下代碼片段捕獲了這一結(jié)構(gòu)：

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
            #正方形的第一個(gè)小三角形
            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)
            #正方形的第二個(gè)小三角形
            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

請(qǐng)注意，與函數(shù)set_vertices不同，函數(shù)set_indices只需要調(diào)用一次。這是因?yàn)槿切雾旤c(diǎn)的索引實(shí)際上并沒(méi)有改變——只是位置在改變。

為了將紅球渲染為粒子，我們實(shí)際上不需要準(zhǔn)備任何數(shù)據(jù)，我們之前定義的ball_center和ball_radius變量就是GGUI所需要的全部?jī)?nèi)容。

完整代碼

至此，我們已經(jīng)介紹完本文示例程序的所有核心函數(shù)!下面代碼展示了我們?nèi)绾握{(diào)用這些函數(shù)：

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

需要注意的一個(gè)小細(xì)節(jié)是，我們將在主程序循環(huán)中的每一幀調(diào)用函數(shù)step()30次，而不是調(diào)用一次。這樣做的目的就是讓動(dòng)畫(huà)不會(huì)運(yùn)行得太慢。把上述所有代碼放在一起，整個(gè)程序應(yīng)該是這樣的：

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

注意到，上述代碼總行數(shù)僅有91行!

挑戰(zhàn)任務(wù)

我希望你喜歡本文中提供的上述示例程序!如果的確如此，下面幾個(gè)不同挑戰(zhàn)等級(jí)的任務(wù)留給你：

• 【容易】隨便地調(diào)整參數(shù)：觀察stiffness，damping和dt參數(shù)的修改如何改變程序的行為。
• 【容易】把程序中的 vsync=True修改為vsync=False。這將取消程序的每秒60幀限制，進(jìn)而觀察程序在你的機(jī)器上運(yùn)行的速度變化情況。
• 【中等難度】在布料和球之間實(shí)現(xiàn)稍微復(fù)雜的交互：使其在不穿透球的情況下滑下球。
• 【中等難度】添加更多球：使布料與多個(gè)球相互作用。
• 【高級(jí)難度】完成第二個(gè)挑戰(zhàn)后，嘗試用另一種編程語(yǔ)言或Python實(shí)現(xiàn)同一個(gè)程序，但不使用Taichi。觀察你能獲得的最大FPS(幀/秒)是多少，以及為了獲得類似的性能你需要寫(xiě)多少代碼。

小結(jié)

最后，讓我們回顧一下Taichi讓我們?cè)谏厦?1行Python代碼中實(shí)現(xiàn)的功能：

• 模擬了具有超過(guò)一萬(wàn)個(gè)質(zhì)量點(diǎn)和大約十萬(wàn)個(gè)彈簧的質(zhì)量彈簧系統(tǒng)。
• 使用@ti.kernel注釋，通過(guò)CUDA GPU或CPU上的多線程自動(dòng)并行化模擬
• 通過(guò)GPU渲染器實(shí)時(shí)渲染結(jié)果

Taichi不僅讓我們能夠用少量代碼實(shí)現(xiàn)所有這些復(fù)雜的功能，還省去了我們學(xué)習(xí)CUDA、多線程編程或GPU渲染的麻煩。借助于Taichi，任何人都可以編寫(xiě)高性能的程序。他們可以專注于代碼的算法方面，而將性能方面留給編程語(yǔ)言本身。這讓我們聯(lián)想到Taichi的座右銘：人人并行編程!

要了解更多關(guān)于Taichi的信息，請(qǐng)?jiān)L問(wèn)它的要了解更多關(guān)于Taichi的信息，請(qǐng)?jiān)L問(wèn)它的??Github頁(yè)面??，在那里你可以找到詳細(xì)的文檔以及許多Taichi項(xiàng)目的例子，所有這些內(nèi)容都很有趣。最后，如果你也篤信為并行計(jì)算開(kāi)發(fā)一種友好而強(qiáng)大的語(yǔ)言的使命，那么非常歡迎你作為開(kāi)源貢獻(xiàn)者加入Taichi。

在我的下一篇文章中，我將討論Taichi的內(nèi)部工作原理以及它如何在不同的平臺(tái)上與GPU交互從而實(shí)現(xiàn)計(jì)算和渲染。屆時(shí)，你將開(kāi)始快樂(lè)的Taichi編程!

譯者介紹

朱先忠，51CTO社區(qū)編輯，51CTO專家博客、講師，濰坊一所高校計(jì)算機(jī)教師，自由編程界老兵一枚。早期專注各種微軟技術(shù)(編著成ASP.NET AJX、Cocos 2d-X相關(guān)三本技術(shù)圖書(shū))，近十多年投身于開(kāi)源世界(熟悉流行全棧Web開(kāi)發(fā)技術(shù))，了解基于OneNet/AliOS+Arduino/ESP32/樹(shù)莓派等物聯(lián)網(wǎng)開(kāi)發(fā)技術(shù)與Scala+Hadoop+Spark+Flink等大數(shù)據(jù)開(kāi)發(fā)技術(shù)。

原文標(biāo)題：A Beginner's Guide to High-Performance Computing in Python，作者：Dunfan Lu