跳转至主要内容
Version: v1.6.0

语法糖

别名

有时为命名繁琐的全局变量和函数创建别名可以提高可读性。 在 Taichi 中,这可以通过将 ti.static() 分配给 kernel 和函数本地变量实现,这会迫使 Taichi 使用标准的 python 指针分配。

例如,观察下面这个简单的内核:

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

fields 和函数可以通过 ti.static 重命名新的别名:

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

还可以为类成员和方法创建别名,这有助于避免出现面向对象编程代码塞满 self 的情况。

例如,考虑使用类内核来计算某个 field 的二维拉普拉斯算子:

@ti.kernel
def compute_laplacian(self):
for i, j in a:
self.b[i, j] = (self.a[i + 1, j] - 2.0*self.a[i, j] + self.a[i-1, j])/(self.dx**2) \
+ (self.a[i, j + 1] - 2.0*self.a[i, j] + self.a[i, j-1])/(self.dy**2)

使用 ti.static(),这可以简化为:

@ti.kernel
def compute_laplacian(self):
a, b, dx, dy = ti.static(self.a, self.b, self.dx, self.dy)
for i, j in a:
b[i, j] = (a[i+1, j] - 2.0*a[i, j] + a[i-1, j])/(dx**2) \
+ (a[i, j+1] - 2.0*a[i, j] + a[i, j-1])/(dy**2)
note

ti.static can also be used in combination with:

  • if(编译时分支)和
  • for(编译时展开)

更多详细信息请访问 元编程

在这里,我们将其用于编译时常量值,也就是说 field/函数句柄 在编译时是常量。

本文有帮助吗?