想象你经营着一座城市的中央厨房,周围每家餐厅都有自己的记账方式:川菜馆写“炝炒”,西餐厅写“sauté”,甜品店写“加热”。当订单汇入中央厨房时,你不可能要求每位厨师学会所有餐厅的暗号,于是你定下规矩:所有订单都必须先改写成同一种标准菜谱卡。这张菜谱卡不关心菜来自哪家店,也不关心最后用燃气灶还是电陶炉,它只关心“用什么食材、做什么动作、产出什么菜”。XLA 里的 HLO IR(High Level Operations,高层操作中间表示)就是这张标准菜谱卡。它把 PyTorch、TensorFlow 等前端框架的计算图,统一翻译成一张与硬件无关的算子图;这张图再经过一系列“主厨审菜谱”式的优化 Pass,最后被翻译成 GPU、TPU 或 CPU 能执行的代码。理解 HLO,你就拿到了 XLA 编译流水线的通用语言。
理解 HLO,等于拿到 AI 编译器流水线的“通用语言”;无论前端是 PyTorch 还是 TensorFlow,后端是 GPU 还是 TPU,你都能读懂模型加速与部署优化的关键链路,这是 AI 编译器工程师的进阶基石。
可视化演示
展示HLO IR以指令为节点操作数为边的计算图结构、HLO指令按算子语义的elementwise/归约/变换等分类体系以及XLA在HLO层完成融合与常量折叠后降级到目标后端的优化流程,交互切换三种模式探索。
本节配有交互动效(进入视口后自动加载,约 3KB)。
'+CAPS.mode1+'
初识HLO
想象你是城市里中央厨房的负责人。周围的餐厅各有各的记账方式:中餐管它叫“炝炒”,西餐叫“sauté”,甜品店写“加热”。当这些订单汇入中央厨房时,你不可能让每一位厨师都去学所有餐厅的暗号,于是你定下一条规矩:无论订单从哪里来,都必须先改写成同一种标准菜谱卡。XLA 里的 HLO(High Level Operations)就是这张标准菜谱卡。
在这张菜谱卡上,每道菜被拆成一个个动作:取料、切配、加热、装盘;对应到计算里,就是一个个 HLO 指令:add、convolution、dot、reshape。原料和半成品在动作之间流动,对应张量数据在节点之间流动。因为数据只能从输入走向输出,不能绕回形成死循环,整张图就是一张有向无环图(DAG)。
HLO 是高层中间表示,它刻意屏蔽了“菜是哪个餐厅点的”和“锅是什么品牌”这些细节。前端只需要把模型翻译成 HLO,后端只需要把 HLO 翻译成自己的指令,两者就可以独立演进。这正是 XLA 能支持多框架、多硬件的原因。
有了这张标准菜谱卡,前厅的点菜方式再怎么变,后厨的烹饪体系都能稳定运转。理解了“HLO 是平台无关的算子图”,你就抓住了 XLA 编译流水线的第一块基石。
- 节点集合 V表示所有 HLO 算子指令,例如 add、convolution、dot、reshape;每个节点相当于菜谱卡上的一个动作。
- 边集合 E表示张量数据流向,从生产者节点连到消费者节点,说明一个算子的输出会成为另一个算子的输入。
- 有向无环数据只能从输入向输出流动,不能形成循环,XLA 因此可以静态地分析、优化并调度整个计算过程。
诞生与流转
客人点餐后,服务员怎么把订单交给中央厨房?中间有一个传菜员:XLA 前端。他是第一个和框架打交道的人。PyTorch 或 TensorFlow 在内存里摆出的计算图,会经过 XLA 前端转换成 HLO 模块。转换过程由 HLO builder 完成,相当于把各家的手写订单誊写成统一菜谱卡。
HLO 不仅存在于内存中,它有自己的协议缓冲(protobuf)定义。菜谱可以被打印、传真、塞进信封送到其他厨房,对应 HloModuleProto 的序列化和反序列化。这种设计让 HLO 可以跨进程、跨语言传递,不再依赖 Python 运行时。可以说,一旦 HLO 生成,框架的使命暂时完成,剩下的优化和编译都由 XLA 接管。
一张菜谱卡从餐厅传到中央厨房,再从中央厨房传到分店,靠的是统一格式。HLO 的 protobuf 序列化就是这条“跨店传真通道”,让模型编译不再被前端语言绑定。
- G\_framework前端框架手里的“手写订单”,比如 PyTorch 的 FX Graph 或 TensorFlow 的 GraphDef。
- Phi(\Phi)XLA 前端执行的整体映射,它把框架算子逐一翻译成 HLO 算子,同时补全形状、布局等信息。
- H\_hlo生成的标准菜谱卡,也就是 HloModule;后续所有优化和代码生成都以它为输入。
HLO内部结构
标准菜谱卡不是一张无边表格,而是一本有层级结构的菜谱集。翻开这本菜谱集,最外层是 HloModule,相当于整本菜谱;每一道菜是一个 HloComputation,相当于一个函数;每一段动作是一个 HloInstruction,相当于一个算子指令。XLA 在编译时以这种三层结构组织所有信息。
每个 HloInstruction 会记录三件核心的事:做什么(opcode)、用谁的结果(operands)、产出什么形状(shape)。shape 还包含 layout 和 element type,告诉编译器数据在内存里怎么摆放。有了这些信息,主厨才能在不真正下锅的情况下,先在纸面上模拟一遍流程,判断哪里能合并、哪里能省略。
有了清晰的层级结构,优化器就像一位能够“隔空审菜谱”的主厨。它不需要真正跑一遍模型,就可以在 HLO 图上做等价变换,这正是编译优化高效的原因。
- Module\_m 个 Computation一本菜谱包含多道菜,对应一个 HLO 模块包含多个函数子图,例如入口函数、While 循环体、条件分支函数。
- Computation 内 n 个 Instruction一道菜由多个步骤组成,对应一个 HloComputation 内的指令序列,指令之间有数据依赖。
- 指令的三元信息每个指令都可看作 (opcode, operands, shape),这决定了它在菜谱中是什么动作、依赖哪些半成品、最终产出什么形状。
优化Pass
菜谱标准化之后,中央厨房的主厨团队开始审菜谱。他们不会改变菜的风味,只会去掉冗余动作:比如明明已经清蒸过的食材,不需要再过一遍油炸。XLA 中承担这些“审菜谱”动作的模块叫 Pass。每个 Pass 都只做一种类型的改进,多个 Pass 串成一条管道。
代数化简把 x+0 改成 x、x×1 改成 x;算子融合把连续几个逐元素操作揉成一个 fused instruction,减少中间张量写回内存,也减少内核启动次数;布局优化则把张量内存排列改成硬件更喜欢的顺序。所有 Pass 都遵守一条铁律:语义保持,优化前后的计算结果在数学上等价。
主厨审菜谱时,每一步改完都要确认口味没变。Pass 管道里的每个环节也必须保证“数学口味”一致,这正是编译器的安全性底线。
- IR\_in进入优化管道的原始 HLO 模块,是刚从前端翻译出来的“初稿菜谱”。
- p\_1 到 p\_n每个 Pass 都是一个独立的等价变换函数,例如代数化简、算子融合、布局优化。
- 复合运算 \circ表示函数按顺序套用,复合结果等于把优化管道完整执行一遍。
Lowering
菜谱定稿后,中央厨房要把菜谱翻译成具体灶台操作手册。GPU 的灶台有数以千计的小火头,TPU 的灶台像一条流水线,CPU 的灶台擅长向量化炒菜。不同硬件各有脾气,因此 XLA 不能只给一份通用菜谱就收工,而是进入 Lowering 阶段:把 HLO 逐步降低到更贴近硬件的表示,再生成机器码。
这个降低不是一步到位的。HLO 会先被翻译到目标后端使用的低层 IR,例如 LLVM IR,或 OpenXLA 中的 Linalg 与 LLVM 组合;随后再进行寄存器分配、指令选择、内存调度等代码生成工作。分层设计让“与硬件无关的优化”和“与硬件相关的优化”各居其位,互不干扰。
不同厨房用不同锅具,但菜谱卡本身不用重写。Lowering 让 XLA 既能共享高层优化,又能针对每种硬件生成最高效的指令,做到“一次写菜谱,处处可烹饪”。
- H\_opt经过 Pass 管道打磨后的最终菜谱,已经完成所有平台无关优化。
- Lowering(\mathcal{L})从高层 IR 到低层 IR 再到机器码的多层翻译,包含指令映射、寄存器分配等步骤。
- T\_target目标 CPU、GPU 或 TPU 能直接执行的原生指令,也就是最终端上桌的菜。
Torch桥接
过去用 PyTorch 训练时,像一位急性子厨师:拿到一条指令就立刻执行一条,反复开关火、拿锅放锅,性能损耗很大。PyTorch 2.0 想要改成“整桌菜统一规划”:torch.compile 把原本迫不及待的 eager 执行改造成先捕图、再优化、后执行。
其中 Dynamo 负责捕获 Python 层计算图,AOTAutograd 负责把反向传播也纳入图中,XLA 后端则把这张图转换成 HLO IR。站在 XLA 视角,Dynamo 和 AOTAutograd 都是“前厅服务”,真正进入中央厨房的第一件标准物品还是 HLO。理解这个桥接关系,才不会把 PyTorch 的 Graph 和 HLO 混为一谈。
在 PyTorch 2.0 的编译链路里,Dynamo 和 AOTAutograd 解决“如何把点菜单写清楚”,而 HLO 解决“如何让所有厨房都能照着做”。这条桥接是理解现代 PyTorch 训练加速的关键。
- G\_torchPyTorch eager 模式下算子逐条执行的计算图,像服务员随手记下的点菜单。
- Dynamo通过字节码分析把 Python 层控制流转换为可静态分析的图,避免 Python 解释器逐算子解释。
- Xi(\Xi)XLA 前端把捕获得的计算图翻译成 HLO 模块,包括算子映射、形状推断和 layout 指定。
调试与实战
再熟练的主厨也会偶尔失手。这时他不会凭空猜,而是拿出保存的菜谱快照,对比某一步改动后味道为什么变了。XLA 也提供了同样的能力:通过命令行开关把 HLO 在每轮 Pass 前后的文本 dump 出来,开发者可以像看 git diff 一样定位问题。
HLO 文本格式是一种人类可读的表示,结构完全对应 HloModule、HloComputation 和 HloInstruction。综合实践里,常用来排查三类问题:Pass 是否把合法计算改错、Lowering 是否映射到错误的硬件指令、融合是否带来内存访问异常。性能调优则常用加速比和内核数量来判断优化是否有效。
HLO dump 是编译器的“后厨监控录像”。真实项目中,遇到性能异常或数值不对时,先看 HLO 前后对比,再定位是优化 Pass 还是硬件映射的问题,往往比盲目改代码高效得多。
- t\_before优化前的模型运行耗时,可以理解为旧菜谱从头做到尾的总时间。
- t\_after经过 HLO 优化和代码生成后的模型运行耗时。
- 比值 speedup大于 1 表示优化有效,等于 2 表示速度快了一倍。
AI 应用场景
大模型推理加速
阶段:模型编译与推理优化
解决的问题:PyTorch eager 执行模式逐算子启动开销大,GPU 利用率低。XLA 通过 HLO 算子融合减少内核启动次数和中间张量访存。
具体案例:某推荐模型在 NVIDIA A100 上推理,原先有 800 多个算子,通过 OpenXLA 将连续 elementwise 算子融合后只剩约 300 个内核,端到端延迟降低 40%。
跨硬件部署
阶段:模型部署与硬件迁移
解决的问题:模型从 GPU 迁移到 TPU 往往需要重写算子层实现。HLO 作为平台无关中间表示,让“写一次前端口,适配多个后端”成为可能。
具体案例:某 NLP 公司在 PyTorch 中训练模型,通过 torch\_xla 将 HLO 模块部署到 TPU v4 进行推理,同时保留一套 GPU 版本,两套后端共享同一份高层优化逻辑。
编译问题定位
阶段:编译调试与性能调优
解决的问题:优化后精度下降或性能不符合预期时,难以判断是前端翻译、优化 Pass 还是代码生成的问题。HLO dump 提供了关键现场证据。
具体案例:工程师开启 --xla_dump_hlo_module 后,对比某几轮 Pass 的 HLO 文本,发现一个融合 Pass 把带有 side-effect 的随机数算子错误合并,从而快速回滚并修正 Pass。
Python 演示
class HloInst: def __init__(self, op, name, operands=(), shape=()): self.op = op self.name = name self.operands = list(operands) self.shape = shape def __repr__(self): return "%s:%s" % (self.name, self.op) class HloComp: def __init__(self, name): self.name = name self.insts = [] def add(self, inst): self.insts.append(inst) return inst def is_elementwise(op): return op in ("add", "mul", "max") def print_hlo(comp): for inst in comp.insts: opers = ", ".join(o.name for o in inst.operands) print("%s = %s(%s) shape=%s" % (inst.name, inst.op, opers, inst.shape)) def fuse_elementwise(comp): old_insts = comp.insts[:] comp.insts = [] group = [] for inst in old_insts: if is_elementwise(inst.op): group.append(inst) continue if len(group) > 1: comp.insts.append(HloInst("fused", comp.name + "_fused", tuple(group))) else: comp.insts.extend(group) group = [] comp.insts.append(inst) if len(group) > 1: comp.insts.append(HloInst("fused", comp.name + "_fused", tuple(group))) else: comp.insts.extend(group) return comp def main(): comp = HloComp("demo") x = comp.add(HloInst("var", "x", shape=(16,))) y = comp.add(HloInst("var", "y", shape=(16,))) z = comp.add(HloInst("var", "z", shape=(16,))) a = comp.add(HloInst("add", "a", (x, y), shape=(16,))) b = comp.add(HloInst("mul", "b", (a, z), shape=(16,))) c = comp.add(HloInst("max", "c", (b, z), shape=(16,))) d = comp.add(HloInst("add", "d", (c, y), shape=(16,))) out = comp.add(HloInst("reduce", "out", (d,), shape=())) print("before fusion:") print_hlo(comp) fuse_elementwise(comp) print("\nafter fusion:") print_hlo(comp) if __name__ == "__main__": main()
class HloInst:表示一个 HLO 指令节点,包含操作类型、操作数和形状,类似菜谱上的一步动作。self.operands = list(operands)保存上游节点引用,即这一步依赖哪些半成品。def is_elementwise(op):判断算子是否为逐元素运算,逐元素操作更容易被融合。def fuse_elementwise(comp):实现简化版 Fusion Pass,把连续逐元素指令合并为一个 fused 指令。comp.insts = []重建指令列表;遇到非逐元素算子时把已累积的分组合并。print_hlo(comp)把 HLO 指令打印成文本,对应 XLA 的 HLO dump 体验。if __name__ == "__main__": main()提供可执行入口,方便读者运行查看融合前后差异。
常见误区
误区我一直以为 torch.compile 生成的前端计算图就是 HLO,毕竟都是节点和边,结构看起来差不多。
正解HLO 是 XLA 自己定义的中间表示,用 protobuf 描述,包含 Layout、算子语义、硬件无关优化等前端图没有的信息。PyTorch FX Graph 只是前端快照,必须经过 XLA 前端翻译才能变成 HLO。类比:客人的原始点菜单和中央厨房的标准菜谱卡不是同一种东西。
为何常见:前端图和 HLO 都由节点和边组成,初学者容易只看形状而忽略各自的语义层次和格式定义。
误区编译器不就应该多堆几个优化 Pass 吗?反正都是让程序更快,越多越强。
正解Pass 数量需要平衡编译时间、调试复杂度和最终性能。每个 Pass 都有分析、变换、验证成本,且 Pass 之间存在依赖和互斥;盲目堆叠可能让编译时间翻倍而性能提升甚微。主厨不会为了菜谱好看而多加十道无效工序。
为何常见:直觉上把优化看作单方向增益,忽略了编译器本身是资源受限的流水线,也忽略了 Pass 之间的潜在冲突。
误区HLO 优化完,下一步应该直接变成 CPU 或 GPU 指令,不然怎么叫编译?
正解XLA 不会一步到位,而是分层 Lowering:HLO 先降到目标后端使用的低层 IR,例如 LLVM IR 或 Linalg,再做指令选择、寄存器分配等代码生成。这样硬件相关优化可以插在中间层,避免每种硬件都重复实现一遍前端优化。
为何常见:把编译理解成简单翻译,忽略了现代编译器是多层抽象栈,每一层都有明确职责。
误区XLA 不是给 GPU 加速用的吗?所以 HLO 只在 GPU 编译路径里存在。
正解HLO 是平台无关的中间表示,XLA 的 CPU、GPU、TPU 后端都共享 HLO 前端。CPU 后端同样通过 HLO 加 LLVM 生成高效代码。菜谱卡不会因为厨房从燃气灶变成电陶炉,就只存在于某一种厨房里。
为何常见:XLA 最初伴随 TPU 和 GPU 的高性能案例被大家熟知,于是把实现场景误认为本质属性。
面试题
Q1. HLO IR 在 XLA 编译流程中所处的位置是?
- A.前端框架的原生计算图
- B.独立于前后端的中间表示← 正确答案
- C.后端硬件可直接执行的机器码
- D.用户手动编写的 Python 脚本
答案:B。HLO 是 XLA 定义的平台无关中间表示,既不属于 PyTorch 等前端框架,也不是后端机器码。它像中央厨房的标准菜谱卡,前厅和后厨都围绕它协作。
Q2. HloModule 与 HloComputation 的关系是?
- A.Module 包含多个 Computation 子图
- B.Computation 是 Module 中的函数级子图← 正确答案
- C.Module 等价于 Computation 的指令列表
- D.二者都等价于单个张量
答案:B。HloModule 是整本菜谱,HloComputation 是其中一道菜或一个函数子图,二者不是等价关系。Module 可以包含多个 Computation,每个 Computation 内含若干 HloInstruction。
Q3. XLA 中算子融合(fusion)主要解决什么问题?
- A.减少算术计算误差
- B.减少内核启动与内存读写← 正确答案
- C.增加模型网络层数
- D.消除反向传播梯度
答案:B。连续逐元素算子如果分开执行,会反复把中间结果写回显存、启动多个内核。融合把它们合并成一个 HLO 指令,相当于把多个烹饪动作合并成一道工序,访存和启动开销都下降。
Q4. 下列哪些属于 HLO 优化 Pass?
- A.代数化简← 正确答案
- B.算子融合← 正确答案
- C.词法分析
- D.布局优化← 正确答案
答案:ABD。HLO 优化 Pass 包括代数化简、算子融合、布局优化等。词法分析是传统编译器前端处理源码的工具,不属于 HLO 层优化。
Q5. XLA 前端从 PyTorch 中主要接收什么?
- A.编译后的二进制文件
- B.高层计算图描述← 正确答案
- C.GPU 驱动指令
- D.模型权重文件
答案:B。XLA 前端接收的是框架产出的高层计算图,而不是二进制、驱动指令或单纯的权重文件。它把这张图翻译成 HLO 后,再交给后续优化和代码生成。
Q6. HLO 文本格式的主要作用是什么?
- A.直接部署到移动端
- B.方便开发者阅读和调试 IR← 正确答案
- C.替代 Python 解释器
- D.加速 GPU 显存的读取
答案:B。HLO 文本格式以人类可读方式展示 Module、Computation 和 Instruction,适合对比 Pass 前后差异。它不是部署格式,也不会直接提升显存读写速度。
Q7. 关于 HLO Lowering 的过程,下列说法正确的是?
- A.HLO 先被优化再逐步降低到后端表示← 正确答案
- B.Lowering 后可以直接跳过代码生成
- C.不同后端有各自的指令映射规则← 正确答案
- D.Lowering 阶段也会考虑硬件特性← 正确答案
答案:ACD。HLO 先经过 Pass 优化,再进行 Lowering,逐层降到目标 IR 并最终代码生成,所以 B 错误。不同 GPU、CPU、TPU 后端确实有各自的映射规则,Lowering 也把硬件特性纳入考量。
Q8. 在 PyTorch 2.0 编译栈中,Dynamo、AOTAutograd、Inductor/XLA 的协作顺序是?
- A.Dynamo 导出图,AOTAutograd 求导,Inductor 生成后端代码← 正确答案
- B.AOTAutograd 求导,Dynamo 导出图,Inductor 生成后端代码
- C.Inductor 生成代码,Dynamo 导出图,AOTAutograd 求导
- D.Dynamo 生成代码,Inductor 求导,AOTAutograd 导出图
答案:A。Dynamo 先捕获 Python 层计算图,AOTAutograd 再处理反向求导,最后 Inductor 或 XLA 后端生成代码。顺序颠倒会导致前端还没有图谱就开始求导或生成代码。
Q9. 一个 HLO 指令节点通常不包含以下哪项?
- A.操作类型
- B.操作数引用
- C.输出形状
- D.训练超参数← 正确答案
答案:D。HloInstruction 记录 opcode、operands、shape 等编译必要信息,而训练超参数属于框架层配置,不会进入 HLO 指令节点。
Q10. 关于 HLO 优化 Pass 的语义保持,正确理解是?
- A.Pass 可以任意改写算子语义以提速
- B.合法 Pass 保证优化前后计算结果等价← 正确答案
- C.代数化简使用数学恒等式等价变换← 正确答案
- D.融合总能逐位保持浮点运算结果
答案:BC。合法 Pass 必须保持语义等价,A 错误;代数化简依赖 x+0=x 等恒等变换,C 正确。融合虽然减少了中间结果,但可能改变浮点运算顺序,因此不能保证逐位一致,D 过于绝对。