StableHLO类型推导：MLIR类型系统推导方法

可视化演示

\u5c55\u793aStableHLO\u7b97\u5b50\u7684\u7c7b\u578b\u63a8\u5bfc\u8fc7\u7a0b\u4e0e\u5e7f\u64ad\u8bed\u4e49\uff0c\u5305\u542b\u5f20\u91cf\u5f62\u72b6\u9759\u6001\u63a8\u65ad\u4e0eNumPy\u5e7f\u64ad\u89c4\u5219\u7684\u53ef\u89c6\u5316\u6f14\u793a\u3002

类型系统基础：溶液的双重身份

小雨来到试剂柜前，发现每一瓶溶液都贴着两张标签：一张写「体积规格」（比如 2×3，表示长方体试剂瓶的底面是 2 行 3 列），另一张写「浓度类型」（比如 f32，表示用 32 位浮点记录浓度）。这两张标签共同决定了这瓶溶液的身份。张量在 StableHLO 里也有这样两张标签，合在一起就是它的「身份证」。

这张身份证在编译器里的正式写法是 tensor<2x3xf32>：尖括号里前半段 2x3 是形状（shape），后半段 f32 是元素类型（element type）。任何 StableHLO 算子一开始接触张量，就要先核验这张身份证。形状的每个维度既可以是静态的（写死数字），也可以是动态的（用 ? 表示，留到运行时才知道）；元素类型则决定每个数值占用多少存储空间、如何参与运算。两者组合成类型，这是 MLIR 所有验证与推导的地基。

静态维度好比标签上写死「精确体积 2×3」，动态维度则像「加蒸馏水至 1000mL 刻度」——前者立刻可知，后者要等操作进行到那一步才能确定。张量带着这张或清晰或模糊的标签走进算子，接下来就看实验手册如何规定动作了。

算子契约：实验手册规定动作

实验室里有一排「实验手册」，每本都写死了三栏：材料、条件、产物。比如 add 手册——材料：A、B 两瓶溶液；条件：浓度类型必须一致；产物：留白，待推断。而 broadcast_in_dim 手册更讲究，材料栏之外还有一行小字：「注明每个输入维度搬到输出的第几位」，这就是属性（attribute）。每本手册都构成一份「契约」，规定操作成立的全部前提。

在 StableHLO 中，每个算子都遵循这样的契约：声明自己的输入操作数（operands）、属性（attributes）、区域（regions）和输出结果（results）。输出类型往往由推导完成，但契约里写明了推导所需的全部前提。属性是推导形状的关键「实验条件」：stablehlo.broadcast_in_dim 用 broadcast_dimensions 属性说明输入的每一维对应输出形状的哪一维；stablehlo.dot_general 用 batch_dims 与 contracting_dims 指出批量维度与收缩维度；stablehlo.reduce 用 dimensions 属性列出哪些维度要被归约掉。

区域（region）则是手册里的「循环小节」，像 while 算子携带一个 body 区域，区域内部还有自己的类型约束——循环携带的张量每次迭代类型必须保持稳定。这些前提合在一起，就是推断产物规格的「充分条件」。手册把材料与条件写清楚，产物栏留给统一机制去推导——这就是 StableHLO 算子的设计哲学：不重复描述已知，把类型推导交给下一节的「推演仪器」。

MLIR推导机制：化学家工作台

小雨桌上有一台「推演仪器」：把输入试管插进左槽，翻到手册对应细则，摇动把手，右槽就会掉出一张写着产物规格的纸片。这台仪器不是水晶球，而是一套严谨的机械传动：读类型、套规则、写结论。在 MLIR 中，这套传动就是类型推导机制，由两类接口驱动：InferTypeOpInterface 提供静态方法 inferReturnTypes，直接返回完整的结果类型列表；InferShapedTypeOpInterface 提供 inferReturnTypeComponents，返回更灵活的「形状组件」（ShapedTypeComponents）。

形状组件的价值在于「可以不把话说死」：当某些输出维度是动态的，组件只记录已知的维度，对未知维度保留约束而不填具体数字。推导在什么时候触发？最典型的是算子被构造时（builder 调用）和 IR 验证阶段（verifier 调用）。如果推理失败，接口通过 emitError 产生一条包含算子位置、输入类型和冲突细节的诊断消息，编译随即停止——化学家发现配方有冲突，绝不会硬做实验。

有了这台仪器，小雨不用每次「把实验做完」才知道产物规格——编译器同样不用「把模型跑完」就知道中间张量的类型。接下来我们看看仪器里最常用的四套「配方」。

典型推导规则：四大试管配方

配方一：逐元素同型。 add 这类逐元素算子，像两瓶同规格溶液直接混合。条件是两瓶溶液规格完全一致，产物规格也随之确定。

配方二：广播对齐。 当两个输入形状不同但可兼容时，广播规则把「小溶液」按维度扩展成「大网格」再混合。维度的对齐从右向左，大小为 1 的维度自动扩展到对方大小。

配方三：矩阵乘维度。 矩阵乘像「串联反应器」：A 管出口的 K 号管径必须接上 B 管入口的 K 号管径。

配方四：归约去维度。 reduce 像「把指定行的数据汇总成一行」，输出形状就是删掉这些维度。

这四大配方覆盖了 StableHLO 绝大多数算子的形状推理。把它们固化成公式后，就能交给一个统一的函数去执行——这就是下一节的形状函数。

形状函数体系：配方背后的计算器

手册角落挂着一台小小的「计算器」，上面刻着配方对应的数学式。小雨只要把输入规格拨进去，计算器就咔哒咔哒吐出产物形状。这台计算器，就是形状函数（shape function）。StableHLO 为每个算子导出一个 compute_shape_fn：当算子无法直接给出输出类型时，MLIR 的 shape function 机制把「如何算形状」从算子定义中分离出来，做成可复用、可测试的纯函数。

形状函数支持符号化推导：当输入包含动态维度（?），它把未知维度当作符号变量参与运算，产出一个符号表达式而不是立即报错。符号化推导让编译器可以「纸上谈兵」，把不确定的维度作为变量保留，等到后续 pass 再做具体化（specialization）。对于 while 循环，类型推导还施加一条「循环不变约束」——每次迭代携带的（carried）张量类型必须与上一次完全一致，否则推导直接拒绝，防止出现越跑形状越乱的失控反应器。

形状函数把「实验操作」抽象成「数学计算」，让蒸馏、配料这些工艺脱离实物，在纸面上先推演一遍。这正是编译器静态分析的浪漫：不用点火，就能知道反应的产物规格。

推导工程价值：为什么重要

知道产物规格后，小雨一口气解决了三件大事：提前预订烧杯（内存）、规划工序（优化）、核对禁配表（校验）。类型推导的工程价值，正是把这三件事从运行时提前到编译期。第一是编译期正确性：形状不匹配、元素类型冲突在编译期就被拦下，像「禁配表」防止把不相容的试剂倒在一起，避免模型跑了一半才崩溃。

第二是优化触发器：算子融合、向量化、重排都需要精确的形状信息。知道每个中间张量的形状，编译器才能判断「这两个算子的布局是否兼容、能否合并」，像工序规划必须知道产物尺寸。第三是内存规划依据：推导输出张量大小后，编译器可以为显存池预留空间。

第四是跨方言稳定 ABI：StableHLO 与 HLO 互相转换时，类型推导保证两侧形状语义一致，像不同实验室共用一套容器规格标准，任何一侧的改动都不会悄悄破坏另一侧。有了推导，编译器不用等运行时才「大吃一惊」——这正是工业级编译器的底气。

实践排查指南：当推导失败时

小雨第一次接 while 循环反应器，仪器吐出一张红条：「循环体携带值形状漂移！」她盯着红条愣了半天，最后乖乖把 body 区域里的 reshape 删掉，故障消失。推导失败并不可怕，可怕的是不会读诊断信息。最常见的失败场景有四类：操作数元素类型不一致（f32 与 s32 混用）、广播维度无法对齐（2x3 与 3x1）、必要属性缺失（broadcast_in_dim 忘写 broadcast_dimensions）、循环携带值形状漂移。

错误消息解读有章可循：inferReturnTypes 产出的诊断消息包含算子位置和冲突细节，例如 "requires the same element type for all operands"。读消息时先定位「哪个算子推导失败」，再看「哪个前提被违反」，最后顺着 IR 回溯到出错的那一层。调试工具链方面，mlir-opt --debug 可以逐 pass 观察 IR，-print-ir-after-all 打印每个 pass 之后的 IR，帮助找出类型「从哪一步开始失真」。而 stablehlo-opt -verify-diagnostics 支持在测试文件里用 // expected-error 声明错误位置，把验证诊断写进回归测试，防止「改一个算子弄坏十个下游」。化学家靠红条和检修记录积累经验，编译器工程师靠诊断消息和回归测试一次次加固推导引擎——每一次红条，都是通往更可靠编译器的台阶。

---

AI 应用场景

Python 演示

from itertools import zip_longest  class TensorType:    """模拟 StableHLO 的张量类型：形状 + 元素类型。"""     def __init__(self, shape, elem_type):        self.shape = list(shape)   # 动态维度用 None 表示        self.elem_type = elem_type     def __repr__(self):        dims = []        for d in self.shape:            dims.append(str(d) if d is not None else "?")        return "tensor<{}x{}>".format("x".join(dims), self.elem_type)  def _element_must_match(types, op_name):    first = types[0].elem_type    for t in types[1:]:        if t.elem_type != first:            raise ValueError(                "{} 要求所有输入元素类型一致, 但收到 {} 与 {}".format(                    op_name, first, t.elem_type))  def infer_add(a, b):    """stablehlo.add: 逐元素加法, 输出形状 = 广播(a.shape, b.shape)."""    _element_must_match([a, b], "stablehlo.add")    s1, s2 = list(reversed(a.shape)), list(reversed(b.shape))    out = []    for d1, d2 in zip_longest(s1, s2, fillvalue=1):        if d1 is None or d2 is None:            out.append(None)                 # 动态维度保持动态        elif d1 == 1:            out.append(d2)        elif d2 == 1:            out.append(d1)        elif d1 == d2:            out.append(d1)        else:            raise ValueError("广播失败: 维度 {} 与 {} 不兼容".format(d1, d2))    return TensorType(list(reversed(out)), a.elem_type)  def infer_dot(a, b):    """stablehlo.dot: 批量矩阵乘 [...]MxK x [...]KxN -> [...]MxN."""    _element_must_match([a, b], "stablehlo.dot")    m, k1 = a.shape[-2], a.shape[-1]    k2, n = b.shape[-2], b.shape[-1]    if k1 != k2:        raise ValueError("矩阵乘契约维度不匹配: {} vs {}".format(k1, k2))    batch = infer_add(TensorType(a.shape[:-2], a.elem_type),                      TensorType(b.shape[:-2], b.elem_type)).shape    return TensorType(batch + [m, n], a.elem_type)  def infer_reduce(t, dims):    """stablehlo.reduce: 删除归约维后的形状."""    kept = [d for i, d in enumerate(t.shape) if i not in dims]    return TensorType(kept, t.elem_type)  def demo_failure():    a = TensorType([2, 3], "f32")    b = TensorType([3, 1], "f32")    try:        infer_add(a, b)    except ValueError as ex:        print("推导失败:", ex)  def build_demo_graph():    """构造一条 mini StableHLO 数据流图: add -> dot -> reduce."""    x = TensorType([2, 3], "f32")    y = TensorType([3], "f32")    w = TensorType([3, 4], "f32")    add_out = infer_add(x, y)    dot_out = infer_dot(add_out, w)    red_out = infer_reduce(dot_out, [1])    print("add   :", add_out)    print("dot   :", dot_out)    print("reduce:", red_out)    demo_failure()  if __name__ == "__main__":    build_demo_graph()

1. class TensorType用 shape 列表和 elem_type 两个字段模拟张量的「身份证」，动态维度用 None 表示。
2. infer_add(a, b)实现逐元素算子的类型推导：先校验元素类型一致，再把两个形状右对齐逐维做广播。
3. zip_longest(s1, s2, fillvalue=1)右对齐的关键：短形状缺失的前导维度补 1，对应广播规则中「视为 1」。
4. out.append(None)任何一侧是动态维度时，输出该维保持动态，不强行填死。
5. raise ValueError(...)维度既不相等又非 1 时抛出推导失败，模拟 MLIR 的 emitError。
6. infer_dot(a, b)矩阵乘推导：校验内维 K 对齐，输出末两维取 M 与 N，批量维复用广播推理。
7. infer_reduce(t, dims)归约推导：按编号删除指定维度，剩余维度顺序不变。
8. demo_failure()演示 2x3 与 3x1 的广播冲突，说明「非 1 且不等」的维度是推导失败的重灾区。

常见误区

面试题