语法糖

别名

Creating aliases for global variables and functions with cumbersome names can sometimes improve readability. In Taichi, this can be done by assigning kernel and function local variables with ti.static(), which forces Taichi to use standard python pointer assignment.

例如,考虑下面这个简单的内核:

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

张量和函数使用 ti.static 别名为新名称:

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

还可以为类成员和方法创建别名,这有助于防止含有 self 的面向对象编程代码混乱。

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

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

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

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

注解

ti.static 还可与 if``(编译时分支)和 ``for (编译时展开)结合使用。 更多相关详细信息,请参见 Metaprogramming

在这里,我们将其用于 编译时常量值 ,即 张量/函数句柄 在编译时是常量。