语法¶

内核¶

内核参数必须有类型提示。内核最多只能有 8 个参数，例如，

@ti.kernel
def print_xy(x: ti.i32, y: ti.f32):
    print(x + y)

一个内核可以有一个标量返回值。如果内核有一个返回值，那它必须有类型提示。这个返回值会自动转换到所提示的类型。例如，

@ti.kernel
def add_xy(x: ti.f32, y: ti.f32) -> ti.i32:
    return x + y  # 等价于： ti.cast(x + y, ti.i32)

res = add_xy(2.3, 1.1)
print(res)  # 3，因为返回值类型是 ti.i32

注解

目前，我们只支持返回一个标量。返回 ti.Matrix 或者 ti.Vector 是不支持的。Python 方式的返回元组也是不支持的。例如：

@ti.kernel
def bad_kernel() -> ti.Matrix:
    return ti.Matrix([[1, 0], [0, 1]])  # 错误

@ti.kernel
def bad_kernel() -> (ti.i32, ti.f32):
    x = 1
    y = 0.5
    return x, y  # 错误

在 Taichi 内核中，我们也支持 模板参数（template arguments） (参见 Template metaprogramming) 和 外部数组参数（external array arguments） (参见 Interacting with external arrays)

警告

当使用可微编程时，对内核数据结构有一些约定。参见 Differentiable programming (WIP) 中的 内核简化规则（Kernel Simplicity Rule） 。

请不要在可微编程中使用内核返回值，因为这种返回值并不会被自动微分追踪。取而代之，可以把结果存入全局变量（例如 loss[None]）。

函数¶

使用 @ti.func 来装饰您的 Taichi 函数。这些函数只能在 Taichi 作用域内调用。不要在 Python 作用域内调用它们。

@ti.func
def laplacian(t, i, j):
    return inv_dx2 * (
        -4 * p[t, i, j] + p[t, i, j - 1] + p[t, i, j + 1] + p[t, i + 1, j] +
        p[t, i - 1, j])

@ti.kernel
def fdtd(t: ti.i32):
    for i in range(n_grid): # 并行
        for j in range(n_grid): # 在每个并行的线程中分别顺序执行
            laplacian_p = laplacian(t - 2, i, j)
            laplacian_q = laplacian(t - 1, i, j)
            p[t, i, j] = 2 * p[t - 1, i, j] + (
                c * c * dt * dt + c * alpha * dt) * laplacian_q - p[
                           t - 2, i, j] - c * alpha * dt * laplacian_p

警告

目前不支持具有多个 return 语句的函数。请用 局部变量 暂存结果，以便最终只有一个 return 语句：

# 错误示范 - 两个返回语句
@ti.func
def safe_sqrt(x):
  if x >= 0:
    return ti.sqrt(x)
  else:
    return 0.0

# 正确示范 - 一个返回语句
@ti.func
def safe_sqrt(x):
  rst = 0.0
  if x >= 0:
    rst = ti.sqrt(x)
  else:
    rst = 0.0
  return rst

警告

目前，所有函数都是强制内联的。因此，不能使用递归。

注解

函数的参数是以值传递的。

注解

Unlike functions, kernels do not support vectors or matrices as arguments:

@ti.func
def sdf(u):  # functions support matrices and vectors as arguments. No type-hints needed.
    return u.norm() - 1

@ti.kernel
def render(d_x: ti.f32, d_y: ti.f32):  # kernels do not support vector/matrix arguments yet. We have to use a workaround.
    d = ti.Vector([d_x, d_y])
    p = ti.Vector([0.0, 0.0])
    t = sdf(p)
    p += d * t
    ...

标量算术¶

Taichi 支持的标量函数：

ti.sin(x)¶

ti.cos(x)¶

ti.asin(x)¶

ti.acos(x)¶

ti.atan2(x, y)¶

ti.cast(x, data_type)¶

ti.sqrt(x)¶

ti.rsqrt(x)¶

ti.floor(x)¶

ti.ceil(x)¶

ti.tan(x)¶

ti.tanh(x)¶

ti.exp(x)¶

ti.log(x)¶

ti.random(data_type)¶

abs(x)¶

int(x)¶

float(x)¶

max(x, y)¶

min(x, y)¶

pow(x, y)¶

注解

Python 3 中 / （浮点数除法）和 // （整数除法）是区分开来的。例如，1.0 / 2.0 = 0.5，1 / 2 = 0.5，1 // 2 = 0，4.2 // 2 = 2。Taichi 也遵循了这个设计：

true divisions on integral types will first cast their operands to the default float point type.

floor divisions on float-point types will first cast their operands to the default integer type.

为避免这样的隐式转换，你可以手动使用 ti.cast 将你的操作数转换为你需要的类型。参见默认精度获取数字类型的更多细节。

注解

当这些标量函数被作用在 Matrices 或向量上时，它们会被逐个作用到所有元素，例如：

B = ti.Matrix([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
C = ti.Matrix([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]])

A = ti.sin(B)
# is equivalent to
for i in ti.static(range(2)):
    for j in ti.static(range(3)):
        A[i, j] = ti.sin(B[i, j])

A = ti.pow(B, 2)
# is equivalent to
for i in ti.static(range(2)):
    for j in ti.static(range(3)):
        A[i, j] = ti.pow(B[i, j], 2)

A = ti.pow(B, C)
# is equivalent to
for i in ti.static(range(2)):
    for j in ti.static(range(3)):
        A[i, j] = ti.pow(B[i, j], C[i, j])

A += 2
# is equivalent to
for i in ti.static(range(2)):
    for j in ti.static(range(3)):
        A[i, j] += 2

A += B
# is equivalent to
for i in ti.static(range(2)):
    for j in ti.static(range(3)):
        A[i, j] += B[i, j]