8 · Grappler、XLA 与 MLIR：图如何被重写和降低

优化发生在多个 IR 层

TensorFlow 不是把 Python 图直接交给一个“编译器”。同一模型可能经过多次表示转换：

Python / FuncGraph
  → GraphDef / FunctionDef
  → Grappler graph rewrites
  → TF Graph / TF dialect (MLIR)
  → StableHLO / MHLO / HLO
  → backend executable

每一层都可能丢失或重编码信息。资源变量、控制依赖、形状、多输出、函数库和 side effects 是最容易在转换边界出问题的对象。

Grappler：图级、运行时友好的重写

tensorflow/core/grappler 的 optimizer 以 GraphDef 为主要输入，包含 constant folding、layout、memory、function、dependency、auto mixed precision 等优化。Meta optimizer 根据配置选择和排序多个 optimizer，并在变换前后维护 graph properties、cost model 和 function library。

Grappler 的工程特点是“可插拔但有约束”：optimizer 可以增删节点和属性，但必须保证 graph validity、控制依赖和 output semantics。很多优化不是纯局部 pattern rewrite，而是需要推断 shape、成本或跨函数边界。

tf.function(jit_compile=True) 的边界

tf.function 先完成 tracing，得到 ConcreteFunction；jit_compile=True 再把可编译区域交给 XLA。并非所有 op、字符串/资源副作用和动态行为都能进入同一 XLA cluster，所以编译路径可能是“部分 cluster + 外部 TensorFlow op”的混合图。

MLIR：把方言当成阶段合同

TensorFlow 的 MLIR 代码位于 tensorflow/compiler/mlir 和 tensorflow/core/ir。TF dialect 仍保留 TensorFlow 语义；shape inference、canonicalization、control-flow、resource lifting 等 pass 逐步降低到更接近硬件的表示。每个 pass 的价值不只在于“优化多少”，还在于把下一阶段需要的约束显式化。

阅读 pass 时抓三个问题：输入/输出 dialect 是什么？它允许改动哪些 side effect？失败时是拒绝整个函数、回退到 TensorFlow runtime，还是生成不完全 lowered 的 IR？

HLO 与设备代码

XLA 最终围绕 HLO/StableHLO 做算子融合、布局、内存规划和后端 lowering。GPU、CPU、TPU 对同一个 HLO 的成本模型和 codegen 不同，故“图被融合”不是跨硬件恒定的结果。性能调试应同时收集 cluster 形状、编译时间、kernel launch 数和设备内存。

源码锚点

• tensorflow/core/grappler/optimizers/meta_optimizer.cc：optimizer 编排。
• tensorflow/core/grappler/graph_properties/：shape/cost 属性推导。
• tensorflow/compiler/mlir/tensorflow/：TensorFlow dialect 与 pass。
• tensorflow/compiler/mlir/tf2xla/：TF 到 XLA 的 lowering。
• tensorflow/compiler/xla/：HLO、backend、executable 与编译缓存。
• tensorflow/core/ir/：TensorFlow IR 与 GraphDef import/export。

观察编译行为

使用 f.experimental_get_compiler_ir(...)(stage='hlo') 可以把用户函数与中间表示联系起来；配合 profiler 看 kernel fusion 和 launch 数，才能判断编译是否真正改善了瓶颈。