你好,世界!

我们将通过一个分形程序的例子来介绍 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
  2. import taichi as ti
  3. ti.init(arch=ti.gpu)
  4. n = 320
  5. pixels = ti.var(dt=ti.f32, shape=(n * 2, n))
  6. @ti.func
  7. def complex_sqr(z):
  8. return ti.Vector([z[0] ** 2 - z[1] ** 2, z[1] * z[0] * 2])
  9. @ti.kernel
  10. def paint(t: ti.f32):
  11. for i, j in pixels: # 对于所有像素,并行执行
  12. c = ti.Vector([-0.8, ti.sin(t) * 0.2])
  13. z = ti.Vector([float(i) / n - 1, float(j) / n - 0.5]) * 2
  14. iterations = 0
  15. while z.norm() < 20 and iterations < 50:
  16. z = complex_sqr(z) + c
  17. iterations += 1
  18. pixels[i, j] = 1 - iterations * 0.02
  19. gui = ti.GUI("Fractal", (n * 2, n))
  20. for i in range(1000000):
  21. paint(i * 0.03)
  22. gui.set_image(pixels)
  23. 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)
  3. # 在 GPU 上运行, 使用 NVIDIA CUDA 后端
  4. ti.init(arch=ti.cuda)
  5. # 在 GPU 上运行, 使用 OpenGL 后端
  6. ti.init(arch=ti.opengl)
  7. # 在 GPU 上运行, 使用苹果 Metal 后端(仅对 OS X)有效
  8. ti.init(arch=ti.metal)
  9. # 在 CPU 上运行 (默认)
  10. 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
  5. s = 0
  6. for j in range(5): # 在每个并行的线程中顺序执行
  7. s += j
  8. y[i] = s
  9. @ti.kernel
  10. def fill_3d():
  11. # 在区间 3 <= i < 8, 1 <= j < 6, 0 <= k < 9 上展开并行
  12. for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
  13. x[i, j, k] = i + j + k

注解

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

  1. @ti.kernel
  2. def foo():
  3. for i in range(10): # 并行 :-)
  4. @ti.kernel
  5. def bar(k: ti.i32):
  6. if k > 42:
  7. for i in range(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. @ti.kernel
  5. def bar(k: ti.i32):
  6. # 最外层作用域是 `if` 语句
  7. if k > 42:
  8. for i in x: # 语法错误。结构 for 循环 只能用于最外层作用域

警告

并行循环不支持 break 语句:

  1. @ti.kernel
  2. def foo():
  3. for i in x:
  4. ...
  5. break # 错误:并行执行的循环不能有 break
  6. @ti.kernel
  7. def foo():
  8. for i in x:
  9. for j in y:
  10. ...
  11. 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))
  3. pixels[42, 11] = 0.7 # store data into pixels
  4. 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))
  3. import numpy as np
  4. arr = np.random.rand(1024, 512)
  5. pixels.from_numpy(arr) # load numpy data into taichi tensors
  6. import matplotlib.pyplot as plt
  7. arr = pixels.to_numpy() # store taichi data into numpy arrays
  8. plt.imshow(arr)
  9. plt.show()
  10. import matplotlib.cm as cm
  11. cmap = cm.get_cmap('magma')
  12. gui = ti.GUI('Color map')
  13. while gui.running:
  14. render_pixels()
  15. arr = pixels.to_numpy()
  16. gui.set_image(cmap(arr))
  17. gui.show()

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