语法

内核

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

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

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

  1. @ti.kernel
  2. def add_xy(x: ti.f32, y: ti.f32) -> ti.i32:
  3.     return x + y  # 等价于： ti.cast(x + y, ti.i32)
  5. res = add_xy(2.3, 1.1)
  6. print(res)  # 3，因为返回值类型是 ti.i32

注解

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

  1. @ti.kernel
  2. def bad_kernel() -> ti.Matrix:
  3.     return ti.Matrix([[1, 0], [0, 1]])  # 错误
  5. @ti.kernel
  6. def bad_kernel() -> (ti.i32, ti.f32):
  7.     x = 1
  8.     y = 0.5
  9.     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 作用域内调用它们。

  1. @ti.func
  2. def laplacian(t, i, j):
  3.     return inv_dx2 * (
  4.         -4 * p[t, i, j] + p[t, i, j - 1] + p[t, i, j + 1] + p[t, i + 1, j] +
  5.         p[t, i - 1, j])
  7. @ti.kernel
  8. def fdtd(t: ti.i32):
  9.     for i in range(n_grid): # 并行
  10.         for j in range(n_grid): # 在每个并行的线程中分别顺序执行
  11.             laplacian_p = laplacian(t - 2, i, j)
  12.             laplacian_q = laplacian(t - 1, i, j)
  13.             p[t, i, j] = 2 * p[t - 1, i, j] + (
  14.                 c * c * dt * dt + c * alpha * dt) * laplacian_q - p[
  15.                            t - 2, i, j] - c * alpha * dt * laplacian_p

警告

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

  1. # 错误示范 - 两个返回语句
  2. @ti.func
  3. def safe_sqrt(x):
  4.   if x >= 0:
  5.     return ti.sqrt(x)
  6.   else:
  7.     return 0.0
  9. # 正确示范 - 一个返回语句
  10. @ti.func
  11. def safe_sqrt(x):
  12.   rst = 0.0
  13.   if x >= 0:
  14.     rst = ti.sqrt(x)
  15.   else:
  16.     rst = 0.0
  17.   return rst

警告

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

注解

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

注解

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

  1. @ti.func
  2. def sdf(u):  # functions support matrices and vectors as arguments. No type-hints needed.
  3.     return u.norm() - 1
  5. @ti.kernel
  6. def render(d_x: ti.f32, d_y: ti.f32):  # kernels do not support vector/matrix arguments yet. We have to use a workaround.
  7.     d = ti.Vector([d_x, d_y])
  8.     p = ti.Vector([0.0, 0.0])
  9.     t = sdf(p)
  10.     p += d * t
  11.     ...

标量算术

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 或 向量 上时，它们会被逐个作用到所有元素，例如：

  1. B = ti.Matrix([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  2. C = ti.Matrix([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]])
  4. A = ti.sin(B)
  5. # is equivalent to
  6. for i in ti.static(range(2)):
  7.     for j in ti.static(range(3)):
  8.         A[i, j] = ti.sin(B[i, j])
  10. A = ti.pow(B, 2)
  11. # is equivalent to
  12. for i in ti.static(range(2)):
  13.     for j in ti.static(range(3)):
  14.         A[i, j] = ti.pow(B[i, j], 2)
  16. A = ti.pow(B, C)
  17. # is equivalent to
  18. for i in ti.static(range(2)):
  19.     for j in ti.static(range(3)):
  20.         A[i, j] = ti.pow(B[i, j], C[i, j])
  22. A += 2
  23. # is equivalent to
  24. for i in ti.static(range(2)):
  25.     for j in ti.static(range(3)):
  26.         A[i, j] += 2
  28. A += B
  29. # is equivalent to
  30. for i in ti.static(range(2)):
  31.     for j in ti.static(range(3)):
  32.         A[i, j] += B[i, j]