Skip to content
返回

Volvox 深度学习框架中的元胞(Cell)概念与节点融合设计

发布于
Cell := {
   inputs: Dict Node Cell,
   output: Node,
}

语义:

核心意图

  1. inputs_node:Cell 能融合到的”最前面的节点层”
    • 通常是 Reindex 节点(作为 Cell 的”输入适配层”)
    • 可能为空(None),表示某个输入路径上没有 Reindex
    • 这一层的节点数量 = Cell 的输入数量
    • 它们的输入才是真正的”外部依赖”(来自其他 Cell 或 Input/Constant)
  2. output_node:Cell 的”计算体 + 输出”
    • elemwise:融合的 Elemwise/Cast 链(可能为空)
    • reduce:最终的 Reduce 节点(可能为空)
    • 约束:如果 elemwise 为空,则 inputs_node 只能有一个元素(直接连到 reduce)
  3. 推导关系:
    • inputs(Cell 依赖)从 inputs_node 的输入推导
    • output_shape/dtype 从 output_node 推导
  4. 融合的终止条件(边界条件) 你说的两种情况,我帮你补充完整: 硬边界(必须停止融合)
  5. 遇到 Reduce 节点
    • Reduce 作为 output_node:可以融合(Reduce 是 sink)
    • Reduce 作为中间节点:停止(Reduce 的下游不融合,除非你做 epilogue)
  6. 遇到多 consumer 节点(fan-out)
    • 节点被多个下游使用 → 它应该是某个 Cell 的 output_node,停止向上融合
    • 否则会重复计算这个节点
  7. 遇到 Input/Constant
    • 它们是叶子节点,加入 inputs_node,停止向上
  8. 遇到不兼容的索引/形状变换
    • 多个 Reindex 无法合成一个等价的 index_map
    • Overflow conditions 冲突
  9. 遇到其他 Cell 的 output_node
    • 这是跨 Cell 依赖,应该加入 inputs_node,停止 软边界(可以停,也可以继续,取决于策略)

我重新仔细思考了一下,Cell应该是 [Pasted ~7 lines] 原本的inputs输入可以被自动推断(Cell本身不用管),刚才对inputs_node的描述可能不清晰,他是指能融合到Cell的最前面的节点,一般是reindex,但是也有可能前面没有reindex,那就留空(inputs本身也是从这里推断出来的,如果这里为空才去elem搜,另外需要澄清如果elem为空就意味着inputs_node只能有一个,且为reduce输入,如果唯一input_node也为None,那其实Cell相当于就是个Reduce了)

在基于 Node 的 IR 上做「融合」时,需要一个明确的最小融合单元和融合规则。本文整理在 volvox 里用 Cell 表示「融合完成后的一个最小单元」的设计:Cell 的逻辑形态、链上各段如何由多节点融合而成,以及何时终止融合。

Node 系统简述

Node 是树形结构,用代数数据类型和访问者模式做遍历与变换。形态大致如下(类型名与 Haskell 风格一致):

通过 NodeVisitor avisitNode 对每种节点做递归访问,得到统一的结果类型 a。在此之上可以扩展分析、改写和融合等逻辑。

Cell 与融合目标

Cell 定义为「融合完成后的一个最小执行单元」。目标是在不改变语义的前提下,把一条计算链上能一起做的节点合并成尽量少的逻辑段,减少遍历与中间结果。

逻辑形态上看,一条链可以看成:

也就是说:Cell 内允许多个 Reindex、多个 Elemwise/Cast、多个 Reduce 分别先各自融合,再在链上按 Reindex → Elemwise/Cast → Reduce 这样的顺序组成一段完整计算。多段 Elemwise 不会强制拆成多个 Cell,而是先融合成一段再参与链的组成。

串行同类节点能否融合

链上同一类型的节点若串行出现(前一个的输出是后一个的输入、且无其它使用者),可以考虑合并为一个逻辑节点,从而在一个 Cell 内用「一段」表示。

多 Elemwise 串行 可以融合为一个 Elemwise:在一次遍历中按顺序执行多个 Kernel 的表达式,等价于多次逐元素遍历。实现上把多个 Kernel 合成为一个复合 Kernel 即可。

多 Reindex 串行 可以融合为一个 Reindex:索引为 ix2 ∘ ix1(先做 ix1 再做 ix2),溢出与默认值在实现上合并处理。这样链上只保留一个逻辑 Reindex 节点。

多 Reduce 串行 可以融合为一个 Reduce:在合并后的维度和归约器上做一次归约。实现时需要处理好维度与 Reducer 的对应关系。

因此,在「Cell 的逻辑形态」里,链上的每一段(Reindex / Elemwise / Reduce)都允许是多节点融合后的一个逻辑单元,而不是「至多一个 Elemwise、多了就拆 Cell」。这样既和「多个 Elemwise 可完美融合成一个」一致,也显式承认串行 Reindex、串行 Reduce 同样可融合,实现时只需做好索引复合与 Reducer/维度合并。

融合终止条件

在从某条链的某个节点开始「向前/向后」做融合时,需要在以下情况终止,不再把更多节点纳入当前 Cell:

其它实现层面的约束(如显存、 kernel 大小等)也可以作为额外终止条件,但逻辑上以上两条是保证正确性与融合边界清晰的最小集合。

小结

这样在 volvox 的 Node 系统上,可以用 Cell 作为融合的抽象单位,并在此基础上实现 fuse pass 和代码生成。


建议修改

上一篇
E-graphs 学习笔记:动机与 Explanations
下一篇
在 PyO3 里优雅地使用 `#[pyclass]` enum