2022-12-10发表2023-12-16更新Taichi10 分钟读完 (大约1517个字)0次访问

课程笔记：核函数执行流程 | Taichi

本文是 Taichi 官方教学视频的笔记：

Overview

下面是一个简单的 Taichi 程序的示例：

数据容器的编译和指令的编译是分开的：

Kernel Compilation

在编译 kernel 时，经历了以下步骤：

这个过程和其他的一些使用 Python 作为前端的 JIT/AOT DSL 差不多，都是将 Python AST 转换为自己的 IR 然后进行编译器后端优化。

接下来以这样一个简化的 Taichi kernel 为例看一下 kernel 具体是如何完成编译的：

@ti.kernel
def test(arg0: int, arg1: int):
    tmp = 0

    if arg1:
        tmp = tmp + 4

    for i in range(10):
        tmp += i

test(10, True)

在进行 kernel 函数调用的时候，就开始了编译：

class Kernel:
    def __call__(self):
        self.ensure_compiled

    def ensure_compiled(self):
        self.materialize()

    def materialize(self):
        # Field compilation
        self.runtime.materialize()

        # Python AST -> Frontend IR
        transform_tree(ast_tree, ast_builder)

        # Frontend IR -> CHI IR -> CodeGen
        k = prog.create_kernel(ast_builder)

        # Prepare Runtime Executable
        self.get_function_body(k)

Python AST -> Frontend IR

python/taichi/lang/ast/ast_transformer.py

Frontend IR 的层级结构如下：

taichi/ir/frontendir.h

Frontend IR -> 浅层 CHI IR

Frontend IR 基本上只承担了一个承接的作用，接下来需要 lowering 成 CHI IR 才能做进一步的优化。逻辑上可以将 CHI IR 分为浅层和深层。对于到浅层 CHI IR 的转换，只是做了语义上的对等转换：

taichi/transforms/lower_ast.cpp

taichi/ir/statements.h

taichi/transforms/compile_to_offloads.cpp

上图提到的 offload 是 CHI IR 浅层和深层的分界线。

浅层 CHI IR -> 深层 CHI IR

taichi/transforms/offload.cpp

offload 主要是对 kernel 做了进一步的划分。以上图为例，图中红色蓝色需要相互依赖，是串行结构，被 offload 成 serial。绿色部分是循环结构，的通常情况下是可以展开的，被 offload 成 serial range_for，表示是可以进行并行的。

offload 之后，把 IR 结构做了一些调整，kernel 层级下面不再是 stmt，而是 offload，offload 下面才是 stmt：

taichi/ir/statements.h

CHI IR -> CodeGen

在进行 codegen 时，CHI IR 会被翻译成功能对等的 LLVM IR Builder API，从而完成代码生成：

taichi/codegen/llvm/codegen_lvm.h

Field Compilation

ti.field 的编译和 ti.kernel 编译的过程类似，每个阶段可以大概对应起来：

这里插入介绍下 SNode。SNode 是 Taichi 提供的通用的数据容器，可以方便地构造或稠密或稀疏的多维 field，并形成清晰的层级。注意区分三个概念：SNode container、SNode cell 和 SNode component。

一个 SNode container 可以包含多个 SNode cell。cell 的数量推荐设为 2 的幂次。举个例子，S = ti.root.dense(ti.i, 128) 创建了一个 SNode S，每个 S container 包含 128 个 S cell。
一个 SNode cell 可以包含多个 SNode components。举个例子，P = S.dense(ti.i, 4); Q = S.dense(ti.i, 4) 插入两个 components（一个 P container 和一个 Q container）到每个 S cell。
请注意，每个 SNode component 都是较低级别 SNode 的 SNode container。

Taichi 中的分层数据结构，无论是稠密的还是稀疏的，本质上都是一棵具有交错的 container 和 cell 级别的树。更加详细的内容可以参考官方文档。

这样 Python 中的 ti.field 就转换成了 C++ 中浅层的数据结构。

Access Fields in Kernel

上图将前面提到的两个编译步骤对应了起来。接下来以一段代码为例，看一下示例中的 x[3] = 100 是怎么转换下去的：

x = ti.field(ti.i32)
y = ti.field(ti.f32)
common = root.dense(2)
common.dense(2).place(y)
common.dense(2).place(x)

@ti.kernel
def test0():
    x[3] = 100