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