Mojo 编译器 KGEN 并行展开Parallel Elaboration从参数求解到并行调度的核心算法解析【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo本文基于 Mojo 编译器设计文档 ParallelElaboration.md系统讲解 KGEN 管线中最复杂的 pass——Elaborator展开器的并行化实现它如何将参数化 IR 通过生成器实例化转化为具体 IR如何用挂起/恢复的迭代式任务模型替代递归遍历以及如何借助 AsyncRT 工作队列实现多版本multi-versioning、fork 传播、搜索隔离与循环检测的并行调度。读完本文你将理解kgen.param.fork、kgen.param.evaluate等操作的底层展开语义并能在 Mojo/lib/Elaborator 源码中找到ParamNode/ImplNode的对应实现具备独立阅读和排查 elaborator 相关问题的能力。展开器的定位KGEN 管线中从参数化 IR 到具体 IR 的桥梁Elaborator 定位Elaborator 是 KGEN 管线中复杂度最高的 pass 之一其职责是执行生成器实例化generator instantiation——相当于 C 的模板实例化但涉及更多机制。其核心算法在概念上是一次图遍历因此可以被并行化。文档假设读者已了解 elaboration 与参数parameter的基本概念。从源码结构看展开器实现在 Mojo/lib/Elaborator/Elaborator.h 中Elaborator类的注释明确写道constructs the expansion tree as it walks the IR and specializes operations. This outputs IR that has been fully specialized/concretized, with the appropriate functions multi-versioned——即边遍历 IR 边构建展开树输出完全特化/具体化后的 IR并对函数做相应的多版本化。文档将 elaboration 清晰地划分为三个基本正交的部分参数解析Parameter Resolution调用图实例化与多版本化Callgraph Instantiation and Multi-VersioningJIT 与求值JIT and Evaluation参数解析把生成器体视为一组参数方程参数解析是指从单个生成器及其函数体的视角看参数实例化的过程。生成器的函数体中包含一张参数图parameter graph描述参数使用与定义之间的关系。图中可以有多层嵌套的作用域但为简化起见可先忽略。相对于生成器参数图可以理解为一组方程每个参数的值都由求值一个参数表达式parameter expression来确定——例如add(x, 2)是参数表达式apply(:(index) - index mul2, 4)也是——前提是给定一组输入参数生成器的入边。以文档给出的例子kgen.generator fooa, b - e() { kgen.param.declare c add(a, b) kgen.call otherc - d() : () - () kgen.param.result_bindsub(d, b) kgen.return }其中参数方程为c a bd other(c)e d - b对other的调用可以看作对某个参数化函数的求值。用给定的输入参数实例化生成器只需把这些输入值代入方程计算其余所有声明的参数然后把计算出的值替换到生成器函数体的各处操作、属性和类型中。从源码看参数图的承载结构是ParameterUseDefGraph每个实现节点都持有一份见 Mojo/lib/Elaborator/IREvaluatorContext.h 中ImplNodeBase::paramGraph的定义。调用图实例化与多版本化fork 的向上传播上一节讲的是单个生成器层面的参数解析但 elaborator 更大的关切是调用图如何被处理并实例化以及它与 KGEN 生成器独有的**多版本化multi-versioning**特性如何相互作用。展开前的调用图pre-elaboration callgraph在参数化生成器之间连边展开后的调用图post-elaboration callgraph在具体函数concrete function之间连边kgen.generator parametrica() - index { %0 kgen.param.constant a kgen.return %0 : index } kgen.generator main() { %0 kgen.call parametric1() : () - index %1 kgen.call parametric2() : () - index kgen.return }在这个例子中参数化调用图有 2 个节点main和parametrica。具体调用图有 3 个main、parametric1、parametric2边从main指向后两者。调用图的这种膨胀expansion构成了 elaborator 主算法的骨架也是我们想要并行化的部分。Elaborator 会创建parametrica的两个版本对应两组不同的输入参数。但 KGEN 还允许同一组输入参数的生成器可能产生多个或零个实例化——这是由 fork 造成的kgen.generator parametrica() - index { kgen.param.fork result [a, add(a, 1)] %0 kgen.param.constant result kgen.return %0 : index } kgen.generator main() { kgen.call parametric1 : () - () kgen.return }这里parametrica在result的两个可能取值上发生 fork。它在main中如何被消解fork 会沿着膨胀图一路向上传播最终产生两个不同的main版本kgen.generator parametric,a1,result1() - index { %0 kgen.param.constant 1 kgen.return %0 : index } kgen.generator parametric,a1,result2() - index { %0 kgen.param.constant 2 kgen.return %0 : index } kgen.generator main,parametric,a1,result1() { kgen.call parametric,a1,result1() : () - () kgen.return } kgen.generator main,parametric,a1,result2() { kgen.call parametric,a1,result2() : () - () kgen.return }关键点在于fork 可能在任何生成器函数体的参数解析过程中发生这意味着参数解析必须在生成的两个版本上从同一位置继续。JIT 与求值搜索根终止 fork 传播细心的读者会注意到如果 elaborator 需要把 fork 沿调用图向上传播它在哪里终止答案是求值evaluation。求值就是从多个具有相同签名的具体函数中选定一个具体函数的过程。这些候选函数可以是完全不同的函数也可以是同一函数的不同实例化如上文。求值通常涉及对不同实现的基准测试benchmarking为了速度但 evaluator 本身是用户编写的代码。evaluator 以及候选函数都由 KGENExecutionEngineJIT 编译然后把 evaluator 的函数指针传给它它随后可以任意处理但必须返回被选中候选的索引。kgen.generator two_candidates() { kgen.param.fork N [1, 2] kgen.return } kgen.generator evaluator(%fns: !kgen.pointer() - (), %sz: index) - index { // Always pick the second one. %idx1 index.constant 1 kgen.return %idx1 : index } kgen.generator foo() { kgen.param.evaluate selected: () - () [two_candidates] with [(!kgen.pointer() - (), index) - index: evaluator] kgen.call_param[() - (): select]() kgen.return }在这个例子中 evaluator 被简化为固定选择。一般来说对 evaluator 的唯一要求是搜索必须尽可能在隔离环境中进行。如果编译器一边在编译代码、一边在 benchmark 其他代码结果将不准确。从源码结构看fork 在展开中的传播可以在生成一个多版本化的具体函数上观察到例如kgen.func someFunc,foo,a1,N1与kgen.func someFunc,foo,a1,N2的命名编码了祖先参数值这与上文中main,parametric,a1,result1的混淆命名规则一致。核心算法递归的深度优先遍历在展开完成之前展开前调用图和膨胀图在任何时刻都不完全已知。某个生成器最终产生哪些实例化是其输入参数的函数kgen.generator pickInstantiationc: i1() { kgen.param.if c { kgen.call foo1() : () - () kgen.param.yield } else { kgen.call bar() : () - () kgen.param.yield } kgen.return }我们拿到的是调用图和膨胀图的根节点即所谓的primary generators对应导出函数或通常是main。访问一个节点并做参数解析就会揭示膨胀图的出边。重要的是生成器实例化是与参数解析天然递归交织的过程kgen.generator foo() - out() { kgen.param.result_bind1 kgen.return } kgen.generator main() { kgen.call foo[] - out() : () - () %0 kgen.param.constant out kgen.return }对main做参数解析会揭示out foo()必须递归进入foo的实例化foo实例化完成后弹回main继续参数解析。可以看到 elaboration 的基本算法是用给定输入参数访问一个生成器创建一个与它有相同函数体的函数开始参数解析。在给定具体输入的前提下求解函数中的参数方程。遇到生成器实例化通常是 call 操作时如果该生成器尚未被访问过就递归进入。当该生成器完成后弹回来继续参数解析。这会以深度优先DF顺序遍历并揭示膨胀图。为了处理递归每个生成器实例化在被访问时都可以被标记为进行中in-progress之后再次访问一个进行中的实例化就说明存在递归。递归可以被适当打断简单情况下直接引用那个进行中的具体函数即可。Fork 的实现ParamNode 与 ImplNode前文已经说明参数化生成器如何产生多个实例化且膨胀图的边存在于实例化之间生成器 输入参数的配对。但有了 fork 之后每个生成器实例化本身可以有多个实现。在 elaborator 中生成器实例化被表示为ParamNode它的每个实现是一个ImplNode。在源码中这两个结构分别定义于 Mojo/lib/Elaborator/IREvaluator.h 与 Mojo/lib/Elaborator/IREvaluatorContext.hImplNode的注释写道This struct represents a concrete instantiation of a generator -- generators may have multiple concrete instantiations -- and contains the current state of elaboration for that concrete instanceParamNode则是expansion tree 的节点。当一个生成器实例化被首次访问时创建一个ImplNode表示展开的初始状态。假设 elaborator 遇到一个 fork// The generator inside the ParamNode. kgen.generator fooa() { kgen.param.fork N [a, add(a, 1)] %0 kgen.param.constant N kgen.return } // The state of foo1 when the fork is encountered. kgen.func foo,a1() { kgen.param.fork N [1, 2] %0 kgen.param.constant N kgen.return }Elaborator 会克隆当前的ImplNode包括当前函数体及其全部展开状态// The original implementation. kgen.func foo,a1,N1() { kgen.param.declare N 1 %0 kgen.param.constant N kgen.return } // The forked implementation. kgen.func foo,a1,N2() { kgen.param.declare N 2 %0 kgen.param.constant N kgen.return }Elaborator 会继续处理当前ImplNode但在其完成后会继续循环处理其他未完成的实现直到不再出现新 fork 且所有实现都完成。此时父ParamNode才被认为完成。当某个函数的展开遇到具有多个实现的生成器实例化如上文的foo1时该函数会以完全相同的方式被 fork// Multiple implementations of foo are propagated by multi-versioning // someFunc as well. kgen.func someFunc,foo,a1,N1() { kgen.call foo,a1,N1() : () - () kgen.return } kgen.func someFunc,foo,a1,N2() { kgen.call foo,a1,N2() : () - () kgen.return }正如前文所述fork 传播可以在搜索根search root处被打断——由kgen.param.evaluate操作表示——它能把例如具有多个实现的生成器实例化缩减为 1 个。错误与传播约束作为快速失败的编译期优化Elaboration 在处理某个实现和处理某个生成器实例化时都可能失败。失败原因多种多样但通常来自用户编写的静态断言kgen.generator foobara() { kgen.param.assert gt(a, 1), a must be bigger than 1! kgen.return }Elaborator 会先创建第一个实现但在处理该函数时会遇到静态断言并可能失败。若失败该实现即被视为失败。若某生成器实例化的所有实现都失败则该实例化也失败。引用了失败实例化的函数同样失败依此类推。注意如果一个生成器有多个实例化fork 只会使用成功的实现。例如下面这个例子只会产生 1 个候选kgen.generator baz() { kgen.param.fork a [1, 2] kgen.call foobara() : () - () kgen.return }处理baz会创建 2 个实现但a 1的那个会失败因为foobar1的实例化会失败。因此baz最终只有 1 个有效实现。从源码看错误状态确实内建于节点之中ImplNodeBase持有std::optionalErrorTree error与原子标志hasError见 Mojo/lib/Elaborator/IREvaluatorContext.h允许在可恢复的情况下延迟错误处理ParamNode一侧则用DoneState { NOT_DONE, DONE, ERROR }原子标记区分完成态与错误态见 Mojo/lib/Elaborator/IREvaluator.h。约束Constraints生成器可以直接携带约束——即只依赖输入参数的静态断言kgen.generator foobara() constraints [gt(a, 1), a must be bigger than 1!] { kgen.return }这本质上是一个编译期优化如果一个生成器实例化在克隆函数体、开始处理实现之前就能检查其有效性那么展开在这里就会更快地失败。状态保存与恢复、挂起与递归把递归改写成迭代前几节勾勒了完整 elaboration 算法的组成部分省略了一些细节如 bindings整体图景已经清晰。但存在两个问题递归会使编译器栈溢出。递归的 DF 算法不容易并行化。递归进入生成器实例化是写这个算法的自然方式但调用图可能很深用户写了大量代码时很容易让 elaborator 栈溢出。而且它也不是容易并行化的形态。因此第一步是把算法改写为迭代式。由于在 fork 展开一个实现节点时我们已经保存了节点的展开状态稍加改造、把状态从调用栈移到ImplNode上就可以把算法中递归的部分改造成挂起suspend并退出某个ImplNode的展开去处理那个实例化等它完成后再把挂起的ImplNode重新入队。这就构成了并行化的基础因为这一步可以异步执行。概括一下当展开遇到一个新的生成器实例化时当前正在处理的实现被挂起并作为**等待者waiter**挂到新实例化上。该实例化被放入队列当它完成时所有等待者被重新压入队列并恢复。在源码中挂起状态被显式地挂在实现节点上ImplNode持有stackcurrent stack of worklists and scopes、dependencies延迟处理的生成器实例化列表、blockers下游阻塞本节点展开的节点它们必须全部完成本节点才能继续丢掉任何一个都会使展开死锁等字段见 Mojo/lib/Elaborator/IREvaluator.h正是把调用栈搬进节点的直接体现。apply的一个特殊点一个麻烦在于生成器实例化可能嵌套在参数表达式的任何位置——任何操作的类型和属性内部。这意味着任何操作的展开都可能挂起。实际上这并不会发生因为lift-and-fold-apply优化会把apply算子提出来变成kgen.param.apply。文档明确提到这部分处理需要重新审视。递归的三个附加难题前文讨论了在 DF 遍历下如何处理递归——只要设置visited标志即可。但递归还有几个额外问题递归的生成器实例化不能有多个实现。递归的生成器实例化不能有结果参数result parameters。bindings 需要不动点迭代才能在环中传播。如果递归实例化有多个实现fork 会向无穷大爆炸。这与有效实现的数量无关elaborator 在递归实例化展开完成之前无法知道哪些实现有效而这又要求先假设它们都有效。如果递归生成器有结果参数就会在膨胀图层面给参数的 use-def 图引入环。技术上下面这种写法是有良好定义的kgen.generator foo() - x() { kgen.call foo[] - y() %0 kgen.param.constant y kgen.param.result_bind2 kgen.return }但出于关注点分离elaborator 不会穿透生成器参数方程的抽象去支持这种极其罕见乃至不存在的用法。最后环中的 bindings 是 elaborator放弃的部分因为正确地在环中传播 bound 函数需要不动点迭代而当全图未知时这根本不可能。另一个问题是当遍历不再深度优先时检测递归变得相当棘手。这一点很关键因为引入并行化之后图将不再以深度优先方式遍历。通往并行化任务模型、原子计数与工作队列耗尽并行化 elaborator 时需要做的核心决策是什么是核心任务前文已经确立处理一个实现节点ImplNode就是要被并行化的主任务。这些任务可以挂起、派生其他任务、被恢复。文档感叹If only C17 had coroutines!——C17 要是能有协程就好了。这意味着我们可以从根节点开始并行 elaboration并行处理 fork。思路很直接每个ParamNode持有一个原子变量表示进行中的ImplNode数量每完成一个就递减减到零时父ParamNode完成所有等待它的任务通过AsyncValueRefChain的 waiter 列表保存被恢复。周而复始。ParamNode的状态也必须是原子的因为只允许一个任务启动某生成器的特化。在源码中这个状态 等待者计数的原子合并被精确实现为ParamNodeState见 Mojo/lib/Elaborator/IREvaluatorContext.h一个 64 位原子高 2 位表示状态FRESH 0、IN_PROGRESS 1、DONE 3低 62 位是等待任务数。addWaiter()、markInProgress()、markDone()全部通过一次原子操作同时完成状态迁移 计数注释明确说明这样做的目的preventing erroneous waiters from being counted due to a race——防止竞态导致错误地计入等待者。一个巨大的头疼之处是处理搜索与环。我们不能再依赖深度优先遍历来找环且必须保证搜索至少就编译器进程而言在隔离状态下执行——也就是说不能一边让 elaborator 编译代码、一边 benchmark JIT 代码。这要求 elaborator 确保至少就展开而言没有其他任务在运行即耗尽工作队列exhausting the workqueue。elaborator 把 AsyncRT 用作虚拟工作队列但要跟踪有多少活跃任务。每次调度一个任务活跃任务数加一每完成一个减一。为防止计数器短暂归零ParamNode完成时必须先增加将要释放的任务数再emplace 它的 chain。等待者数量和ParamNode的状态必须事务性地修改——这正是通过把两者揉进同一个原子实现的对应上文ParamNodeState。对应的全局结构是ExpansionGraph见 Mojo/lib/Elaborator/Elaborator.h其中numWorkItems原子初始为 1表示当前调度在 elaborator 工作列表上的任务总数worklistCh是当所有活跃工作项完成时被触发的 chain用于在运行 evaluator 前饿死starve工作队列因为在其他线程编译时求值不可靠。这与文档中 AsyncRT 工作队列 quiesce的设计一一对应。耗尽工作队列与搜索Elaborator 的主线程运行到工作列表被耗尽为止当工作项数量归零时emplace 一个 chain主线程等待这个 chain。为了确保搜索在隔离状态下执行当处理kgen.param.evaluate操作时它会并行地发出编译命令产出一个运行搜索的 functor。该 functor 被保存在 elaborator 中当前任务挂起。这意味着工作队列最终会耗尽而没有完成整个图的展开。主线程可以通过检查是否所有 primary generators 都完成来诊断这种情况。若否它检查是否存在延迟的deferred搜索函数并由主线程串行地处理它们释放挂起的任务。环检测与递归Elaborator 不能再假设访问到一个进行中的生成器实例化就意味着递归因为很可能只是另一个线程正在处理那个ParamNode。这意味着一个环会导致工作队列耗尽但没有延迟搜索、也没有完成所有 primary generators。在这种情况下主线程可以从未完成的 primary 节点出发做一次裁剪过的 DFS沿进行中的节点向下找到环发生的位置并打断它或者在上述递归约束不满足时抛出错误。打断边会释放更多任务主线程循环运行直到全部完成。源码中对应的结构是ImplNode的sccRemovedDeps被 diagnoseAndBreakRecursion 移除的 SCC 内部边与sccCh表示 SCC 完成的 chain说明环检测以 SCC 为单位进行。解锁完整并行度异步分发与完成任务按上述写法并行算法会并行处理 primary generators 和 fork也就是说在如下例子中kgen.generator main() { kgen.call foo1() : () - () kgen.call foo2() : () - () kgen.call bar() : () - () kgen.return }main的展开会先遇到foo1、挂起并等待foo1完成、恢复后再遇到foo2、再挂起……这浪费了大量并行度因为foo1、foo2和bar本可以并行展开。最后一步是在函数展开过程中遇到的所有生成器实例化都异步分发。值得注意的是这只对没有结果参数的生成器可行即使在参数 use-def 图内部更细粒度的并行也是可以做到的但结果参数是不常用的特性不值得。重要的是如果任何一个异步分发async的生成器实例化产生了多个实现fork不能在 elaborator 还在处理函数体时执行必须推迟到处理完成之后。为此引入第二种任务ImplNode完成completion任务。在展开函数时异步分发的实例化被跟踪为以 1 为初值的原子计数器中的未完成依赖每完成一个依赖计数器减一当 elaborator 处理完函数体中其他所有内容时计数器再减一次允许完成任务运行。完成任务把所有已完成的依赖按序处理生成 fork 并为其调度必要的完成任务。等待某个ParamNode的完成任务也必须被加入其 waiter 之一。源码层面ImplNodeBase::numDependenciesthis atomic tracks the number of in-flight so-called dependencies. Upon hitting zero, the elaborator will complete processing of this node by handling the calls independencies与ImplNode::dependencies正是这一机制的直接落地见 Mojo/lib/Elaborator/IREvaluatorContext.h 和 Mojo/lib/Elaborator/IREvaluator.h。总结与源码导航本文按照 ParallelElaboration.md 的原始脉络完整呈现了并行 elaboration 的设计可以归纳为一条主线语义层参数解析 解参数方程调用图膨胀 实例化展开fork 实现的多版本化kgen.param.evaluate 搜索根终止 fork 传播。算法层递归 DF 遍历 → 改写为挂起/恢复的迭代任务模型状态保存在ImplNode上。并行层任务 处理一个ImplNodeParamNode用原子状态 等待者计数见ParamNodeState做事务性迁移AsyncRT 工作队列 耗尽工作队列保证搜索隔离主线程串行处理 deferred 搜索与 SCC 环打断无结果参数实例化的异步分发 ImplNodecompletion 任务解锁完整并行度。如需继续深入推荐沿以下路径阅读源码Mojo/lib/Elaborator/Elaborator.hExpansionGraph工作列表计数、worklistCh、quiesceChain与Elaborator主类Mojo/lib/Elaborator/IREvaluator.hImplNodeWorkItem 栈、dependencies、blockers、SCC 边与ParamNodeemplace、setToError、waiterMojo/lib/Elaborator/IREvaluatorContext.hImplNodeBase、ParamNodeState位级原子的状态/等待者编码、ParamNodeBaseMojo/lib/Elaborator/Elaborator.cpp 与 Mojo/lib/Elaborator/ParametricElaborator.cpp展开主流程与参数化求值的具体实现AsyncRT/docs/AsyncRTRuntime.md 与 AsyncRT/docs/WorkQueue.mdelaborator 所依赖的异步运行时与工作队列机制。文档结尾也坦承本文档提供的是驱动 elaborator 的并行算法的高层概览希望足以让读者进入 elaborator 的荆棘丛导航其中仍有许多细节未展开但核心算法及其工作方式已被完整描述。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考