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

快速上手

目标受众

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

IMPORTANT

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

先决条件

Python

3.6/3.7/3.8/3.9/3.10(64位)

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

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

在 Arch Linux 系统上,需要从 Arch 的用户仓库安装 ncurses5-compat-libsyaourt -S ncurses5-compat-libs

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

Hello, world!

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

使用 python3 fractal.pyti example fractal (你可以在 命令行工具 找到更多有关 Taichi CLI 的信息) 来运行下方的 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 is the key to sparse computation in Taichi, as it will only loop over active elements in a sparse 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 provides a CPU-based GUI system for you to render your results on the screen.

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,描述问题细节。 我们随时为你提供帮助!