Cell := {
inputs: Dict Node Cell,
output: Node,
}
语义:
- inputs_node 到 output_node 之间的所有节点都属于这个 Cell(通过 output_node 反向可达的闭包)
- 这些节点满足融合条件(后面详细定义)
- inputs_node 本身不属于 Cell(它们是外部依赖,由其他 Cell 或 Input/Constant 提供)
- output_node 属于 Cell(它是 Cell 的计算结果) 重要澄清:
- inputs_node 中的节点不在 Cell 内部执行,它们是”端口/占位符”,代表”从外部拿到的值”
- Cell 实际执行的是:从 inputs_node 的使用者(children)开始,到 output_node 结束
核心意图
- inputs_node:Cell 能融合到的”最前面的节点层”
- 通常是 Reindex 节点(作为 Cell 的”输入适配层”)
- 可能为空(None),表示某个输入路径上没有 Reindex
- 这一层的节点数量 = Cell 的输入数量
- 它们的输入才是真正的”外部依赖”(来自其他 Cell 或 Input/Constant)
- output_node:Cell 的”计算体 + 输出”
- elemwise:融合的 Elemwise/Cast 链(可能为空)
- reduce:最终的 Reduce 节点(可能为空)
- 约束:如果 elemwise 为空,则 inputs_node 只能有一个元素(直接连到 reduce)
- 推导关系:
- inputs(Cell 依赖)从 inputs_node 的输入推导
- output_shape/dtype 从 output_node 推导
- 融合的终止条件(边界条件) 你说的两种情况,我帮你补充完整: 硬边界(必须停止融合)
- 遇到 Reduce 节点
- Reduce 作为 output_node:可以融合(Reduce 是 sink)
- Reduce 作为中间节点:停止(Reduce 的下游不融合,除非你做 epilogue)
- 遇到多 consumer 节点(fan-out)
- 节点被多个下游使用 → 它应该是某个 Cell 的 output_node,停止向上融合
- 否则会重复计算这个节点
- 遇到 Input/Constant
- 它们是叶子节点,加入 inputs_node,停止向上
- 遇到不兼容的索引/形状变换
- 多个 Reindex 无法合成一个等价的 index_map
- Overflow conditions 冲突
- 遇到其他 Cell 的 output_node
- 这是跨 Cell 依赖,应该加入 inputs_node,停止 软边界(可以停,也可以继续,取决于策略)
- Cell 内节点数量超过阈值(控制 register pressure)
- Cell 的输入数量超过阈值(控制参数传递开销)
- 估计的计算复杂度超过阈值(避免单个 kernel 太重)
我重新仔细思考了一下,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 风格一致):
- Input:带元数据的输入。
- Constant:常量(数值 + 数据类型)。
- Elemwise:多子节点 + 逐元素核(Kernel)。
- Reindex:子节点 + 输出形状 + 索引表达式 + 溢出/默认值 + 溢出条件列表。
- Reduce:子节点 + 输出形状 + 索引 + 归约器 + 溢出条件。
- Cast:子节点 + 目标数据类型。
通过 NodeVisitor a 和 visitNode 对每种节点做递归访问,得到统一的结果类型 a。在此之上可以扩展分析、改写和融合等逻辑。
Cell 与融合目标
Cell 定义为「融合完成后的一个最小执行单元」。目标是在不改变语义的前提下,把一条计算链上能一起做的节点合并成尽量少的逻辑段,减少遍历与中间结果。
从逻辑形态上看,一条链可以看成:
- 入口:Input 或 Constant。
- 中间段(顺序可交错、某段可为空):
- Reindex 段:若干 Reindex 融合成「一个逻辑 Reindex」;
- Elemwise/Cast 段:若干 Elemwise、Cast 融合成「一段逐元素/类型转换」;
- Reduce 段:若干 Reduce 融合成「一个逻辑 Reduce」。
- 出口:可以是最后的 Elemwise/Cast,也可以是最后的 Reduce;链尽量拉长,若前面已满足「可融合」条件且能接到一个 Reduce,则把该 Reduce 也纳入 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:
- 遇到 Reduce:Reduce 是天然的边界(读写模式、维度变化与 Elemwise/Reindex 不同),作为 Cell 的出口或段界。
- 被多个节点依赖:若某节点的输出被多于一个后继使用,该节点不能只属于一个 Cell,应作为多个 Cell 的共享输入,在此处切断融合。这类衔接不同 Cell 的节点可称为 Seam(接缝)。
其它实现层面的约束(如显存、 kernel 大小等)也可以作为额外终止条件,但逻辑上以上两条是保证正确性与融合边界清晰的最小集合。
小结
- Cell:融合后的最小单元,对应一条「入口 → Reindex 段 → Elemwise/Cast 段 → Reduce 段(可选)」的链,链尽量长。
- 同类串行可融合:多 Elemwise、多 Reindex、多 Reduce 在串行且无多依赖时,可分别融合为一段,再组成 Cell。
- 终止条件:Reduce 作为边界;被多处使用的节点不继续向内融合,作为共享输入。
这样在 volvox 的 Node 系统上,可以用 Cell 作为融合的抽象单位,并在此基础上实现 fuse pass 和代码生成。