xtl
Back to writing

/ MLC chentianqi

[MLC-07] Computational Graph Optimization and Pattern Rewriting

Notes for lesson seven: Relax ASTs, visitors, operator pattern rewriting, Linear + ReLU fusion, and mapping graph optimization to TensorIR calls.

1 minMLC · Graph Optimization · Relax · Operator Fusion

前几讲优化单个张量函数，第七讲把视角抬到计算图：怎样通过检查与改写 IRModule 中的表达式树，改变算子组合方式，并把优化后的图重新连接到 TensorIR 实现。

1. 计算图也是 IR#

以一个简单的 Relax 函数为例：

@R.function
def main(x, y):
    with R.dataflow():
        lv0 = R.multiply(x, y)
        gv0 = R.add(lv0, y)
        R.output(gv0)
    return gv0

这段代码背后不是 Python 执行流，而是一棵由 Function、SeqExpr、DataflowBlock、VarBinding 和 Call 节点组成的 AST。图优化 pass 的工作是遍历这些节点，匹配模式并构造等价的新节点。

IRModule
  -> Relax Function
    -> SeqExpr
      -> DataflowBlock
        -> VarBinding
          -> Call(op, args)

理解这层对象模型后，“融合算子”不再是神秘操作：它就是把一组相关 binding 替换为一个语义等价的 call 或子函数。

2. 模式匹配与改写#

课程中的一个基础示例是：

multiply(x, y)
      |
add(result, y)
      |
      v
ewise_fma(x, y, y)

一个 rewrite pass 通常包含：

  1. 访问当前 Call 节点；
  2. 判断 op、参数个数和数据依赖是否满足模式；
  3. 确认中间值没有不安全的额外使用；
  4. 生成替代表达式；
  5. 让 mutator 重建外层函数和模块。

这里最重要的是语义条件，而不是字符串匹配。只有当中间值的使用关系、shape、dtype 和副作用约束满足时，替换才安全。

3. Linear + ReLU 融合#

典型图模式为：

matmul(x, w)
    -> add(bias)
    -> relu

可被改写为：

linear_relu(x, w, bias)

融合可能减少：

  • kernel launch 数量；
  • 中间张量的分配与写回；
  • 全局内存往返；
  • 调度器需要管理的独立执行单元。

但融合也有边界：如果中间 matmul 结果还被另一个消费者使用，直接删除它就会错误；如果某个后端已有更优的 unfused library path，也不应盲目融合。

4. 为什么要创建子函数#

融合后的计算不一定立即降成一个内联表达式。课程强调可以把匹配到的子图抽成一个新函数：

main graph
  -> call fused_linear_relu(...)
 
fused_linear_relu
  -> TensorIR implementation or generated code

这样能保留图层的模块边界，也让后续 kernel schedule、自动调优和 target-specific lowering 有明确的承载位置。

5. 映射到 TensorIR calls#

图优化完成后，需要把高层图中的 fused call 映射到具体 TensorIR 函数：

Relax fused operator
      |
      | call_tir
      v
TensorIR primitive function
      |
      v
target code or runtime library

这再次体现 MLC 的分层：图 pass 决定组合关系，TensorIR pass 决定底层实现。两者可以独立迭代，也可以在 cost model 或 target capability 的驱动下协同选择。

6. 对 LLM 图优化的联想#

LLM 的图优化比 linear + relu 复杂得多，但思路相同：识别可融合的 RMSNorm、QKV projection、RoPE、attention 或 MLP 子图，并确保 shape、并行切分、缓存副作用和数值语义不被破坏。

能融合，不等于应该融合。
 
还要看：目标硬件、动态 shape、量化边界、并行通信、可复用库和调试可观测性。

7. 本讲总结#

图优化 = 在高层 IR 中匹配语义模式并改写组合结构。
算子优化 = 在低层 IR 中变换循环、布局、内存与硬件映射。
 
二者通过 call_tir / IRModule 连接，形成端到端编译链路。

参考#