跳转至主要内容
Version: develop

快速上手

目标受众

从事高性能计算、图形学、仿真等领域的工程师和研究人员; 希望加速 Python 中计算密集任务程序的一般开发者; 希望使用 Python 开发但部署到其它环境的开发人员。

IMPORTANT

如你是对编译器、计算机图形或高性能计算感兴趣的开发者, 并想要为 Taichi 编程语言 贡献新功能或修复漏洞, 请查看 开发者安装指南 一文,以获取更多从源代码构建 Taichi 的信息。

先决条件

Python

3.7/3.8/3.9/3.10 (64-bit)

note

如果你使用的是搭载 M1 芯片的 MacBook,Taichi 建议从 Miniforge conda 安装 Python。

支持的系统和后端

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

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

安装

要开始使用 Taichi 语言,只需使用 pip 安装:

pip install taichi

根据所使用的操作系统的不同,还需要安装一些其他的依赖项:

On Arch Linux, you need to install ncurses5-compat-libs package from the Arch User Repository: yaourt -S ncurses5-compat-libs

如果在安装 Taichi 时遇到任何问题,请参考 安装疑难解答

成功安装 Taichi 后,系统中会增加一个 CLI(命令行界面)工具,帮助快速执行多个日常任务。 要调用 CLI,请运行 tipython3 -m taiichi

示例

太极提供了一组内置的示例。 你可以运行 ti example -h 来打印帮助信息并获取可用示例的名称列表。

例如,要运行基础的 fractal(分形)示例,尝试在 shell 中运行 ti example fractal。 (也可以运行 ti example fractal.py

你可以运行 ti example -p fractal,打印该示例的源代码,或运行 ti example -P fractal,打印带语法高亮的源代码。

也可以运行 ti example -s fractal,把该示例保存到当前的工作目录中。

Hello, world!

我们将通过一个基础的分形程序示例来介绍 Taichi。

通过 python3 fractal.pyti example fractal 命令运行下面这段 Taichi 代码,你将得到一段 朱利亚集(Julia set) 动画:

image

fractal.py
import taichi as ti

ti.init(arch=ti.gpu)

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

@ti.func
def complex_sqr(z):
return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])

@ti.kernel
def paint(t: float):
for i, j in pixels: # Parallelized over all pixels
c = ti.Vector([-0.8, ti.cos(t) * 0.2])
z = ti.Vector([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))

i = 0
while gui.running:
paint(i * 0.03)
gui.set_image(pixels)
gui.show()
i = i + 1

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

import taichi as ti

Taichi 是一个嵌入在 Python 中的领域特定语言(DSL)。

让 Taichi 能像 Python 包一样易于使用——基于这个目标,我们完成了大量工程,使得每个 Python 程序员能够以最低的学习成本编写 Taichi 程序。

你甚至可以选择最喜欢的 Python 包管理系统、Python IDE 以及其他 Python 包和 Taichi 一起使用。

# Initialize Taichi and run it on CPU (default)
# - `arch=ti.gpu`: Run Taichi on GPU and has Taichi automatically detect the suitable backend
# - `arch=ti.cuda`: For the NVIDIA CUDA backend
# - `arch=ti.metal`: [macOS] For the Apple Metal backend
# - `arch=ti.opengl`: For the OpenGL backend
# - `arch=ti.vulkan`: For the Vulkan backend
# - `arch=ti.dx11`: For the DX11 backend
ti.init(arch=ti.cpu)
info
  • 在规定 arch=ti.gpu 时,Taichi 首先尝试在 CUDA 上运行。 如果设备不支持 CUDA,Taichi 转而在 Metal、OpenGL、Vulkan,或 DX11 上运行。
  • 如果没有支持的 GPU 后端(CUDA、Metal、OpenGL、Vulcan,或 DX11),Taichi 转而在 CPU 后端运行。
note

在 Windows 或 ARM 设备(如 NVIDIA Jetson)上使用 CUDA 后端时,Taichi 会默认分配 1 GB GPU 内存用于存储 field。

要覆盖此行为,请执行以下任一操作:

  • 在初始化 Taichi 时使用 ti.init(arch=ti.cuda, device_memory_GB=3.4) 来分配 3.4 GB GPU 内存;
  • 在初始化 Taichi 时使用 ti.init(arch=ti.cuda, device_memory_fraction=0.3) 来分配 30% 的 GPU 内存。

在 Windows 或 ARM 设备以外的平台上,Taichi 依靠其按需内存分配器来调整内存分配。

Field

Taichi 是一门面向数据的编程语言,以稠密或空间稀疏 field 为一等公民。

在上文的代码中,pixels = ti.field(dtype=float, shape=(n * 2, n)) 分配了一个叫做 pixels 的二维稠密 field,其尺寸为 (640, 320),元素数据类型为 float

函数与 kernel

计算发生在 Taichi 的内核函数中。

Taichi的内核由装饰器@ti.kernel来定义。 可以从 Python 中调用 kernel 来进行计算。 如果 kernel 有参数的话,必须指定参数类型。

Taichi 函数由装饰器 @ti.func 定义。 只能从 Taichi kernel 或其他 Taichi 函数中调用。

关于 Taichi kernel 和函数的更多细节请参阅 kernel 和函数 章节。

Taichi 的 kernel 与函数所用语法与 Python 语法看起来全然一致,然而 Taichi 的前端编译器会将其转换为编译型、静态类型、有词法作用域、并行执行且可微分的语言。

info

Taichi 作用域与 Python 作用域

任何被 @ti.kernel@ti.func 修饰的部分都处于 Taichi 作用域中,这部分代码由 Taichi 编译器编译。 其余部分处于 Python 作用域中。 它们是 Python 原生代码。

warning
  • 必须从 Python 作用域调用 Taichi kernel。
  • 必须从 Taichi 作用域调用 Taichi 函数。
tip

如果用 CUDA 做类比的话,ti.func 相当于 __device__ti.kernel 相当于 __global__

note
  • 不支持嵌套 kernel。
  • 支持嵌套函数。
  • 目前不支持递归函数。

并行执行的 for 循环

Taichi 自动并行执行 kernel 里位于最外层作用域的 for 循环。 有两种形式的 for 循环:

  • range-for 循环
  • struct-for 循环

range-for 循环和 Python 的 for 循环类似,区别仅在于前者在位于最外层作用域时会被并行执行。 range-for 循环支持嵌套。

@ti.kernel
def fill():
for i in range(10): # Parallelized
x[i] += i
s = 0
for j in range(5): # Serialized in each parallel thread
s += j
y[i] = s

@ti.kernel
def fill_3d():
# Parallelized for all 3 <= i < 8, 1 <= j < 6, 0 <= k < 9
for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
x[i, j, k] = i + j + k
note

Taichi 并行执行的是最外层作用域的循环,而不是最外层循环。

@ti.kernel
def foo():
for i in range(10): # Parallelized :-)
...

@ti.kernel
def bar(k: ti.i32):
if k > 42:
for i in range(10): # Serial :-(
...

struct-for 循环在遍历(稀疏)field 中的元素时尤其有用。 例如在上述 fractal.py 的代码中,for i, j in pixels 将遍历所有像素点坐标, 即 (0, 0), (0, 1), (0, 2), ... , (0, 319), (1, 0), ..., (639, 319)

note

struct-for 循环是 Taichi 进行 稀疏计算 的关键,因为这种循环只会遍历稀疏 field 中的活跃元素。 在稠密 field 中,所有元素都是活跃元素。

WARNING

结构 for 循环只能使用在内核的最外层作用域。 Taichi 并行执行的是最外层作用域的循环,而不是最外层循环。

x = [1, 2, 3]

@ti.kernel
def foo():
for i in x: # Parallelized :-)
...

@ti.kernel
def bar(k: ti.i32):
# The outermost scope is a `if` statement
if k > 42:
for i in x: # Not allowed. Struct-fors must live in the outermost scope.
...
WARNING

并行循环支持 break 语句。

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

for i in range(10):
...
break # Error!

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

GUI 系统

Taichi 提供基于 CPU 的 [GUI 系统](../visualization/gui_system. md) ,用于在屏幕上渲染结果。

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

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

仍有问题?

如果本文没能解答你遇到的问题,欢迎在 Github 创建 issue,描述问题细节。 我们随时为你提供帮助!