你好，世界！

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

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

# fractal.py

import taichi as ti

ti.init(arch=ti.gpu)

n = 320
pixels = ti.var(dt=ti.f32, 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: ti.f32):
  for i, j in pixels: # 对于所有像素，并行执行
    c = ti.Vector([-0.8, ti.sin(t) * 0.2])
    z = ti.Vector([float(i) / n - 1, float(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("Fractal", (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 中的领域特定语言(Domain-Specific Language, DSL )。为了使 Taichi 能像 Python 包一样易于使用，基于这个目标我们做了大量的工程工作——使得每个 Python 程序员能够以最低的学习成本编写 Taichi 程序。你甚至可以选择你最喜欢的 Python 包管理系统、Python IDE 以及其他 Python 包和 Taichi 一起结合使用。

可移植性

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

# 在 GPU 上运行，自动选择后端
ti.init(arch=ti.gpu)

# 在 GPU 上运行， 使用 NVIDIA CUDA 后端
ti.init(arch=ti.cuda)
# 在 GPU 上运行， 使用 OpenGL 后端
ti.init(arch=ti.opengl)
# 在 GPU 上运行， 使用苹果 Metal 后端（仅对 OS X）有效
ti.init(arch=ti.metal)

# 在 CPU 上运行 (默认)
ti.init(arch=ti.cpu)

注解

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

平台	CPU	CUDA	OpenGL	Metal
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 循环可以嵌套。

@ti.kernel
def fill():
  for i in range(10): # 并行执行
    x[i] += i

    s = 0
    for j in range(5): # 在每个并行的线程中顺序执行
      s += j

    y[i] = s

@ti.kernel
def fill_3d():
  # 在区间 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

注解

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

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

@ti.kernel
def bar(k: ti.i32):
    if k > 42:
        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 循环只能使用在内核的最外层作用域。

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

@ti.kernel
def foo():
    for i in x:
        …

@ti.kernel
def bar(k: ti.i32):
    # 最外层作用域是 `if` 语句
    if k > 42:
        for i in x: # 语法错误。结构 for 循环 只能用于最外层作用域
            …

警告

并行循环不支持 break 语句：

@ti.kernel
def foo():
  for i in x:
      ...
      break # 错误：并行执行的循环不能有 break

@ti.kernel
def foo():
  for i in x:
      for j in y:
          ...
          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:

import taichi as ti
pixels = ti.var(ti.f32, (1024, 512))

pixels[42, 11] = 0.7  # store data into pixels
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.:

import taichi as ti
pixels = ti.var(ti.f32, (1024, 512))

import numpy as np
arr = np.random.rand(1024, 512)
pixels.from_numpy(arr)   # load numpy data into taichi tensors

import matplotlib.pyplot as plt
arr = pixels.to_numpy()  # store taichi data into numpy arrays
plt.imshow(arr)
plt.show()

import matplotlib.cm as cm
cmap = 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()

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