想象你面前摆着一本厚厚的手写菜谱。有的页写“加盐少许”,有的页写“油温适当”。于是同一道菜,A 厨师做出来偏咸,B 厨师做出来夹生——因为“少许”和“适当”没有精确的定义。现在把菜谱换成 AI 编译器里的算子:add、convolution、dot_general 这些算子就是一道道的“菜”,而“算子语义”就是关于每道菜最精确的规定——食材是什么品种、每份几克(形状)、按什么顺序、用什么火候处理。StableHLO(Stable High-Level Operations)正是一本由全行业共同维护的《标准菜谱集》:它把 AI 模型中的高级算子用精确、稳定、可移植的语义固定下来。不管是 TensorFlow 训练出的模型,还是 JAX 写出的程序,只要翻译成这本菜谱,任何一家“厨房”(编译器后端)都能照着做出同样味道的菜。这篇文章,就带你翻开这本菜谱,看懂它为什么这么写、怎么写才对。
理解算子语义,是 AI 编译器工程师的“内力”。它能帮你快速定位模型兼容性问题、设计可靠的降级与优化 pass,是通往 StableHLO 贡献者与编译器架构师之路的必经关卡。
目录
可视化演示
展示StableHLO算子语义中逐元素计算、归约聚合与广播张量构造规则,交互切换查看算子形式化定义。
本节配有交互动效(进入视口后自动加载,约 3KB)。
逐元素:add/multiply语义 对每个元素独立计算无依赖 输入输出shape保持一致 out[i]=a[i]+b[i] 逐元素语义可并行优化
算子语义基础:把“少许”变成“克数”
一百多年前,法国厨神埃科菲干了一件“得罪全行业”的事:他把菜谱里的“加盐少许”“文火慢炖”“适量黄油”统统改写成“盐 3 克”“95 摄氏度炖 40 分钟”“黄油 20 克”。同行骂他死板,食客却因此受益——从此在巴黎任意一家餐厅点同一道菜,端上来的味道几乎一致。StableHLO 就是 AI 编译器界的埃科菲。
在机器学习编译器里,“菜”是高层的算子:加法、矩阵乘、卷积、归约……“菜谱”则是这些算子的语义(Op Semantics)。所谓算子语义,就是这个算子在任意合法输入上的完整行为契约:输入必须有什么形状和类型,输出形状如何由输入推出,以及数值上必须遵守哪些规则(舍入、溢出、NaN 传播)。StableHLO 最初从 XLA 的 HLO 中抽取稳定算子而来,如今是 TensorFlow、JAX、PyTorch 共同认可的可移植中间表示。它把每个算子写成一条精确的“标准配方”,对应如下泛型签名:
正因为有了这本精确到“克数”的手册,各家编译器后端才能放心大胆地把算子揉碎重组、做等价变换,而不必担心打翻哪口锅。语义精确,是一切优化的地基。
- 算子应用Op(……) 表示一次算子调用,比如 add(A, B) 或 dot_general(A, B)。
- 操作数x_i 是参与计算的张量,来自前序算子或常量的结果。
- 类型签名Type(x_i) 表示第 i 个输入的类型,包括形状与元素类型。
- 输出确定箭头的含义是语义保证——只要输入合法,输出类型就被唯一确定。
形状与类型系统:食材的容器与品种
菜谱里写“取鸡蛋 3 枚”,可厨房里盛鸡蛋的托盘有一排排格子:2 格一排、3 格一排。张量的形状(shape)就是这种“托盘”:一个 tensor\<2x3xf32\> 表示这块数据排成 2 行 3 列,每格是一个 32 位浮点数。
在 StableHLO 语义中,每个算子的输入输出都标注了“托盘”与“品种”。品种就是元素类型:f32、f64、i32、i1(布尔)等。类型与形状合起来构成张量类型(TensorType),它是语义公式里的 Type(x)。形状又分为静态与动态:静态形状的每一维都是确定整数,编译期就能分配内存;动态形状用问号表示未知维度,如 tensor\<?x3xf32\>,常见于批次大小未知的推理场景;有界形状则给维度一个上限,让编译器既保留灵活性又能提前规划内存。
逐元素算子(如 add)对两个输入有广播规则:维度对齐时,相等可直接用;某边为 1 可扩展;否则非法。
托盘和品种搞清楚了,后厨才知道每道菜准备多大的盘、装什么料。StableHLO 的形状与类型系统,就是这套全行业通用的“餐具标准”。
- 逐维对齐broadcast(d1, d2) 比较两个张量在同一维度上的尺寸。
- 尺寸相等相等则直接采用该尺寸作为公共维度。
- 单边扩展某维为 1 时,把 1 复制扩展到另一方的尺寸。
- 冲突失败其他组合语义非法,编译器必须报错而不是猜测。
核心算子族:招牌菜的配方
标准菜谱集里有几百道菜,但后厨翻得最多的永远是那几道招牌:大火快炒、文火慢炖、把食材剁碎后搅匀。
对应到 StableHLO,招牌算子族有三类。逐元素算子(add、multiply、relu 等)对张量的每个位置独立计算,输出与输入形状对齐;归约类算子(reduce、reduce_window 等)把张量沿某些维度压缩,输出比输入“瘦”一圈;还有一类最常用的“剁碎搅匀”——dot_general,它是通用矩阵乘的抽象,允许用户指定一组“契约维度”:两个输入在这些维度上做乘法并累加消去,其余维度原样保留成输出。一个 2x3 的矩阵与 3x4 的矩阵相乘,契约维度是 A 的第 1 维与 B 的第 0 维,长度均为 3,消去后得到 2x4。
知道招牌菜怎么做,后厨才知道备什么样的料。理解了 add、reduce、dot_general 的语义,再看 StableHLO 程序就像翻菜谱目录:每道菜在做什么、盘子(形状)如何变化,一目了然。
- 输入张量A 和 B 是两个多维张量,每个维度有各自的含义。
- 契约维度调用时指定一组“相乘后求和消去”的维度,即 k_1……k_r。
- 内积求和对契约维度的全部下标做多重累加。
- 输出形状剩余非契约维度按序拼接,构成输出张量的形状。
MLIR 方言集成:菜谱的统一装帧
光有菜谱内容还不够,各家出版社得用统一的排版、统一的章节编号,菜谱才能在书架间流转。StableHLO 的“装帧规范”就是 MLIR 的方言(Dialect)机制。
在 MLIR 中,一种方言就是一套算子、类型、属性的集合。StableHLO 作为 dialect 注册进 MLIR,每个算子是一个继承 mlir::Op 的 C++ 类,自带创建(build)、打印(print)、解析(parse)与验证(verify)能力。IR 的文本形式是这本“菜谱”的印刷版式:
除了算子本体,方言还定义了自定义类型(如表示有界形状的 BoundedTensorType、token 类型)以及接口(Interface)。算子实现 InferTypeOpInterface 后,形状推导逻辑得以复用;方言注册时声明允许的属性与区域,让 IR 在解析阶段就完成第一道安检。统一的装帧让“菜谱”能在不同出版社之间流转,MLIR 的方言与接口机制,则让 StableHLO 这本书既能被人读,也能被机器安全地拆解、重组、翻译。
- 结果名\%y 是这条指令定义的 SSA 值,后续算子可直接引用。
- 算子标识stablehlo.add 表示“stablehlo 方言里的 add 算子”。
- 操作数列\%a、\%b 是输入张量,冒号前还可携带属性字典(如 dot_general 的契约维度配置)。
- 类型签名冒号后的类型串是验证器检查的依据,语义不符就拒绝编译。
降级与合法化:把菜谱翻译成后厨手册
标准菜谱是给所有餐厅看的,但真正动手时,每家后厨有自己的灶台、锅具和上菜流程。把“标准菜谱”翻译成“某家后厨实操手册”的过程,在编译器里叫降级(Lowering)。
StableHLO 是高级方言,可移植、语义清晰,但硬件不认识它。编译器要做的是沿一条方言链逐级向下:先经过 CHLO(兼容层)把隐式语义显式化——比如把隐式广播变成显式的 broadcast_in_dim、把动态形状操作变成显式版本;再降到 MHLO;继而映射到 Linalg 的循环与访存结构;最后生成 LLVM/目标指令。
后厨手册没有玄学,每一条都是从标准菜谱一步步推出来的。降级链路正是 StableHLO“可移植”承诺的兑现方式:前端只写一遍,后端各取所需。
- 源头程序P_StableHLO 是各前端框架共同产出的可移植“标准菜谱”程序。
- 语义显式化第一站 CHLO/MHLO,把广播等隐式行为翻译成显式算子。
- 结构映射再降到 Linalg,把张量算子映射成循环加访存的底层结构。
- 终极代码最终生成 LLVM 或专用后端的指令。
版本兼容保证:菜谱的修订制度
菜谱再权威也会出新版:修正笔误、补充新菜。关键是修订规则要透明——老读者手里的旧版书不能作废,新书也不能让老厨师看不懂。
StableHLO 为此设计了一套版本机制。程序在序列化时会带上版本号,编译工具链(consumer)在加载时检查自己能否安全消费这个版本(producer)。规则用一句话说:新编译器保证能读旧程序(向后兼容),旧编译器也保证能读比自己新一个版本的程序(前向兼容)。
序列化格式同样重要:StableHLO 程序以可移植二进制格式保存,内嵌版本号与目标版本,方便工具链在不丢失语义的前提下跨版本搬运。新增算子、行为澄清、格式演进,都分门别类落在版本策略里。版本制度让这本标准菜谱“活了多年还能继续加印”——新增菜式不影响老菜谱,老厨师也能看懂八成新菜,AI 生态因此稳步演进,而不是每次升级都推倒重来。
- 两个角色V_p 是产生程序的框架版本(写菜谱的人),V_c 是编译程序的工具链版本(读菜谱的人)。
- 向后兼容只要 V_c >= V_p,新编译器必定能读旧程序。
- 前向兼容编译器即使比程序旧一个版本,也保证能读懂。
- 兼容窗口两者合起来构成生态契约,越界即拒绝加载。
优化编译效能:给厨房提效
语义精确不只为了“不出错”,更为了“做得更快”。一个知道菜谱精确步骤的厨师,可以提前备菜、合并步骤、甚至把两道菜一锅出。
编译器利用 StableHLO 语义做三类典型优化。其一,算子融合:把多个连续算子合并成一个 kernel,省去中间张量的分配与读写——因为语义精确,编译器能证明“先加后激活”与“一边加一边激活”结果一致。其二,形状特化:动态形状算子编译成通用内核后,运行时一旦拿到具体形状,就派发到为该形状特化的内核。其三,常量折叠:输入全是常量的算子,编译期直接算出结果。
埃科菲的菜谱之所以伟大,不仅因为标准,更因为标准让厨房分工与提速成为可能。对 AI 编译器而言,算子语义就是这份“标准”,它把正确性与性能焊在了同一条地基上。
---
- 原始模式add 先算出中间结果 z,relu 再读 z,两轮算子、两轮访存。
- 融合条件语义保证逐元素算子可合并时,两个算子变成一个。
- 融合结果生成一个 fused 算子,一轮循环同时完成两步。
- 收益来源省掉中间张量的分配、写入与再次读入,内存流量减半。
AI 应用场景
多前端模型的统一接入
阶段:模型中间表示转换
解决的问题:传统上每个框架各写一条后端适配链路,接入新硬件要同时适配 TF、JAX、PyTorch 三套 IR,工作量大且易出偏差。借助 StableHLO,各框架只需把模型翻译成这本“标准菜谱”,所有支持 StableHLO 的硬件后端一次接入、处处可用。
具体案例:某公司推理引擎要支持 PyTorch 导出的模型在自研 NPU 上运行。团队用 torch-mlir 将模型转为 StableHLO,再经合法化降级到 Linalg 生成 NPU 指令,省去了原先维护三套方言转换的工量,新模型接入时间从数周缩短到两天。
硬件后端的语义一致性验证
阶段:测试与验证
解决的问题:算子实现是否符合语义,肉眼难查。StableHLO 提供参考解释器作为“标准味道”,后端实现可用随机输入对拍验证,把语义偏差消灭在合入之前。
具体案例:编译器团队把 StableHLO 参考解释器当作黄金实现,构造包含动态形状与广播的测试样例,对拍自研 CPU 后端的输出,成功定位了 reduce_window 算子 padding 边界的越界错误。
模型数值一致性排查
阶段:模型部署与调试
解决的问题:同一模型在不同硬件上结果不一致,难以区分是框架差异还是后端 bug。回归到 StableHLO 语义,可以在中间表示层面锁定差异归属。
具体案例:平台工程师发现同一模型在 GPU 与 NPU 上精度差异明显,依据 StableHLO 语义文档中的舍入规则,排查出 NPU 后端把 round-to-nearest-even 与 floor 混用的实现 bug。
Python 演示
# -*- coding: utf-8 -*-"""StableHLO 算子语义小型演示: 形状推导与广播规则.""" class TensorType: """张量类型: 形状列表 + 元素类型, 类比菜谱中的食材容器.""" def __init__(self, shape, elem="f32"): self.shape = list(shape) self.elem = elem def __repr__(self): return "tensor<%sx%s>" % ("x".join(map(str, self.shape)), self.elem) def broadcast_shapes(s1, s2): """按 StableHLO 广播语义计算公共形状.""" out = [] for i in range(max(len(s1), len(s2))): d1 = s1[len(s1) - 1 - i] if i < len(s1) else 1 d2 = s2[len(s2) - 1 - i] if i < len(s2) else 1 if d1 == d2: out.append(d1) elif d1 == 1: out.append(d2) elif d2 == 1: out.append(d1) else: raise ValueError("incompatible dims %d vs %d" % (d1, d2)) return tuple(reversed(out)) def infer_elementwise(a, b): """逐元素算子(如 stablehlo.add)的形状推导.""" return TensorType(broadcast_shapes(a.shape, b.shape), a.elem) def infer_dot_general(a, b, lhs_contracting, rhs_contracting): """stablehlo.dot_general 形状推导: 契约维度求和消去.""" for la, rb in zip(lhs_contracting, rhs_contracting): if a.shape[la] != b.shape[rb]: raise ValueError("contract dim mismatch (%d vs %d)" % (a.shape[la], b.shape[rb])) out = [d for i, d in enumerate(a.shape) if i not in lhs_contracting] out += [d for i, d in enumerate(b.shape) if i not in rhs_contracting] return TensorType(out, a.elem) if __name__ == "__main__": a = TensorType([2, 3]) b = TensorType([3, 4]) print("dot_general([2x3]x[3x4]) ->", infer_dot_general(a, b, [1], [0])) x = TensorType([1, 3]) y = TensorType([2, 1]) print("add([1x3]+[2x1]) ->", infer_elementwise(x, y)) try: broadcast_shapes([2, 3], [2, 4]) except ValueError as e: print("broadcast_error:", e)
def broadcast_shapes(s1, s2):把广播规则实现成函数,这就是“语义”落到代码的样子。d1 = s1[len(s1) - 1 - i] if i < len(s1) else 1从右往左逐维对齐,维度不足视为 1,隐式广播由此开始。elif d2 == 1:一方为 1 则复制扩展,对应广播三连判断中的“单边扩展”。raise ValueError("incompatible dims %d vs %d" % (d1, d2))语义非法时报错,对应编译器拒绝编译而非猜测。def infer_dot_general(a, b, lhs_contracting, rhs_contracting):dot_general 要求两侧契约维度成对出现且长度相等。out = [d for i, d in enumerate(a.shape) if i not in lhs_contracting]输出形状 = 非契约维度按序拼接。print("dot_general([2x3]x[3x4]) ->", infer_dot_general(a, b, [1], [0]))可执行入口,契约维度消去后得到 [2, 4]。if __name__ == "__main__":保证模块可被导入复用,也可直接运行验证。
常见误区
误区“算子语义不就是给程序员看的说明书吗?我实现算子时只要输入输出形状对上,数值差不多就行了吧?”
正解语义是编译器与硬件之间的法律契约。StableHLO 规范把舍入方式、溢出行为、NaN 传播、并行不变量都写死。后端实现若“差不多”,会造成相同模型在不同设备上数值不一致,甚至触发未定义行为。语义精确到可以直接写编译器验证器(verifier),任何偏离都会被机器抓住。
为何常见:平时写业务代码习惯了“文档仅供参考”,没有意识到编译器会用语义做等价变换与形式化证明。
误区“StableHLO 既然是中间表示,那 GPU/TPU 驱动里应该有对应的算子实现吧?”
正解StableHLO 是不可执行的高级方言,必须先降级到 MHLO/Linalg 等低层表示,再做循环级优化与代码生成。绝大多数硬件后端根本不认识 stablehlo.add,它们消费的是 LLVM 或专用 IR。把 StableHLO 理解成菜谱而不是菜本身。
为何常见:受“IR 直接执行”的直觉误导,忽略了编译器多级降级的存在。
误区“动态形状在 StableHLO 里能用,但既然它性能差,应该是个快要被淘汰的特性吧?”
正解动态形状是 StableHLO 的头等公民,它支持用问号表示未知维度,还有带上限的有界形状。编译期会尽量对动态形状做特化,实在无法静态化才生成带运行时检查的通用代码。动态形状解决的是“编译期确实不知道”的问题,与性能取舍是两回事。
为何常见:把“性能差”误解为“设计上不允许”,没有区分表达能力和优化策略。
误区“给 StableHLO 加个 op 就是改改版本号,旧程序反正能跑,不用考虑什么兼容矩阵吧?”
正解StableHLO 有严格的版本策略:新增算子属于哪一档、行为澄清与格式变更落在哪个版本窗口,都要按 producer/consumer 兼容矩阵验证。乱升版本会撕裂整个编译生态。向后兼容加一版本前向兼容是底线,越界必须明确报错。
为何常见:把版本号当成计数器,忽略了它在分布式编译生态中的契约意义。
面试题
Q1. StableHLO 最核心的定位是什么?
- A.一种编译后的硬件指令集
- B.一种端到端的训练框架
- C.一种可移植的算子语义规范← 正确答案
- D.一种模型权重压缩格式
答案:C。StableHLO 是一本“标准菜谱”,用精确语义定义高级算子的行为,让不同框架与硬件后端共享同一种中间表示。它既不是指令集,也不是训练框架或权重格式。
Q2. stablehlo.add 对两个输入张量的形状要求是?
- A.两者形状必须完全一致
- B.元素总个数必须相同
- C.广播后形状对齐即可← 正确答案
- D.第一个维度必须相等
答案:C。逐元素算子遵循广播语义:维度相等或某边为 1 时可扩展对齐。它不要求形状完全一致,也不要求元素总个数或第一个维度相等。
Q3. 动态形状在 StableHLO 中如何表示?
- A.使用问号记录未知维度← 正确答案
- B.把维度数字减一表示未知
- C.形状在运行期自动补全
- D.动态形状完全禁止使用
答案:A。StableHLO 用问号表示未知维度,如 tensor\<?x3xf32\>。动态形状是受支持的头等公民,编译期派发通用内核,运行时结合具体形状再做特化派发。
Q4. StableHLO 在 MLIR 生态中的角色是?
- A.直接生成硬件指令的底层方言
- B.跨框架统一的高层方言层← 正确答案
- C.替代全部前端框架的编程入口
- D.运行期动态解释执行引擎
答案:B。StableHLO 是 MLIR 里的高级可移植方言,位于前端与后端之间。它不直接生成指令,也不替代前端框架,更不是解释器,而是统一的“中间菜谱”。
Q5. dot_general 的契约维度会被如何处理?
- A.原样保留到输出张量中
- B.参与乘积并被求和消去← 正确答案
- C.强制与批次维度合并
- D.自动扩展为不变维度
答案:B。契约维度是 dot_general 中“相乘再累加”的维度,类似矩阵乘的内维 k,计算后被消去;输出只保留非契约维度。
Q6. CHLO 方言在降级链路中的主要职责是?
- A.对算子做常量折叠优化
- B.显式化广播等隐式语义← 正确答案
- C.直接调度底层硬件指令
- D.负责内存分配与复用
答案:B。CHLO 是兼容层,把 StableHLO 中隐式的广播、动态形状行为翻译成显式算子,相当于降级到 MHLO 之前的“翻译准备”。
Q7. 根据兼容窗口,消费者版本 Vc 与生产者版本 Vp 必须满足?
- A.Vc 必须严格大于 Vp
- B.Vc 可以小于 Vp-1
- C.Vp-1 必须小于等于 Vc← 正确答案
- D.Vp 必须始终大于 Vc
答案:C。向后兼容允许 Vc >= Vp,前向兼容允许 Vc = Vp - 1,合并即 Vp - 1 <= Vc。消费者不能比生产者老超过一个版本,越界即拒绝加载。
Q8. StableHLO 的设计目标包含哪些?
- A.跨前端框架统一表示← 正确答案
- B.算子语义长期稳定← 正确答案
- C.自动搜索训练超参数
- D.适配多样化硬件后端← 正确答案
答案:ABD。StableHLO 追求可移植与稳定:统一前端表示、语义长期稳定、多硬件适配。超参数搜索是训练框架的职责,不在算子语义范畴内。
Q9. 关于动态形状与有界形状,正确的说法有哪些?
- A.维度可用问号表示未知← 正确答案
- B.有界形状为维度设上限← 正确答案
- C.有界形状可辅助内存规划← 正确答案
- D.动态形状必然大幅变慢
答案:ABC。问号表示未知维度;有界形状(如 1 到 1024)让编译期能按上限规划内存、生成通用代码。动态形状性能取决于特化程度,不是必然变慢。
Q10. 精确的算子语义给编译器带来的契机有哪些?
- A.进行算术等价变换← 正确答案
- B.实施安全的算子融合← 正确答案
- C.验证程序整体正确性← 正确答案
- D.免去类型检查的环节
答案:ABC。语义精确让编译器敢做等价变换、算子融合与正确性验证。类型检查不会被免除,反而因为语义精确可以做得更细、更早。