你好，世界！

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

Running the Taichi code below (python3 fractal.py or ti example fractal) will give you an animation of Julia set:

https://github.com/yuanming-hu/public_files/raw/master/graphics/taichi/fractal.gif

  1. # fractal.py
  3. import taichi as ti
  5. ti.init(arch=ti.gpu)
  7. n = 320
  8. pixels = ti.var(dt=ti.f32, shape=(n * 2, n))
  10. @ti.func
  11. def complex_sqr(z):
  12.   return ti.Vector([z[0] ** 2 - z[1] ** 2, z[1] * z[0] * 2])
  14. @ti.kernel
  15. def paint(t: ti.f32):
  16.   for i, j in pixels: # 对于所有像素，并行执行
  17.     c = ti.Vector([-0.8, ti.sin(t) * 0.2])
  18.     z = ti.Vector([float(i) / n - 1, float(j) / n - 0.5]) * 2
  19.     iterations = 0
  20.     while z.norm() < 20 and iterations < 50:
  21.       z = complex_sqr(z) + c
  22.       iterations += 1
  23.     pixels[i, j] = 1 - iterations * 0.02
  25. gui = ti.GUI("Fractal", (n * 2, n))
  27. for i in range(1000000):
  28.   paint(i * 0.03)
  29.   gui.set_image(pixels)
  30.   gui.show()

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

import taichi as ti

Taichi 是一种嵌入在 Python 中的领域特定语言(Domain-Specific Language, DSL )。为了使 Taichi 能像 Python 包一样易于使用，基于这个目标我们做了大量的工程工作——使得每个 Python 程序员能够以最低的学习成本编写 Taichi 程序。你甚至可以选择你最喜欢的 Python 包管理系统、Python IDE 以及其他 Python 包和 Taichi 一起结合使用。

可移植性

Taichi 既能在 CPU，也能在 GPU 上运行。你只需根据你的硬件平台初始化 Taichi：

  1. # 在 GPU 上运行，自动选择后端
  2. ti.init(arch=ti.gpu)
  4. # 在 GPU 上运行， 使用 NVIDIA CUDA 后端
  5. ti.init(arch=ti.cuda)
  6. # 在 GPU 上运行， 使用 OpenGL 后端
  7. ti.init(arch=ti.opengl)
  8. # 在 GPU 上运行， 使用苹果 Metal 后端（仅对 OS X）有效
  9. ti.init(arch=ti.metal)
  11. # 在 CPU 上运行 (默认)
  12. ti.init(arch=ti.cpu)

注解

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

平台CPUCUDAOpenGLMetal
Windows可用可用可用不可用
Linux可用可用可用不可用
Mac OS X可用不可用不可用可用

（可用: 该系统上有最完整的支持；不可用: 由于平台限制，我们无法实现该后端）

在参数 arch=ti.gpu 下，Taichi 将首先尝试在 CUDA 上运行。如果你的设备不支持 CUDA，那么 Taichi 将会转到 Metal 或 OpenGL。如果所在平台不支持 GPU 后端（CUDA、Metal 或 OpenGL），Taichi 将默认在 CPU 运行。

注解

当在 Windows 平台 或者 ARM 设备（如 NVIDIA Jetson）上使用 CUDA 后端时， Taichi 会默认分配 1 GB 显存用于张量存储。如需重载显存分配，你可以在初始化的时候通过 ti.init(arch=ti.cuda, device_memory_GB=3.4) 来分配 3.4 GB 显存，或者使用 ti.init(arch=ti.cuda, device_memory_fraction=0.3) 来分配所有可用显存的 30%.

在其他平台上， Taichi 将会使用它的自适应内存分配器来动态分配内存。

（稀疏）张量

Taichi 是一门面向数据的程序设计语言，其中（稠密、稀疏）张量是第一类公民(First-class Citizen)。在 Sparse computation (WIP) 这一章节，你可以了解到更多关于稀疏张量的详细信息。

在以上代码中，pixels = ti.var(dt=ti.f32, shape=(n * 2, n)) 分配了一个叫做 pixels 的二维张量，大小是 (640, 320) ，数据类型是 ti.f32 （即，C语言中的 float）.

函数与内核

计算发生在 Taichi 的 内核(kernel) 中。内核的参数必须显式指定类型。Taichi 内核与函数中所用的语法，看起来和 Python 的很像，然而 Taichi 的前端编译器会将其转换为 编译型，静态类型，有词法作用域，并行执行且可微分 的语言。

Taichi 的 函数 可以被 Taichi 内核和其他 Taichi 函数调用，你应该使用关键字 ti.func 来进行定义。

注解

Taichi 作用域与 Python 作用域：任何被 @ti.kernel 和 @ti.func 修饰的函数体都处于 Taichi 作用域中，这些代码会由 Taichi 编译器编译。而在 Taichi 作用域之外的就都是 Python 作用域了，它们是单纯的 Python 代码。

警告

Taichi 内核只有在 Python 作用域中才能调用，也就是说，我们不支持嵌套内核。同时，虽然不同函数可以嵌套调用，但 Taichi 暂不支持递归函数 。

Taichi 函数只有在 Taichi 作用域中才能调用。

如果用 CUDA 做类比的话， ti.func 就像是 __device__，ti.kernel 就像是 __global__。

并行执行的for循环

最外层作用域的 for 循环是被 自动并行执行 的。Taichi 的 for 循环具有两种形式， 区间 for 循环，和 结构 for 循环。

区间 for 循环 和普通的 Python for 循环没多大区别，只是 Taichi 最外层的 for 会并行执行而已。区间 for 循环可以嵌套。

  1. @ti.kernel
  2. def fill():
  3.   for i in range(10): # 并行执行
  4.     x[i] += i
  6.     s = 0
  7.     for j in range(5): # 在每个并行的线程中顺序执行
  8.       s += j
  10.     y[i] = s
  12. @ti.kernel
  13. def fill_3d():
  14.   # 在区间 3 <= i < 8, 1 <= j < 6, 0 <= k < 9 上展开并行
  15.   for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
  16.     x[i, j, k] = i + j + k

注解

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

  1. @ti.kernel
  2. def foo():
  3.     for i in range(10): # 并行 :-)
  4.         …
  6. @ti.kernel
  7. def bar(k: ti.i32):
  8.     if k > 42:
  9.         for i in range(10): # 串行 :-(
  10.             …

结构 for 循环 在遍历（稀疏）张量元素的时候很有用。例如在上述的代码中，for i, j in pixels 将遍历所有像素点坐标, 即 (0, 0), (0, 1), (0, 2), ... , (0, 319), (1, 0), ..., (639, 319)。

注解

结构 for 循环是 Taichi 稀疏计算（Sparse computation (WIP)）的关键，它只会遍历稀疏张量中的活跃元素。对于稠密张量而言，所有元素都是活跃元素。

警告

结构 for 循环只能使用在内核的最外层作用域。

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

  1. @ti.kernel
  2. def foo():
  3.     for i in x:
  4.         …
  6. @ti.kernel
  7. def bar(k: ti.i32):
  8.     # 最外层作用域是 `if` 语句
  9.     if k > 42:
  10.         for i in x: # 语法错误。结构 for 循环 只能用于最外层作用域
  11.             …

警告

并行循环不支持 break 语句：

  1. @ti.kernel
  2. def foo():
  3.   for i in x:
  4.       ...
  5.       break # 错误：并行执行的循环不能有 break
  7. @ti.kernel
  8. def foo():
  9.   for i in x:
  10.       for j in y:
  11.           ...
  12.           break # 可以

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 tensor elements using plain indexing syntax. For example, to access a single pixel of the rendered image in Python-scope, simply use:

  1. import taichi as ti
  2. pixels = ti.var(ti.f32, (1024, 512))
  4. pixels[42, 11] = 0.7  # store data into pixels
  5. print(pixels[42, 11]) # prints 0.7

Sharing data with other packages

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

  1. import taichi as ti
  2. pixels = ti.var(ti.f32, (1024, 512))
  4. import numpy as np
  5. arr = np.random.rand(1024, 512)
  6. pixels.from_numpy(arr)   # load numpy data into taichi tensors
  8. import matplotlib.pyplot as plt
  9. arr = pixels.to_numpy()  # store taichi data into numpy arrays
  10. plt.imshow(arr)
  11. plt.show()
  13. import matplotlib.cm as cm
  14. cmap = cm.get_cmap('magma')
  15. gui = ti.GUI('Color map')
  16. while gui.running:
  17.     render_pixels()
  18.     arr = pixels.to_numpy()
  19.     gui.set_image(cmap(arr))
  20.     gui.show()

在 Interacting with external arrays 这一章节获得更多有关细节。