中文文档：《Taichi编程语言》

1. 初始化 Taichi

1	ti.init(arch = ti.cuda)

推荐使用：

1	ti.init(arch = ti.gpu)

会自动识别 cuda / metal / opengl

2. 数据类型

# Signed integers: 
ti.i8/i16/i32/i64

# Unsigned integers: 
ti.u8/u16/u32/u64

# Float-point numbers: 
ti.f32/f64

NOTE：

布尔类型暂时使用 ti.i32
opengl 不支持64位的整型和浮点型，详情查看文档

3. Tensors（张量）

Taichi 中的张量 Tensors 本质是一个多维数组
一个张量元素可以是：
- 零维张量 var ：标量形式
- 一维张量 ti.Vector ：向量形式
- 二维张量 ti.Matrix ：矩阵形式
张量元素总是通过 a[i，j，k] 语法访问。
Taichi 中的张量和矩阵是两个概念
张量可以是空间稀疏的

import taichi as ti

ti.init()
a = ti.var(dt=ti.f32, shape=(42, 63))  
# A tensor of 42x63 scalars
# 标量的张量，其中 42x63 个元素

b = ti.Vector(3, dt=ti.f32, shape=4)  
# A tensor of 4x 3D vectors
# 4个三维向量组成的张量
# 张量的元素是三维向量
# Vector(向量的维度, dt = 数据类型, shape = 张量的大小：向量的个数)  

C = ti.Matrix(2, 2, dt=ti.f32, shape=(3, 5))
# A tensor of 3x5 2x2 matrices
# 张量的元素是 2*2 的矩阵
# 张量的形状是 3*5 ，其中的每个元素是 2*2 的矩阵

loss = ti.var(dt=ti.f32, shape=())
# A (0−D) tensor of a single scalar
# 零维张量：只有单个元素的张量，且该元素是个标量

a[3, 4] = 1
print('a[3,4] =', a[3, 4])
# ”a[3,4] = 1.000000”

b[2] = [6, 7, 8]
print(' b[0] =', b[0][0], b[0][1], b[0][2])
# print (b[0]) is not yet supported .

loss[None] = 3
print(loss[None])  
# 3
# 在访问 loss 的时候可以将下标设置为 None
# 注意不要直接 loss = 3 ，这样变量 loss 的类型将不再是 ti.var

4. Kernels

Taichi 中的 Kernels 是用来计算的函数，用法类似于 python 中的 Functions
Taichi 中的 Kernels 有编译过程，提高代码运行速度
静态类型，在GPU上并行运行，且可微分
使用 Kernels 必须装饰 @ti.kernel
Kernel 的参数和返回值必须声明类型（Taichi 本身是强类型语言）

@ti.kernel
def hello(i: ti.i32):
    a = 40
    print('Hello world!', a + i)

hello(2) # "Hello world! 42"

@ti.kernel
def calc() -> ti.i32:
    s = 0
    for i in range(10):
        s += i
    return s # 45

5. Functions

Taichi 中的 Functions 可以被 Taichi 中的 Kernels 调用
Taichi 中的 Kernels 可以被 Taichi 中的 Functions 调用
Taichi 中的 Functions 不能被 python 中的 Functions 调用
Taichi 中的 Kernels 可以被 python 中的 Functions 调用
Taichi 中的 Functions 可以被 Taichi 中的 Functions 调用
Taichi 中的 Kernels 不能被 Taichi 中的 Kernels 调用

Note ：
编写代码的过程中尽量遵循一个原则：

将 ti.func 作为所有数值处理的子块，其内联的特性也将提高其运算性能；
将 ti.kernel 仅作为向外与 pyhton 生态交互的接口，其作用是分布 ti.func 任务，同时作为程序的外壳暴露给诸如渲染循环等结构。

@ti.func
def triple(x):
    return x * 3

@ti.kernel
def triple_array:
    for i in range(128):
        a[i] = triple(a[i])

NOTE：

Taichi 的 Functions 将是强制 inline 的。目前，不允许递归。
Taichi 的 Functions 最多只能包含一个 return 语句。

6. 标准数据库

ti.sin(x)
ti.cos(x)
ti.asin(x)
ti.acos(x)
ti.atan2(y, x)

ti.cast(x, data_type)
ti.sqrt(x)

ti.floor(x)
ti.ceil(x)
ti.inv(x)

ti.tan(x)
ti.tanh(x)
ti.exp(x)
ti.log(x)

ti.random(data_type)
# 传入数据类型，返回数据类型中的一个数

# python 的符号
a + b
a / b # 结果为浮点数
a // b # 结果为整型

# python 自带的库
abs(x)
int(x)
float(x)
max(x, y, ...)
min(x, y, ...)

x ** y
# x^y

7. 矩阵与线性代数

ti.Matrix 只能用于小矩阵（例如 3×3）
如果使用 64×64 数组，建议使用一个标量的 2D 张量：ti.var(dt=ti.f32, shape=(64, 64))
ti.Vector 向量只能用于一列的元素
区分 Hadamard 积 * （也称对应元素乘积，A 和 B 对应元素相乘获得新的矩阵）和矩阵积 @ （矩阵乘法）

以下方法是返回一个新的矩阵：

A.transpose()
# 矩阵的转置

A.inverse()
# 矩阵的逆

A.trace()
# 矩阵的迹

A.determinant(type)
# 矩阵的行列式（需要设置数据类型）

A.normalized()
# 矩阵标准化

A.cast(type)
# 重置矩阵中的数据类型

R, S = ti.polar_decompose(A, ti.f32)
# 矩阵的极分解

U, sigma , V = ti.svd(A, ti.f32)
# 矩阵的奇异值分解
# sigma 是对角矩阵

ti.sin(A)/cos(A)... (element -wise)
# 将矩阵中的每个元素进行 sin 计算
# 所有的标量运算符均可以在矩阵上运算

8. 并行 for 循环

因为 Taichi 中的循环有两种形式

Range-for loops ：
- 类似于 python 中的 for 循环；
- 只是在最外层的范围内使用它时会被并行化；
- 循环内并行执行的代码块顺序是不确定的；
- 循环的范围可以嵌套。
Struct-for loops ：迭代（稀疏）张量元素。

8.1 Range-for loops

在 Taichi 的 kernel 中最外层作用域的 For 循环被自动并行化。

@ti.kernel
def fill():
    for i in range(10): # 并行执行，0 <= i < 4
        x[i] += i
    
        s = 0
        for j in range(5): # 在最外层for循环的每个线程中串行执行
            s += j

        y[i] = s

@ti.kernel
def fill_3d():
    # Parallelized for all 3 <= i < 8, 1 <= j < 6, 0 <= k < 9
    # ti.ndrange 可以定义上下限
    for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
        x[i, j, k] = i + j + k

NOTE ：只有最外层范围的循环可以并行化；但循环若不是在最外层（例如最外层还有 if 语句），则无法自动并行。

@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 :-(
            ...

8.2 Struct-for loops

import taichi as ti

ti.init(arch=ti.gpu)

n = 320
pixels = ti.var(dt=ti.f32, shape=(n * 2, n))

@ti.kernel
def paint(t: ti.f32):
    for i, j in pixels: # Parallized over all pixels
        pixels[i, j] = i * 0.001 + j * 0.002 + t

paint(0.3)

Struct-for loops 遍历所有张量坐标，即：

(0, 0), (0, 1), (0, 2), …, (0, 319)
(1, 0), …
…
…, (639, 319).

9. 原子操作

在 Taichi 中，自赋值（例如 x[i]+=1 ）是自动原子化的
在并行的修改全局变量时，请确保使用原子操作。（例如，统计x中的所有元素）

@ti.kernel
def sum():
    for i in x:
        # Approach 1: OK
        total[None] += x[i]

        # Approach 2: OK
        ti.atomic_add(total[None], x[i])

        # Approach 3: Wrong result (the operation is not atomic.)
        # 存在非原子的 读加写 的操作
        total[None] = total[None] + x[i]

10. Taichi-scope v.s. python-scope

Taichi-scope ：任何被 ti.kernel 和 ti.func 修饰的函数的函数体。
python-scope ：任何不在 Taichi 里面的函数体。

区别：
在 Taichi-scope 里的代码会被 Taichi 编译，并且在并行设备上运行。
在 python-scope 的代码不会被编译，仅仅只会被 python 解释器运行。

import taichi as ti
ti.init ()

a = ti . var (dt=ti. f32, shape =(42, 63)) # A tensor of 42x63 scalars
b = ti . Vector (3 , dt=ti . f32 , shape=4) # A tensor of 4x 3D vectors
C = ti . Matrix (2 , 2 , dt=ti.f32 , shape =(3, 5) ) # A tensor of 3x5 2x2 matrices

@ti.kernel
def foo () :
    a[3,4] = 1
    print ('a[3,4] =',a[3,4])
    # ”a[3,4] = 1.000000”
    
    b[2] = [6,7,8]
    print('b[0]=',b[0],',b[2] =',b[2])
    # ”b[0] = [[0.000000],[0.000000],[0.000000]],b[2]=[[6.000000],[7.000000],[8.000000]]”
    
    C[2,1][0,1] = 1
    print('C[2,1] =',C[2,1])
    # C[ 2 , 1] = [[0.000000 , 1.000000] , [0.000000 , 0.000000]]
    
foo()

11. 一个 Taichi 程序的几个阶段

初始化: ti.init(...)
Tensor 定义（到计算阶段就无法继续定义）: ti.var, ti.Vector, ti.Matrix
计算阶段 (启动 kernels)
可选: 重新启动太极系统（清除内存，销毁所有变量和 kernels ）：ti.reset()

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: # Parallized 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()

12. Debug 模式

在调试模式下初始化Taichi，启用绑定检查程序（仅限CPU）：ti.init(debug=True, arch=ti.cpu)

import taichi as ti
ti.init(debug=True , arch=ti.cpu)

a = ti.var(ti.i32, shape=(10))
b = ti.var(ti.i32, shape=(10))

@ti.kernel
def shift():
    for i in range(10):
        a[i] = b[i + 1] 
        # 会出现越界的情况，仅仅在debug 的情况下检测出越界
        # Runtime error in debug mode

shift()

12.1 关闭优化

关闭优化，让编译速度变得更快

1	ti.core.toggle_advanced_optimization(False)

13. ODOP

在面向数据编程中加入面向对象特性

Demo ：

1	ti example odop_solar

\bold{a} = GM\bold{r} / \Vert \bold{r} \Vert_2^3

import taichi as ti
import math

ti.init()


@ti.data_oriented
class SolarSystem:
    def __init__(self, n, dt):  # Initializer of the solar system simulator
        self.n = n
        self.dt = dt
        self.x = ti.Vector.field(2, dtype=ti.f32, shape=n)
        self.v = ti.Vector.field(2, dtype=ti.f32, shape=n)
        self.center = ti.Vector.field(2, dtype=ti.f32, shape=())

    @staticmethod
    @ti.func
    def random_vector(radius):  # Create a random vector in circle
        theta = ti.random() * 2 * math.pi
        r = ti.random() * radius
        return r * ti.Vector([ti.cos(theta), ti.sin(theta)])

    @ti.kernel
    def initialize_particles(self):
        # (Re)initialize particle position/velocities
        for i in range(self.n):
            offset = self.random_vector(0.5)
            self.x[i] = self.center[None] + offset  # Offset from center
            self.v[i] = [-offset.y, offset.x]  # Perpendicular to offset
            self.v[i] += self.random_vector(0.02)  # Random velocity noise
            self.v[i] *= 1 / offset.norm()**1.5  # Kepler's third law

    @ti.func
    def gravity(self, pos):  # Compute gravity at pos
        offset = -(pos - self.center[None])
        return offset / offset.norm()**3

    @ti.kernel
    def integrate(self):  # Semi-implicit Euler time integration
        for i in range(self.n):
            self.v[i] += self.dt * self.gravity(self.x[i])
            self.x[i] += self.dt * self.v[i]

    def render(self, gui):  # Render the scene on GUI
        gui.circle([0.5, 0.5], radius=10, color=0xffaa88)
        gui.circles(solar.x.to_numpy(), radius=3, color=0xffffff)


solar = SolarSystem(8, 0.0001)
solar.center[None] = [0.5, 0.5]
solar.initialize_particles()

gui = ti.GUI("Solar System", background_color=0x0071a)
while gui.running:
    if gui.get_event() and gui.is_pressed(gui.SPACE):
        solar.initialize_particles()  # reinitialize when space bar pressed.

    for i in range(10):  # Time integration
        solar.integrate()

    solar.render(gui)
    gui.show()

14. Templates

14.1 template 元编程

@ti.kernel
def offset(x: ti.template(), y: ti.template(), c: ti.f32):
    for i in x:
        y[i] = x[i] + c

当使用 ti.template() 时，每次调用该 kernel 都会实例化一个对应的 Kernel

import taichi as ti
ti.init()

@ti.kernel
def hello(i: ti.template()):
    print(i)

for i in range(100):
    hello(i) # 100 different kernels will be created

@ti.kernel
def world(i: ti.i32):
        print(i)

for i in range(100):
    world(i) # The only instance will be reused现

复制不同阶的 Tensor ：

@ti.kernel
def copy(x: ti.template(), y: ti.template()):
    for I in ti.grouped(y):
        x[I] = y[I]n

NOTE ：

其中使用 ti.grouped(y) 对张量 y 中的所有元素进行了一个打包操作

@ti.kernel
def array_op(x: ti.template(), y: ti.template()):
    for I in ti.grouped(x):
        # I is a vector of size x.dim() and data type i32
        y[I] = I[0] + I[1]

如果 x 是二维张量，代码为：

@ti.kernel
def array_op(x: ti.template(), y: ti.template()):
    for i, j in x:
        y[i, j] = i + j

14.2 template 的反射

可在编译器获得 Tensor 大小的反射：

编译时可以获得 template 的大小

import taichi as ti

tensor = ti.var(ti.f32, shape=(4, 8, 16, 32, 64))

@ti.kernel
def print_tensor_size(x: ti.template()):
    print(x.dim())
    for i in ti.static(range(x.dim())):
        print(x.shape()[i])

print_tensor_size(tensor)

14.3 编译期的分支语句

使用 ti.static(value) 让分支语句在编译期执行：

enable_projection = True

@ti.kernel
def static():
    if ti.static(enable_projection): # No runtime overhead
        x[0] = 1

14.4 编译期的循环展开

使用 ti.static(range(...)) 进行循环展开：

import taichi as ti

ti.init()
x = ti.Vector(3, dt=ti.i32, shape=16)

@ti.kernel
def fill():
    for i in x:
        for j in ti.static(range(3)):
            x[i][j] = j
        print(x[i])

fill()

在向量/矩阵元素上循环，Taichi 的矩阵的索引必须是编译期常数，太极张量的下标可以是运行时变量：
eg：如果x是三维向量的一维张量，则作为 x[tensor_index][matrix_index] 访问。第一个索引可以是变量，而第二个索引必须是常量
因此在循环张量时需要对其进行 unroll 操作

15. 变量别名

@ti.kernel
def my_kernel():
    for i, j in tensor_a:
        tensor_b[i, j] = some_function(tensor_a[i, j])

使用 ti.static() 对变量或者函数进行别名操作

@ti.kernel
def my_kernel():
    a, b, fun = ti.static(tensor_a , tensor_b , some_function)
    for i,j in a:
        b[i,j] = fun(a[i,j])

16. 可微编程

正向编程求 $f(\bold{x})$ ，可微编程求 $\cfrac{\partial f(\bold{x})}{\partial \bold{x}}$

Taichi 支持反向模式自动微分（AutoDiff）
可以计算出反向传播 w.r.t.标量（损失）函数 $f(\bold{x})$ 的导数。

计算梯度的两种方式：

用 ti.Tape(loss) ，用于正向和梯度评估
显式使用 gradient kernels 进行梯度计算，运用上更多控制。

16.1 基于梯度的优化

min_{\bold{x}} \space \space L(\bold{x}) = \cfrac{1}{2} \sum_{i = 0}^{n-1}(\bold{x}_i - \bold{y}_i)^2

16.1.1 指定参与梯度运算的变量

声明张量时分配梯度：

1	x = ti.var(dt=ti.f32, shape=n, needs_grad=True)

只有 needs_grad=True 才能求 L 关于 x 的导数
自变量和因变量都需要激活

16.1.2 定义损失函数

定义损失函数 kernel(s) ：

@ti.kernel
def reduce():
    for i in range(n):
        L[None] += 0.5 * (x[i] - y[i])**2

16.1.3 调用梯度值

定义梯度下降（在该处调用 x.grad[i] ）：

@ti.kernel
def gradient_descent():
    for i in x:
        x[i] -= x.grad[i] * 0.1

也可用于时间积分（将 x.grad[i] 作为差分的数值）：

@ti.kernel
def substep():
        # Collide
        ......
        # 显式时间
        v[p] = (v[p] + ((-x.grad[p] / node_mass) + ti.Vector([0, -10, 0])) * dt) * math.exp(dt * -6)
        x[p] += dt * v[p]

Note ：

务必记得声明为 ti.kernel

16.1.4 使用微分器

对损失函数求导：ti.Tape(loss=L): reduce()

# 使用100次梯度下降迭代进行优化
for k in range(100):
    with ti.Tape(loss=L):
        reduce()
    print('Loss =', L[None])
    gradient_descent()

16.1.5 代码示例

参考示例 ：

1 2	ti example autodiff_minimization ti example autodiff_regression

NOTE ：

可以使用 rich 包：

1	pip install rich

这时可以使用语法高亮

1	ti example -p autodiff_minimization

可以将结果保存在当前目录下；

1	ti example -s autodiff_minimization

autodiff_minimization.py ：

import taichi as ti
import random

n = 8
x = ti.field(dtype=ti.f32, shape=n, needs_grad=True)
y = ti.field(dtype=ti.f32, shape=n)
L = ti.field(dtype=ti.f32, shape=(), needs_grad=True)


@ti.kernel
def reduce():
    for i in range(n):
        L[None] += 0.5 * (x[i] - y[i])**2


# Initialize vectors
for i in range(n):
    x[i] = random.random()
    y[i] = random.random()


@ti.kernel
def gradient_descent():
    for i in x:
        x[i] -= x.grad[i] * 0.1


# 使用100次梯度下降迭代进行优化
for k in range(100):
    with ti.Tape(loss=L):
        reduce()
    print('Loss =', L[None])
    gradient_descent()

for i in range(n):
    # 现在你应该大约已经有 x[i] == y[i]
    print(x[i], y[i])

损失函数：

L = \cfrac{1}{2} (x_i - y_i)^2

步长：0.1

16.2 自动微分的物理应用

16.2.1 应用1：势能梯度产生的力

势能对位移进行求导可以获得力：

\bold{f}_i = -\cfrac{\partial U(x)}{\partial \bold{x}_i}

手动导出梯度非常困难，使用 AutoDiff ：

分配一个零维张量来存储势能： potential = ti.var(ti.f32, shape=())
定义从 x[i] 计算势能的正向核函数。
在 ti.Tape(loss=potential) ，调用正向核函数。
每个粒子上的力是 -x.grad[i] 。

16.2.2 应用2：对整个物理过程进行求导

Demo ：DiffTaichi $(\bold{x}_{t+1}, \bold{v}_{t+1}, \cdots) = \bold{F} (\bold{x}_{t}, \bold{v}_{t}, \cdots)$

for k in range(100):
    with ti.Tape(loss=loss):
        for i in range(steps - 1):
            simulate(i)

进行100次迭代，

在进行端到端的微分时，始终保持时间步长的整个计算存活周期。

即不只分配：ti.Vector(3, dt=ti.f32, shape=(num_particles)) 存储最新粒子；
为整个模拟过程分配： ti.Vector(3, dt=ti.f32, shape=(num_timesteps, num_particles)) 。不要覆盖！

（使用检查点（本课程后面部分）来减少内存消耗。）

迭代过程中 x[i] -= x.grad[i] * alpha ，其中 alpha 是步长，如果使用最速下降或者共轭梯度等方法，步长则需要在迭代过程中进行计算。

17. 可视化

17.1 2D

《Taichi 中文文档》 GUI 系统

1	gui = ti.GUI("Taichi MLS-MPM-128", res=512, background_color=0x112F41)

1	gui.circle/gui.circles(x.to_numpy(), radius=1.5, color=colors.to_numpy())

1	gui.line/triangle/set_image/show/...

17.2 3D

《Taichi 中文文档》导出 PLY 文件

使用 ti.PLYWriter 导出三维粒子和网格

1	ti example export_ply/export_mesh

使用 Houdini / Blender 查看 3D 结果。

镜之国的爱丽丝

00 - Taichi 编程