DZone>Web Dev Zone>Python高性能计算初学者指南
Python高性能计算的初学者指南
Taichi编程语言是一种尝试，它通过构造来扩展Python编程语言，从而实现通用的、高性能的计算。
通过
-
1月24日，22 - Web Dev Zone -教程
喜欢 (3)
评论
保存
鸣叫
4.66K浏览次数
加入DZone社区，获得完整的会员体验。
自从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 byN 的点-质量网格，其中相邻的点由弹簧连接。下面的图片由Matthew Fisher提供，说明了这种结构。
这个质量-弹簧系统的运动受到四个因素的影响。
- 重力
- 弹簧的内力
- 阻尼
- 与中间的红球的碰撞
为了简化本博客，我们忽略了布的自我碰撞。我们的程序开始于时间t = 0 。然后，在模拟的每一步，它以一个小常数推进时间dt 。程序通过评估上述4个因素的影响来估计系统在这一小段时间内发生了什么，并在时间步长结束时更新每个质量点的位置和速度。然后，质量点的更新位置被用来更新屏幕上呈现的图像。
入门
尽管Taichi本身就是一种编程语言，但它是以Python包的形式存在的，只需运行pip install taichi 即可安装。
要在一个Python程序中开始使用Taichi，可以用别名ti 来导入它。
import taichi as ti
如果你的机器有一个支持CUDA的Nvidia GPU，那么Taichi程序的性能将得到最大化。如果是这种情况，请在导入后添加以下一行代码。ti.init(arch=ti.cuda)`` 如果你没有CUDA GPU，Taichi仍然可以通过其他图形API与你的GPU交互，如ti.metal、ti.vulkan和ti.opengl。然而，Taichi对这些API的支持并不像它对CUDA的支持那样完整，所以，现在，使用CPU后端：ti.init(arch=ti.cpu)不要担心，即使Taichi只在CPU上运行，它也是非常快的。在初始化了Taichi之后，我们可以开始声明用于描述质量弹簧布的数据结构。我们添加以下几行代码。
Python
N = 128
x = ti.Vector.field(3, float, (N, N))
v = ti.Vector.field(3, float, (N, N))
这三行声明x 和v 是大小为N byN 的二维数组，其中数组的每个元素都是一个浮点数的三维向量。在Taichi中，数组被称为 "场"，这两个场分别记录了点质量的位置和速度。注意，如果你将Taichi初始化为在CUDA GPU上运行，这些字段/数组将自动存储在GPU内存中。除了布以外，我们还需要定义中间的球。
Python
ball_radius = 0.2
ball_center = ti.Vector.field(3, float, (1,))
这里，球心是一个大小为1的一维字段，其单一分量是一个3维浮点向量。在声明了所需的字段后，让我们用相应的数据初始化这些字段：t = 0 。我们希望确保，对于同一行或同一列的任何一对相邻的点，它们之间的距离等于cell_size = 1.0 / N 。这可以通过下面的初始化程序来保证。
Python
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度角落下，如gif所示。
仿真
在每个时间点上，我们的程序模拟了影响布匹运动的四种情况：重力、弹簧的内力、阻尼和与红球的碰撞。重力是最直接的处理方式。
下面是处理这个问题的代码。
Python
@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-loop。在这种情况下，taichi将并行地更新v 中每个N*N 向量的y 部分。
接下来，我们将处理字符串的内力问题。首先，从前面的插图中注意到，每个质量点最多与八个邻居相连。这些链接在我们的程序中表示如下。
Python
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) 的绝对值成正比。这种相互作用可以通过下面的代码片断体现出来。
Python
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 循环仍应作为顶层for 循环放在substep 函数中，该函数被注释为@ti.kernel 。这样可以确保应用于每个质量点的弹簧力是平行计算的。变量stiffness 是一个常数，控制弹簧抵抗其长度变化的程度。在我们的程序中，我们将使用stiffness = 1600 。在现实世界中，当弹簧振荡时，储存在弹簧中的能量会消散到周围环境中，其振荡最终会停止。为了捕捉这种效果，在每个时间步长中，我们稍微降低每个点的速度大小。
Python
for i in ti.grouped(x):
v[i] *= ti.exp(-damping * dt)
其中damping ，取固定值为2 。
我们还需要处理布和红球之间的碰撞。要做到这一点，我们只需将一个质量点的速度降低到0，只要它与球接触。这可以确保布 "挂 "在球上，而不是穿透它或滑下来。
Python
if (x[i]-ball_center[0]).norm() <= ball_radius:
v[i] = ti.Vector([0.0, 0.0, 0.0])
最后，我们用它的速度来更新每个质量点的位置。
Python
x[i] += dt * v[i]
就这样了!这就是我们需要的所有代码，以执行质量-弹簧-布的并行模拟。
渲染
我们将使用Taichi内置的基于GPU的GUI系统（昵称为 "GGUI"）来渲染布。GGUI使用Vulkan图形API进行渲染，所以请确保你的机器上安装了Vulkan。GGUI支持渲染两种类型的3D对象：三角形网格和粒子。我们将把布渲染成一个三角形网格，并把红球渲染成一个粒子。
GGUI用两个taichi字段表示一个三角形网格：一个是vertices ，一个是indices 。vertices 字段是一个一维字段，其中每个元素的提取是一个三维向量，代表一个顶点的位置，可能由多个三角形共享。在我们的应用中，每个点质量都是一个三角形顶点，所以我们可以简单地将数据从x 复制到vertices 。
Python
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 byN 的质点网格表示，它也可以被看作是一个N-1 byN-1 的小方块网格。这些方块中的每一个都将被渲染成两个三角形。因此，总共有(N - 1) * (N - 1) * 2 个三角形。每个三角形将在vertices 字段中表示为3个整数，它在vertices 字段中记录了三角形顶点的索引。下面的代码片段抓住了这个结构。
Python
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
# 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
注意，与set_vertices 不同，函数set_indices 只需要调用一次。这是因为三角形顶点的指数实际上并没有变化 -- 只是位置在变化而已。
为了将红球渲染成粒子，我们实际上不需要准备任何数据，我们之前定义的ball_center 和ball_radius 变量是GGUI所需要的全部内容。
把所有东西放在一起
在这一点上，我们已经涵盖了程序的所有核心功能!下面是我们将如何调用这些函数。
Python
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()
需要注意的一件小事是，我们将在主程序循环的每一帧中调用它30次，而不是调用一次step() 。这样做的目的只是为了使动画的运行不会太慢。把所有东西放在一起，整个程序看起来应该是这样的。
Python
import taichi as ti
ti.init(arch=ti.cuda) # Alternatively, 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 FPS限制。观察该程序在你的机器上能运行多快。 - [中]在布和球之间实现一个稍微复杂的互动：让它在不穿透球的情况下沿着球滑下来。
- [中]增加更多的球：让布与一个以上的球互动。
- [困难] 完成第二个挑战后，尝试用另一种编程语言，或用Python但不使用Taichi来实现同样的程序。看看你能获得的最大FPS是多少，以及为了获得类似的性能，你需要写多少代码。
临别赠言
让我们回顾一下Taichi使我们在91行Python中完成的工作。
- 模拟一个有超过一万个质量点和大约十万个弹簧的质量-弹簧系统。
- 使用
@ti.kernel注释，通过CUDA GPU或CPU的多线程自动并行化仿真。 - 通过GPU渲染器实时渲染结果。
Taichi不仅让我们用少量的代码实现了所有这些复杂的功能，还让我们免去了学习CUDA、多线程编程或GPU渲染的麻烦。有了Taichi，任何人都可以编写高性能的程序，他们可以专注于代码的算法方面，而把性能方面的问题留给编程语言本身。这使我们想到了Taichi的座右铭。每个人的并行编程。
要了解更多关于Taichi的信息，请访问它的Github页面，在那里你可以找到详细的文档，以及不少Taichi程序的例子，这些例子都很有趣。最后，如果你相信制作一个友好而强大的并行计算语言的使命，我们非常欢迎你加入Taichi成为一个开源的贡献者。
在我的下一篇文章中，我将讨论Taichi的内部工作原理，以及它如何与不同平台上的GPU进行计算和渲染的互动。在那之前，祝你编码愉快
主题。
高性能计算, Taichi, python, webdev, 教程, 编程语言, web开发
在DZone上发表，得到了Dunfan Lu的许可。点击这里查看原文。
DZone贡献者所表达的观点属于他们自己。
DZone上的热门文章
评论
