跳转至主要内容
Version: v1.3.0

Hello, World!

Taichi 是一个基于 Python 的领域特定语言,专为高性能并行计算设计。

在 Python 中编写计算密集型任务时,遵循以下几条额外规则,即可利用 Taichi 的高性能计算能力。 概括地说,Taichi 提供两个装饰器 @ti.func@ti.kernel,它们指示 Taichi 接管任务。 Taichi 的即时(JIT)编译器将被装饰的函数编译成机器代码,其后对这些函数的所有调用都在多核 CPU 或 GPU 上执行。 在典型的计算密集场景中,如数值模拟,Taichi 相较原生 Python 代码可将性能加速 50 ~ 100 倍。

Taichi 还有一个内置的预先(AOT)编译系统,用于将代码导出为二进制/shader 文件。你可在 C/C++ 中调用导出文件,脱离 Python 环境运行。 获取更多信息,请查看 AOT 部署

先决条件

  • Python:3.7/3.8/3.9/3.10(64 位)
  • 操作系统:Windows、OS X 和 Linux(64 位)

安装

Taichi 是一个 PyPI 包:

pip install taichi

Hello, world!

一个基础的分形示例——朱利亚分形(Julia fractal)可以作为了解 Taichi 编程语言基本原理的良好切入点。

fractal.py
import taichi as ti
import taichi.math as tm

ti.init(arch=ti.gpu)

n = 320
pixels = ti.field(dtype=float, shape=(n * 2, n))

@ti.func
def complex_sqr(z): # complex square of a 2D vector
return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])

@ti.kernel
def paint(t: float):
for i, j in pixels: # Parallelized over all pixels
c = tm.vec2(-0.8, tm.cos(t) * 0.2)
z = tm.vec2(i / n - 1, j / n - 0.5) * 2
iterations = 0
while z.norm() < 20 and iterations < 50:
z = complex_sqr(z) + c
iterations += 1
pixels[i, j] = 1 - iterations * 0.02

gui = ti.GUI("Julia Set", res=(n * 2, n))

for i in range(1000000):
paint(i * 0.03)
gui.set_image(pixels)
gui.show()

将以上代码保存到本地机器,并运行此程序。

你会看到以下动画:

image

让我们来深入剖析一下这个简单的 Taichi 程序。

导入 Taichi

import taichi as ti
import taichi.math as tm

开头两行导入 Taichi 和 Taichi 的 math 模块。 math 模块包含:

  • 常用的数学函数。
  • 内置的小维度向量和矩阵。 例如 vec2 表示二维实向量,mat3 表示 3×3 实矩阵。 获取更多信息,请查看 数学模块
ti.init(arch=ti.gpu)

此行调用 ti.init(),根据传入参数定制 Taichi 运行时。 目前,我们只介绍最重要的参数:arch

参数 arch 指定后端来执行已编译的代码。 后端可以是 ti.cputi.gpu。 当指定后端为 ti.gpu 时,Taichi 按照 ti.cudati.vulkanti.opengl/ti.metal 的优先顺序选择后端。 如果没有 GPU 架构可用,Taichi 回退到 CPU 设备。

你也可以直接指定要使用的 GPU 后端。 例如,设置 arch=ti.cuda,在 CUDA 上运行程序。 如果目标架构不可用,Taichi 会报错。 关于 ti.init() 的更多信息,请查看 全局设置

定义 field

让我们看接下来的两行。

n = 320
pixels = ti.field(dtype=float, shape=(n * 2, n))

这一代码片段定义了一个 field,其形状为 (640, 320),元素是浮点数。

field 是 Taichi 最重要、最常用的数据结构。 可以将之类比于 NumPy 的 ndarray 或 PyTorch 的 tensor。 但 Taichi 的 field 更加灵活。 例如,field 可以是 空间稀疏结构,并且能轻松 切换不同的数据布局

我们将在其他基于场景的教程中介绍 field 更高级的功能。 目前,只需了解 pixels 这个 field 是一个稠密二维数组。

kernel 与函数

9~22 行定义了两个函数,一个被 @ti.func 装饰,另一个被 @ti.kernel 装饰。 它们分别被称为 Taichi 函数kernel。 Taichi 函数和 kernel 不是由 Python 解释器执行的,而是由 Taichi 的 JIT 编译器接管,并部署到多核 CPU 或 GPU 并行执行。

Taichi 函数和 kernel 之间的主要差异:

  • kernel 是 Taichi 介入并接管任务的入口点。 你可以在程序中随时随地调用 kernel,但只能从 kernel 中或另一个 Taichi 函数中调用 Taichi 函数。 在上面的例子中,Taichi 函数 complet_sqr 被 kernel paint 调用。
  • kernel 接收的参数及返回的结果都必须带有类型提示。 但 Taichi 函数不需要类型提示。 在上面的示例中,kernel paint 中的参数 t 带有类型提示;而 Taichi 函数 complex_sqr 中的参数 z 没有。
  • Taichi支持嵌套函数,但不支持嵌套 kernel。
  • Taichi 支持 Taichi 函数的递归调用。
tip

​ For those who come from the world of CUDA, ti.func corresponds to __device__ and ti.kernel corresponds to __global__.

如果用 OpenGL 类比,ti.func 相当于 GLSL 的常用函数,ti.kernel 相当于 compute shader

并行执行的 for 循环

高性能的关键在于第 15 行代码:

for i, j in pixels:

这是位于 Taichi kernel 中最外层作用域的 for 循环,因此被自动并行化处理。

Taichi 提供了一个便捷的语法糖:并行执行所有位于 kernel 最外层作用域的 for 循环。 这意味着你可以使用一个普通的循环并行执行任务,而无需了解线程的分配与回收或内存管理。

注意 field pixels 被用作迭代器。 作为 field 的元素索引,ij 分别是 [0, 2*n-1][0, n-1] 范围内的整数。 (i, j) 并行遍历 (0, 0)(0, 1) ... (0, n-1)(1, n-1) ... (2*n-1, n-1)

要记住,不在最外层作用域的 for 循环不会被并行,而是串行执行:

@ti.kernel
def fill():
total = 0
for i in range(10): # Parallelized
for j in range(5): # Serialized in each parallel thread
total += i * j

if total > 10:
for k in range(5): # Not parallelized because it is not at the outermost scope

也可以设置 ti.loop_config(serialize=True) 来串行执行最外层作用域的 for 循环。 获取更多信息,请查看 串行执行一个并行的 for 循环

WARNING

并行循环支持 break 语句:

@ti.kernel
def foo():
for i in x:
...
break # Error!

@ti.kernel
def foo():
for i in x:
for j in range(10):
...
break # OK!

显示结果

18~23 行代码利用 Taichi 内置的 GUI 系统 在屏幕上渲染运行结果:

gui = ti.GUI("Julia Set", res=(n * 2, n))
# Sets the window title and the resolution

for i in range(1000000):
paint(i * 0.03)
gui.set_image(pixels)
gui.show()

程序迭代 pixels 超过 1,000,000 次,储存在 pixels 中的分形图案也相应更新。 调用 gui.set_image() 设置窗口,调用 gui.show() 在屏幕上显示同步结果。

要点总结

祝贺你! 在完整学习了上述的简短示例后,你已经掌握了 Taichi 最重要的特性:

  • Taichi 编译 Taichi 函数和 kernel 并在指定后端运行。
  • Taichi kernel 中最外层作用域的 for 循环被自动并行执行。
  • Taichi 提供了一个灵活的数据容器 field,你可以使用索引遍历 field 中的元素。

Taichi 示例

朱利亚分形是 Taichi 提供的示范 demo 之一。 Taichi gallery 中包含更多精选 demo:

ti gallery

运行上述指令,将出现此窗口:

image

查看 Taichi 示例的完整列表,运行 ti example。 其他有用的命令行包括:

  • ti example -p fractal/ti example -P fractal 打印分形示例的源代码。
  • ti example -s fractal 保存示例到当前工作目录。

支持的系统和后端

下表列出了 Taichi 支持的操作系统和在这些系统上支持的后端:

平台CPUCUDAOpenGLMetalVulkan
Windows✔️✔️✔️N/A✔️
Linux✔️✔️✔️N/A✔️
macOS✔️N/AN/A✔️✔️
  • ✔️:支持
  • N/A:暂不支持

现在,你可以更进一步,开始编写自己的 Taichi 程序了。 以下文档提供了 Taichi 在经典应用场景下的使用指南: