内部设计（WIP）

中间表示（Intermediate representation）

使用 ti.init(print_ir=True) 来将中间表示代码输出到控制台。

表生成

Taichi 中的结构 for 循环会以 并行 的方式遍历一个稀疏数据结构中的所有活跃元素。 这把“在稀疏数据结构中均匀分配负载到处理器核心上”这一任务变得十分具有挑战性。具体来说，天真地把一个不规则树分片很容易产生数个叶节点数量严重不均衡的分区。

对此，我们的策略是循序渐进地对于每一层生成（对于该层）活跃的 SNode 元素的表。这个表的计算将发生在和计算正常计算内核的同一个设备上，并且具体取决于在用户调用 ti.init 函数时所提供的 arch 参数。

表的生成将会把数据结构的叶节点展平成一维的稠密数组，并因此规避不完整树的不规则性。然后，我们就可以直接在表上调用一个正常的 并行 for 循环 。

例如，

  1. # misc/listgen_demo.py
  3. import taichi as ti
  5. ti.init(print_ir=True)
  7. x = ti.var(ti.i32)
  8. ti.root.dense(ti.i, 4).bitmasked(ti.i, 4).place(x)
  10. @ti.kernel
  11. def func():
  12.     for i in x:
  13.         print(i)
  15. func()

以上的代码会生成下面的中间表示（IR）

  1. $0 = offloaded clear_list S1dense
  2. $1 = offloaded listgen S0root->S1dense
  3. $2 = offloaded clear_list S2bitmasked
  4. $3 = offloaded listgen S1dense->S2bitmasked
  5. $4 = offloaded struct_for(S2bitmasked) block_dim=0 {
  6.   <i32 x1> $5 = loop index 0
  7.   print i, $5
  8. }

请注意， func 的使用会生成以下两个表：

  • （任务 $0 和 $1）基于 root 节点 （S0）的表会生成一个关于 dense 节点们（S1）的表；
  • （任务 $2 和 $3）基于 dense 节点们（S1）的表会生成一个关于 bitmasked 节点们（S2）的表。

关于 root 节点的表总会有且仅有一个元素（实例），所以我们永远不会去清空或者重新生成这个表。

注解

关于 place （叶）节点的表 （比如说，在这个例子里它是 S3 )，永远不会被生成。相反，我们可以遍历关于这些节点的父节点们的表，并且于每个父节点，我们在不生成额外的表的情况下直接遍历所有 place 节点。

这种设计的初衷是去平摊生成表所带来的额外开销。因为去对于每个叶节点（place SNode）生成一个表元素会带来过多的开销，并且这些开销极有可能大大超过在叶元素本身上进行的必要的计算。所以，我们选择只生成和这些叶节点的父节点相关的的元素表，这样就能把生成表所带来的开销平摊到多个倒数第二层的 SNode 元素的子元素上。

在上面的例子中，虽然我们有 16 个关于 x 的实例，但是我们只生成了 4 个 bitmasked 节点（和 1 个 dense 节点）。

代码生成

统计量

在某些情况下，在Taichi程序的执行过程中，收集关于内部事件的特定的量化信息是很用帮助的。Statistics 类就是为此设计的。

用法：

  1. #include “taichi/util/statistics.h”
  3. 将1.0加到计数器“codegenoffloaded_tasks”上
  4. taichi::stat.add(“codegen_offloaded_tasks”);
  6. // 将“中间表示”中语句的数量加到计数器“codegen_statements”上
  7. taichi::stat.add(“codegen_statements”, irpass::analysis::count_statements(this->ir));

注意键为 std::string 而值类型为 double。

在Python中使用如下方式来打印出所有的统计量：

  1. ti.core.print_stat()

为什么使用Python作为前端语言

将Taichi嵌入到 Python 中有以下优点：

  • 易于学习。Taichi的语法与Python非常相似。
  • 易于运行。不需要运行前编译（ahead-of-time compilation）。
  • 这样的设计使用户可以重复利用已有的Python基础架构：
    • 集成开发环境（IDEs）。大部分Python的集成开发环境提供的语法高亮，语法检查和自动补全功能可以用于Taichi。
    • 包管理器（pip）。开发好的Taichi程序可以被简单地提交至 PyPI 并且其他用户可以轻松地使用 pip 安装它。
    • 现有的包。用户可以很轻松地与其他Python组件（例如 matplotlib 和 numpy）交互。
  • 只要内核主体可以被Python的解析器解析，那么 Python 内置的处理抽象语法树（AST）的工具让我们可以做一些奇妙的事情。

但是，这样的设计同样存在一些不足之处：

  • Taichi内核必须能被Python解析器解析。这意味着Taichi的语法不能超出Python的语法范畴。
    • 例如，访问Taichi张量时，即使张量是0维度的也必须使用索引。如果 x 是0维的，需要使用 x[None] = 123 来给 x 中的量赋值。这是因为在Python语法中， x = 123 将会将 x 本身（而不是它包含的值）设为常数 123 ，不幸的是，我们无法更改这种行为。
  • Python的性能相对较为低下。这在使用纯Python脚本初始化较大Taichi张量时会导致一些性能问题。初始化较大张量时必须使用Taichi内核。

Virtual indices v.s. physical indices

In Taichi, virtual indices are used to locate elements in tensors, and physical indices are used to specify data layouts in memory.

例如，

  • In a[i, j, k], i, j, and k are virtual indices.
  • In for i, j in x:, i and j are virtual indices.
  • ti.i, ti.j, ti.k, ti.l, ... are physical indices.
  • In struct-for statements, LoopIndexStmt::index is a physical index.

The mapping between virtual indices and physical indices for each SNode is stored in SNode::physical_index_position. I.e., physical_index_position[i] answers the question: which physical index does the i-th virtual index correspond to?

Each SNode can have a different virtual-to-physical mapping. physical_index_position[i] == -1 means the i-th virtual index does not corrspond to any physical index in this SNode.

SNode s in handy dense tensors (i.e., a = ti.var(ti.i32, shape=(128, 256, 512))) have trivial virtual-to-physical mapping, e.g. physical_index_position[i] = i.

However, more complex data layouts, such as column-major 2D tensors can lead to SNodes with physical_index_position[0] = 1 and physical_index_position[1] = 0.

  1. a = ti.var(ti.f32, shape=(128, 32, 8))
  3. b = ti.var(ti.f32)
  4. ti.root.dense(ti.j, 32).dense(ti.i, 16).place(b)
  6. ti.get_runtime().materialize()
  8. mapping_a = a.snode().physical_index_position()
  10. assert mapping_a == {0: 0, 1: 1, 2: 2}
  12. mapping_b = b.snode().physical_index_position()
  14. assert mapping_b == {0: 1, 1: 0}
  15. # Note that b is column-major:
  16. # the virtual first index exposed to the user comes second in memory layout.

Taichi supports up to 8 (constexpr int taichi_max_num_indices = 8) virtual indices and physical indices.