Taichi RFC 解读：AOT 支持所有 SNode——SNode 树类型化与字段本地化的设计之路

Taichi RFC 解读AOT 支持所有 SNode——SNode 树类型化与字段本地化的设计之路【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi导读本文基于仓库中的设计文档 docs/rfcs/20220413-aot-for-all-snode.mdTaichi 官方 RFC作者 Ye Kuang2022-04-13系统解读 Taichi 如何通过SNode 树类型SNodeTree type这一抽象让 AOTAhead-of-Time编译支持任意类型的 SNode 与 Taichi 字段从而把全局变量式的 Taichi 字段改造为可显式传入 kernel 的局部化实体。读完本文你将理解为什么ti.field()的全局化实现会成为 AOT 部署的瓶颈RFC 提出的SNodeTreeBuilder仓库中落地为FieldsBuilder如何实现类型构建与实例化解耦shape、AoS/SoA、梯度字段、Python/C AOT API 分别如何设计以及这套设计与仓库现有实现FieldsBuilder、AOT Module、C 侧 module_loader.h 等之间的对应关系。背景为什么全部 SNode 都能 AOT是个问题在 RFC 写作时的 Taichi 中字段的典型定义与使用方式如下a ti.field(ti.i32) b ti.field(ti.f32) ti.root.pointer(ti.ij, 16).dense(ti.ij, 16).place(a, b) ti.kernel def run(): for I in ti.grouped(a): b[I] a[I] * 4.2这种写法对 Python 用户非常友好但对部署侧AOT 场景提出了三个挑战Taichi 字段目前是全局变量实现的。这导致 Taichi kernel 变得不纯not pure依赖隐式信息。将这样的 kernel 保存进 AOT 模块时还必须把其依赖的全部全局状态一并保存。理想情况下用户应该能创建 Taichi 字段并像参数一样把它们传入 kernel。AOT 模块中缺少 SNode 类型信息。要朝把字段传入 kernel的方向前进字段与 SNode 的类型都必须被保存进 AOT 模块。字段数据不由用户管理。由于字段是全局的Taichi 运行时必须负责创建和管理它们。若把字段局部化、与 Taichi kernel 解耦用户就能自行管理这些字段的内存资源。RFC 由此给出了明确的Goals提供一种 SNode API让 SNode 与 Taichi 字段可以被局部化从而让 kernel 变得纯pure支持显式描述完整的 SNode 树类型使 SNode 类型可被序列化进 AOT 模块从而让 AOT 支持所有种类的 SNode新 SNode API 需兼容既有用法不确定但强烈期望将元素类型与 SNode 类型解耦解决矩阵字段必须以分散方式实现才能支持 SoA 布局的问题。同时明确了一个Non-Goal不打算把稀疏 SNode 的支持从 LLVM codegen 扩展到其他后端尤其是 SPIR-V。事实核对上述三点背景与目标来自 RFC 原文仓库实现侧LLVM 后端 AOT builder 的注释也印证了序列化最小单元是整棵 SNodeTree的结论见下文仓库中的落地佐证。核心设计一第一次尝试为何行不通一个直觉上的方案是允许字段作为 kernel 参数a ti.field(ti.i32) b ti.field(ti.f32) ti.root.pointer(ti.ij, 16).dense(ti.ij, 16).place(a, b) ti.kernel def run(a: ?, b: ?): for I in ti.grouped(a): b[I] a[I] * 4.2 run(a, b)但 RFC 明确指出这对 AOT 并不真正可行因为a和b是一个树类型的属性attributes of a tree type你无法单独 dumpa和b的类型。为了讲清这个问题RFC 用 C 做了等价类比struct AB { int32_t a; float b; }; using TreeType PointerDenseAB;此时你无法把 kernel 声明成void run(? a, ? b)正确做法是把整个TreeType实例作为一个整体传入即void run(TreeType tree)。这背后的原因是在使用 Taichi 的 SNode 系统构造层级结构的同时你也在构造一个SNodeTree类型——该工作由 Taichi 的 FieldsBuilder 完成RFC 原文此处即引用了该实现文件。核心设计二可行方案——类型与实例解耦RFC 的解决思路是显式化 SNode 树及其类型引入SNodeTreeBuilder。每个字段通过add_field()注册到 builder 中add_field()不做任何内存分配只返回一个field handle字段句柄供 kernel 内部从树中取回字段。builder ti.SNodeTreeBuilder() builder.add_field(dtypeti.f32, namex) builder.add_field(dtypeti.i32, namey) builder.tree() .pointer(ti.ij, 4) .dense(ti.ij, 5) .place(x, y) # tree_t stands for tree type. tree_t builder.build()同理SNodeTreeBuilder.build()也不为树分配内存它只构建一棵 SNode 树的类型。之后你可以用tree_t.instantiate()来实例化一棵树。类型-树解耦的设计动机有两点我们显式拿到了 SNode 树类型。这对 AOT 是必须的同时也可用作类型注解提升语言的形式化程度。我们可以从同一个类型实例化出任意多棵树并传给同一个 kernel 而无需重新编译。在 Taichi kernel 内部整棵树可以这样使用ti.kernel def run(tr: tree_t): for I in ti.grouped(tr.x): tr.x[I] tr.y[I] 2.0 tree tree_t.instantiate() run(tree)与既有 API 的唯一变化是字段前需要加上tree.前缀下标操作仍发生在字段上而非树上即tr.x[I]而不是tr[I].x。两种从树中取回字段的方式按名称by nameadd_field()接收name参数。构建完 SNode 树后Taichi 会为该树上的每个已注册字段生成一个属性因此可以直接写tr.x访问名为x的字段。name是字段在树中的唯一标识符注意在place时传入的也是名字。按字段句柄by field handle也可以使用add_field()返回的句柄来访问字段builder ti.SNodeTreeBuilder() x_handle builder.add_field(dtypeti.f32, namex) # boilerplate to generate tree type and instantiate a tree ... ti.kernel def foo(tr: tree_t): x ti.static(tr.get_field(x_handle)) # 1 for i in x: x[i] i * 2.0注意该设计要求 kernel 中的部分第 1 行在 Python 侧求值同时把全局变量x_handle拉进了 kernel某种程度上违背了最初纯化的目标。RFC 对此的取舍是可以要求x_handle作为参数传入 kernel或者干脆把它看作一个无足轻重的 Python 常量。定义shape与ti.field()类似add_field可以接收shape参数。一旦指定builder 会自动在树根下创建一个新的dense字段注意指定shape后就不应再做一次placebuilder ti.SNodeTreeBuilder() builder.add_field(dtypeti.f32, namex, shape(4, 8)) # This would result an error # builder.tree().dense(ti.ij, (4, 8)).place(x) tree_t builder.build()它等价于显式写法builder ti.SNodeTreeBuilder() builder.add_field(dtypeti.f32, namex) builder.tree().dense(ti.ij, (4, 8)).place(x) tree_t builder.build()AoS 与 SoA复合类型与字段视图field view需要在 AoS/SoA 之间切换的两种复合类型是ti.Matrix与ti.Struct。AoS 很直接直接把复合类型用作字段的dtype即可。builder ti.SNodeTreeBuilder() builder.add_field(dtypeti.vec3, namex) # ti.vec3 is a vector of 3 ti.f32s builder.dense(ti.i, 8).place(x) tree_t builder.build()SoA 则麻烦一些。RFC 写作时的现行做法是把复合类型的每个分量当作独立的标量 Taichi 字段如下例必须手动分别 placex的 3 个底层分量# Current way (as of v1.0.1) of doing SoA in Taichi x ti.Vector.field(3, ti.f32) for f in x._get_field_members(): # x consists three scalar f32 fields ti.root.dense(ti.ij).place(f)这种做法在多处引入混乱类型不单纯由dtype决定还取决于字段如何被 place引入了嵌套字段nested field概念而 Taichi 对此缺乏良好抽象。这使得对复合类型字段做某些优化例如在特定平台上向量化 load/save 与标量操作带宽相同变得复杂——没有良好抽象时判断矩阵字段是 AoS 还是 SoA 的检查不得不散布在 CHI IR 的不同 pass 中进一步思考会发现SoA 的x其实不是一个真正的字段而是三个独立标量字段的分组视图grouped view——该视图提供对单个标量字段无意义的矩阵运算。由于类型目前与字段定义耦合Taichi 字段为了支持 SoA 场景不得不实现为一个个独立字段一旦切换到类型 builder 模式就可以先控制类型如何构建再选择字段实现方式。若想把这是一个字段视图显式表达出来RFC 给出了add_field_view设计builder ti.SNodeTreeBuilder() builder.add_field(dtypeti.f32, namev0) builder.add_field(dtypeti.f32, namev1) builder.add_field(dtypeti.f32, namev2) for v in [v0, v1, v2]: builder.tree().dense(ti.ij, 4).place(v) # Checks that # 1. components and dtype are compatible. # 2. If dtype is a vector/matrix, then all the fields in components are homogeneous in their SNode hierarchy. builder.add_field_view(dtypeti.vec3, namevel, components[v0, v1, v2])矩阵字段视图支持常见的矩阵操作等价于把每个分量展开成局部矩阵变量# 1 vel_soa[i, j].inverse() # equivalent to ti.vec3([v0[i, j], v1[i, j], v2[i, j]]).inverse() # 2 vel_soa[i, j][1] 2.0 # equivalent to v1[i, j] 2.0 # 3 vel_soa[i, j] vel_soa[i, j] some_vec3 # equivalent to vel_tmp ti.vec3([v0[i, j], v1[i, j], v2[i, j]]) vel_tmp vel_tmp some_vec3 v0[i, j] vel_tmp[0] v1[i, j] vel_tmp[1] v2[i, j] vel_tmp[2]字段视图还可以嵌套例如用三个已注册字段构造出结构体视图vertex_t ti.types.struct({pos: ti.vec3, normal: ti.vec3}) sphere_t ti.types.struct({center: vertex_t, radius: ti.f32}) builder ti.SNodeTreeBuilder() builder.add_field(dtypeti.vec3, namepos) builder.add_field(dtypeti.vec3, namenormal) builder.add_field(dtypeti.f32, nameradius) builder.add_field_view(dtypesphere_t, namespheres, components[[pos, normal], radius]) ### ^^^^^^^^^^^^^^^^^ Note this is nested梯度与自动微分为支持 autodiffadd_field()仍需要接收needs_grad: bool参数b ti.SNodeTreeBuilder() b.add_field(dtypeti.f32, namex, needs_gradTrue) # AOS b.tree()....place(x, b.grad_of(x)) # or SOA b.tree()....place(x) b.tree()....place(b.grad_of(x))当needs_gradTrue时原始primal字段与伴随adjoint字段定义在同一棵树内需要用b.grad_of(primal_name)来获取伴随字段的句柄。RFC 特意指出备选方案是使用f{primal_name}.grad这种命名约定但感觉太临时/太 hacktoo ad-hoc。如果你不想手动 place 梯度字段也可以在末尾调用builder.lazy_grad()它会自动 place 所有梯度字段。这一设计在仓库中确有对应实现全局 builder 的lazy_grad会触发root.lazy_grad()见 fields_builder.py调试模式下还会在 materialize 时自动分配伴随 checkbit见 impl.py 中root._allocate_adjoint_checkbit()的调用。Python AOT API保存 SNode 树类型RFC 设想的 Python AOT API 如下builder ti.SNodeTreeBuilder() # ... tree_t builder.build() ti.kernel def foo(tr: tree_t): # ... m ti.aot.Module(arch) m.add_snode_tree_type(tree_t, namevel_tree) m.add_kernel(foo) m.save(/path/to/module)在仓库当前的落地实现中见 python/taichi/aot/module.pyti.aot.Module(arch)构造时会通过rtm._finalize_root_fb_for_aot()把全局根 FieldsBuilder 以仅编译类型compile_only的方式 finalize然后由prog.make_aot_module_builder(arch, caps)创建后端对应的 builder字段通过Module.add_field(name, field)加入内部调用self._aot_builder.add_field(...)kernel 通过Module.add_kernel(kernel_fn)加入内部调用self._aot_builder.add(kernel_name, kernel.kernel_cpp)最后Module.save(filepath)落盘并在目录中额外写入__content__与__version__文件记录模块内容清单与 Taichi 版本。可见 RFC 中整棵树类型入库的思想在落地时演化为以字段其背后是整棵 SNodeTree为单位入库但kernel 与字段类型分离、可独立加载的架构与 RFC 一脉相承。C AOT API加载并实例化树RFC 设想的 C 侧 API 清晰地演示了按类型取树 → 分配内存 → 实例化 → 启动 kernel的完整链路auto mod taichi::aot::Module(/path/to/module); auto *tree_t mod-get_snode_tree(vel_tree); taichi::Device::AllocParams alloc_params; alloc_params.size tree_t-get_size(); auto *tree_mem device-allocate_memory(alloc_params); // By doing this, the kernel can verify that the passed in memory matches its // signature. auto *tree taichi::instantiate_tree(tree_t, tree_mem); auto foo_kernel mod-get_kernel(foo); foo_kernel-launch(/*args*/{tree});关键点在于内存由用户宿主程序分配kernel 在启动时可校验传入的内存与其签名是否匹配。这与背景中字段数据不再由 Taichi 运行时管理的目标直接呼应。仓库中aot::Module确实提供了Field *get_snode_tree(const std::string name)接口见 taichi/aot/module_loader.hField类还定义了ArgUnion std::variantbool, int64_t, uint64_t, const Field *作为 kernel 参数联合类型module_loader.h说明以整棵树作为 kernel 参数已成为 AOT 加载侧的正式形态。向后兼容ti.root即全局 builderti.field()返回 thunkRFC 要求新 API 兼容既有用法。当时的现状是ti.root已经实现为一个字段累加器——root 中累积的所有字段会在 kernel 调用时被物化为一棵新的 SNode 树。先看既有写法x ti.field(ti.f32) ti.root.pointer(ti.i, 4).dense(ti.i, 8).place(x) ti.kernel def foo(): for i in x: x[i] i * 2.0其使用新 API 的等价写法为b ti.SNodeTreeBuilder() b.add_field(ti.f32, namex) b.tree().pointer(ti.i, 4).dense(ti.i, 8).place(x) tree_t b.build() tr tree_t.instantiate() ti.kernel def foo(): for i in tr.x: tr.x[i] i * 2.0为实现向后兼容需要两类辅助机制把xold映射到tr.xnew且运行时需要知道xold属于哪棵 SNode 树ti.field()返回的xold在ti.root当前 SNode 树被构建并实例化之前只是一个字段占位符。RFC 给出的可行方案是ti.root就是一个全局的SNodeTreeBuilderti.field()返回一个FieldThunkthunk 即延迟求值的占位对象class FieldThunk: def __init__(self, fid): self.field_id fid self.tree None def bind(self, tree): self.tree tree def field(dtype, name, shapeNone, offsetNone, needs_gradFalse): name name or random_name() handle ti.root.add_field(dtype, name) ft FieldThunk(handle) ti.root._field_thunks.append(ft) return ft在物化 SNodeTree 时tree_t ti.root.build() tree tree_t.instantiate() ti._runtime.global_snode_trees.append(tree) for ft in ti.root._field_thunks: ft.bind(tree) # Make ti.root a new SNodeTreeBuilder to allow for dynamic fields ti.root SNodeTreeBuilder()JIT 编译 Taichi kernel 时把xold变换为x.tree.get_field(x.field_id)其中x是FieldThunk。仓库实现对照这一全局 root builder 延迟 finalize 重建新 builder的模式在仓库中真实存在。Runtime维护unfinalized_fields_builder注册表impl.pymaterialize_root_fb()在首次 kernel 调用或 AOT 时 finalize 全局 root并随后重建一个新的全局FieldsBuilder以支持动态字段impl.py未 finalize 的非 root builder 会在 kernel 编译前被validate_fields_builder()拦截报错。这与 RFC 的每次物化后把ti.root换成新 builder的设想一致。仓库中的落地佐证从 RFC 到实现RFC 是 2022-04 的设计提案其核心思想在仓库中已有相当程度的落地可沿以下路径继续深入阅读字段构建器python/taichi/_snode/fields_builder.py 中的FieldsBuilder是 RFC 中SNodeTreeBuilder的落地形态对外暴露为ti.FieldsBuilder与全局ti.root。它提供dense/pointer/dynamic/bitmasked/quant_array/place/lazy_grad/finalize等接口finalize(compile_onlyFalse)与_finalize_for_aot()即compile_onlyTrue分别对应运行时物化与AOT 仅编译类型两种路径fields_builder.py。注意pointer、dynamic、bitmasked等稀疏类型在构造时会检查当前后端是否支持 sparse extension不支持则抛出TaichiRuntimeError——这正是 RFC Non-Goal稀疏 SNode 暂不扩展到 SPIR-V 等后端在实现层的体现fields_builder.py。AOT 模块python/taichi/aot/module.py 的Module类负责把 kernel/字段/图序列化到磁盘目录并支持.tcm归档打包archive()。LLVM 后端的序列化粒度taichi/runtime/llvm/llvm_aot_module_builder.cpp 的add_field_per_backend()注释明确写道字段指 SNodeTree 中的叶子Place SNode单独序列化叶子或其分支没有意义我们必须序列化的最小单元是整棵 SNodeTree且 SNodeTree 以snode_tree_id作为标识符而非字段名多个字段可能指向同一棵 SNodeTree。这从实现层面印证了 RFC无法单独 dump 字段类型、必须整体保存树类型的核心论断。GFX 后端的树内存管理taichi/runtime/gfx/snode_tree_manager.cpp 的SNodeTreeManager通过materialize_snode_tree()编译 SNode 结构并分配 root buffer通过get_field_in_tree_offset()计算树内字段偏移、get_snode_tree_device_ptr()取得设备指针——对应 RFC C API 中实例化树并管理其内存的职责划分。C 端测试验证tests/cpp/aot/llvm/field_aot_test.cpp 展示了完整的 C 加载流程mod-get_kernel(...)取出 kernel、mod-get_snode_tree(0)按snode_tree_id取树、LLVM::allocate_aot_snode_tree_type()分配树内存随后通过LaunchContextBuilder设置参数并依次 launchinit_fields、check_init_x等 kernel覆盖 CPULlvmAotTest.CpuField与 CUDALlvmAotTest.CudaField在TI_WITH_CUDA且 CUDA 可用时运行两个后端还包含对 pointer 字段 deactivate/activate 的验证——即AOT 支持全部 SNode含稀疏在 LLVM 后端的回归测试。备选方案与 FAQRFC 在 Alternatives 一节坦言不确定是否有更好的设计能覆盖上述全部目标。FAQ 一节当时标注为 TBD待补充本文不臆造其内容。小结这条 RFC 的价值在于指出了 Taichi 从Python 内嵌的全局字段 DSL走向可部署的 AOT 运行时之间最关键的抽象缺口字段类型无法脱离 SNode 树类型而独立存在。其给出的答案——引入显式的树类型构建器、类型与实例解耦、以整棵树为 AOT 序列化与 kernel 参数的最小单元、用FieldThunk兼容旧 API——在仓库的FieldsBuilder、Module、SNodeTreeManager与 LLVM/GFX AOT builder 中均有迹可循。对希望深入理解 Taichi AOT 工作流tests/cpp/aot 目录下有大量相关测试或在其上做二次开发的读者而言这份 RFC 与上述源码共同构成了一条完整的学习路径。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考