跳转至主要内容

起步

欢迎来到 Taichi 语言的文档!

安装#

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

python3 -m pip install taichi
note

目前,Taichi 支持的 Python 版本有 3.6/3.7/3.8/3.9 (64 位)。

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

在 Ubuntu 19.04+ 上,你需要安装 libtinfo5:

sudoapt install libtinfo5

如果你在安装 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 = 320pixels = ti.field(dtype=float, shape=(n * 2, n))
@ti.funcdef complex_sqr(z):    return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])
@ti.kerneldef 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))
for i in range(1000000):    paint(i * 0.03)    gui.set_image(pixels)    gui.show()

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

import taichi as ti#

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

为了使 Taichi 能像 Python 包一样易于使用,让每个 Python 程序员能够以最低的学习成本编写 Taichi 程序,我们做了大量的工作。

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

# 在 GPU 上运行,自动选择后端ti.init(arch=ti.gpu)
# 在 GPU 上运行, 使用 NVIDIA CUDA 后端ti.init(arch=ti.cuda)# 在 GPU 上运行, 使用 OpenGL 后端ti.init(arch=ti.opengl)# 在 GPU 上运行, 使用苹果 Metal 后端(仅对 macOS)有效ti.init(arch=ti.metal)
# 在 CPU 上运行 (默认)ti.init(arch=ti.cpu)
info

不同操作系统所支持的后端:

平台CPUCUDAOpenGLMetalC source
WindowsOKOKOKN/AN/A
LinuxOKOKOKN/AOK
macOSOKN/AN/AOKN/A

(OK:支持;N/A:不可用)

With arch=ti.gpu, Taichi will first try to run with CUDA. If CUDA is not supported on your machine, Taichi will fall back on Metal or OpenGL. If no GPU backend (CUDA, Metal, or OpenGL) is supported, Taichi will fall back on CPUs.

note

When used with the CUDA backend on Windows or ARM devices (e.g., NVIDIA Jetson), Taichi allocates 1 GB GPU memory for field storage by default.

You can override this behavior by initializing with ti.init(arch=ti.cuda, device_memory_GB=3.4) to allocate 3.4 GB GPU memory, or ti.init(arch=ti.cuda, device_memory_fraction=0.3) to allocate 30% of the total GPU memory.

On other platforms, Taichi will make use of its on-demand memory allocator to allocate memory adaptively.

Fields#

Taichi is a data-oriented programming language where dense or spatially-sparse fields are the first-class citizens.

In the code above, pixels = ti.field(dtype=float, shape=(n * 2, n)) allocates a 2D dense field named pixels of size (640, 320) and element data type float.

Functions and kernels#

Computation resides in Taichi kernels and Taichi functions.

Taichi kernels are defined with the decorator @ti.kernel. They can be called from Python to perform computation. Kernel arguments must be type-hinted (if any).

Taichi functions are defined with the decorator @ti.func. They can only be called by Taichi kernels or other Taichi functions.

See syntax for more details about Taichi kernels and functions.

The language used in Taichi kernels and functions looks exactly like Python, yet the Taichi frontend compiler converts it into a language that is compiled, statically-typed, lexically-scoped, parallel and differentiable.

info

Taichi-scopes v.s. Python-scopes:

Everything decorated with @ti.kernel and @ti.func is in Taichi-scope and hence will be compiled by the Taichi compiler.

Everything else is in Python-scope. They are simply Python native code.

caution

Taichi kernels must be called from the Python-scope. Taichi functions must be called from the Taichi-scope.

tip

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

note

Nested kernels are not supported.

Nested functions are supported.

Recursive functions are not supported for now.

Parallel for-loops#

For loops at the outermost scope in a Taichi kernel is automatically parallelized. For loops can have two forms, i.e. range-for loops and struct-for loops.

Range-for loops are no different from Python for loops, except that they will be parallelized when used at the outermost scope. Range-for loops can be nested.

@ti.kerneldef 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.kerneldef 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

It is the loop at the outermost scope that gets parallelized, not the outermost loop.

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

Struct-for loops are particularly useful when iterating over (sparse) field elements. In the fractal.py above, for i, j in pixels loops over all the pixel coordinates, i.e., (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. In dense fields, all elements are active.

caution

Struct-for loops must live at the outer-most scope of kernels.

It is the loop at the outermost scope that gets parallelized, not the outermost loop.

@ti.kerneldef foo():    for i in x:        ...
@ti.kerneldef 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.            ...
caution

break is not supported in parallel loops:

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

Interacting with other Python packages#

Python-scope data access#

Everything outside Taichi-scopes (ti.func and ti.kernel) is simply Python code. In Python-scopes, you can access Taichi field elements using plain indexing syntax. For example, to access a single pixel of the rendered image in Python-scope, you can simply use:

import taichi as tipixels = ti.field(ti.f32, (1024, 512))
pixels[42, 11] = 0.7  # store data into pixelsprint(pixels[42, 11]) # prints 0.7

Sharing data with other packages#

Taichi provides helper functions such as from_numpy and to_numpy to transfer data between Taichi fields and NumPy arrays, so that you can also use your favorite Python packages (e.g., numpy, pytorch, matplotlib) together with Taichi as below:

import taichi as tipixels = ti.field(ti.f32, (1024, 512))
import numpy as nparr = np.random.rand(1024, 512)pixels.from_numpy(arr)   # load numpy data into taichi fields
import matplotlib.pyplot as pltarr = pixels.to_numpy()  # store taichi data into numpy arraysplt.imshow(arr)plt.show()
import matplotlib.cm as cmcmap = cm.get_cmap('magma')gui = ti.GUI('Color map')while gui.running:    render_pixels()    arr = pixels.to_numpy()    gui.set_image(cmap(arr))    gui.show()

See Interacting with external arrays for more details.

What's next?#

Now we have gone through core features of the Taichi programming language using the fractal example, feel free to dive into the language concepts in the next section, or jump to the advanced topics, such as the Metaprogramming or Differentiable programming. Remember that you can use the search bar at the top right corner to search for topics or keywords at any time!

If you are interested in joining the Taichi community, we strongly recommend you take some time to familiarize yourself with our contribution guide.

We hope you enjoy your adventure with Taichi!