图优化：形状推理Shape Inference

可视化演示

演示图优化阶段Shape Inference沿算子链静态推导各层输出形状，为内存分配、算子融合与代码生成提供精确依据。

概念与意义

老木匠的案头功夫，从来不在锯子上。订单进门，他先铺开一张空白图纸，把每一块原木的来料尺寸、每一道工序后的半成品尺寸都写上去。柜门要四块板，每块多厚；抽屉要三条轨，每根多长——这些数字在机器启动前就定下了。图纸上多写一行字，车间里就少一次手忙脚乱的测量。

深度学习模型走进编译器时也是如此。模型先被翻译成一张计算图：节点是算子，边是流动的张量。张量的“形体”由形状（shape）描述——一个维度元组，比如 [B, 3, H, W] 依次是批次、通道、高、宽。形状推理（Shape Inference）就是编译期的一项静态分析：不真正执行任何算子，只根据算子的语义与输入形状，推算出每个输出的形状。

在 XLA 与 PyTorch 2.0 的图优化阶段，形状推理紧跟在图构建之后、内存规划与内核生成之前。它输出的“形状事实表”被下游反复消费：内存分配器按形状算缓冲大小，算子融合器按形状判断合并可能性，代码生成器按尺寸特化循环。若形状只能等到运行时才知道，编译器就只能做最保守的安排——相当于每块料都现量现裁，工效大打折扣。正因如此，形状推理被视为编译优化链条上“承上启下”的关键前置步骤。

机器还没转，但老师的图纸已经写满数字。形状推理的价值正在于此：用一次编译期的“纸上标注”，换取运行期每一步的从容。图纸上多一行尺寸，车间里就少一次等待——这就是为什么它如此重要。

静态推导原理

工坊里每台机器都钉着铭牌：进料尺寸、出料尺寸、吃料深度，写得清清楚楚。老师傅按工序从第一台走到最后一台，翻一块铭牌，记一笔台账，把每个中间件的尺寸顺次填进表格。这趟“翻铭牌、记台账”的流程，就是形状推理的日常。

编译器里的每类算子同样自带“铭牌”——算子元数据中注册了一个形状规则（inference rule），描述输入形状如何映射到输出形状。推理器拿到图的输入形状后，按拓扑序遍历计算图（保证每个算子在推理时，它的输入形状都已被上游算好）：遇到算子就取出它的规则，代入输入形状，算出输出形状，写进形状环境（shape env）——一张“张量名 → 形状”的映射表，供下游算子继续查用。

以卷积为例，其规则由填充（padding）p、卷积核（kernel）k 与步幅（stride）s 决定。三个参数都是编译期常量时，只要输入高宽已知，输出高宽立刻确定。若某算子的规则依赖运行时数据——比如“输出长度等于输入序列中非零元素个数”——规则返回未知（Unknown），该输出被标记为动态，不再参与静态推导。整个遍历自底向上、层层推进，形状像水一样从输入端流向输出端，直到整张图都被“标注”或标记为未知。

老师傅一路翻到底，台账越记越厚，图纸越来越清晰；遇到没有铭牌的机器，就在台账上画个问号。这个问号，就是编译器里的 Unknown——它预示了后面的工序没法全按图纸走，得留一手。

符号形状表达

有的客户不把话说死：“柜子要三米到四米，你们看着办。”老师傅没法标一个死数字，只能在图纸上写一个大大的 H（待定），再补一行小字：H 必须是 4 的倍数，滑轨才装得上。这样的图纸依然能用，只是多了一层“带约束的弹性”。

模型部署里也常有“话不说死”的维度：批次大小随请求量浮动，序列长度随句子变化。若把所有维度都写成整数，batch 从 8 变成 16 就要重新编译一次，代价高昂。符号形状（Symbolic Shape）允许这类维度用符号变量表示——形状里的 S_b、S_h、S_w 不是具体数字，而是一个“待定但受约束”的记号。

符号变量照样参与传播：输入 [S_b, 3, S_h, S_w] 经过卷积得到 [S_b, 16, S_h-2, S_w-2]。遇到除以步幅时会引入整除要求，编译器便收集约束 S_h % 4 == 0。符号表达式既能像普通形状一样参与内存计算（表达式可化简），又给运行时留下适度的自由度——介于“完全静态”和“完全动态”之间，是工程上的经典折中。

图纸上的“H（待定）”配上“4 的倍数”的约束小字，就是符号形状的写照：自由要给，但不能无边无际。约束越清晰，编译器的优化空间就越大。

编译优化收益

尺寸标全的图纸，能同时喂饱车间里的三个角色：仓管按最大尺寸预订货架，裁切师傅按余料规划顺序，打磨工序直接并进切割机——半成品不必搬来搬去。尺寸，是车间运转的公共语言。

编译器里三位“角色”同样依赖形状。第一位是内存分配器：张量占用的字节数等于元素总数乘以每个元素的字节数，形状已知，分配器就能在运行前精确预留缓冲，甚至让多个中间张量复用同一块内存，而不是运行时“现租”。

第二位是算子融合器：两个算子能否融合，先看中间形状是否兼容。卷积后面跟一个逐元素激活函数，只要形状一致就能合并成一个内核，省去中间张量的写回与读入。第三位是代码生成器：固定形状可以特化内核——展开小循环、固定边界、为特定 batch 挑选最优的 GEMM 实现。再加上运行时形状检查被整个消除（每次调用都要比较尺寸的环节直接省掉），三重收益叠加，正是形状推理带来的性能红利。

图纸齐全，仓库、裁切、合并全部提前安排；形状齐全，内存、融合、特化全部胸有成竹。老师傅照图施工，编译器按形优化——“先标后做”的朴素智慧，在编译的世界里价值千金。

PyTorch 2.0实践

工坊接了不少相似订单。老师傅把验证过的工序图挂在墙上，下次来了同尺寸的活，照着做就行。但每批木材到货，伙计仍要量一下尺寸——万一这批板子宽了两毫米，整套工序就得推翻重来。这道“到货核验”，就是形状守卫。

PyTorch 2.0 的 torch.compile 延续了同样的逻辑。Dynamo 把 Python 函数捕获成 FX 计算图，后端在图上完成形状推理与一系列优化，得到一份可复用的编译产物；产物上挂着一串形状守卫（guards）——对输入形状做出的编译期假设，例如“batch 等于 4 且高度等于 224”。编译产物上还记录了特化后的 TensorType（如 [4, 224, 224, 3]），它是“这张图纸对应哪批料”的精确签名。

下一次调用时，运行时先逐条求值守卫：实测本次输入形状，与假设逐一比对。全部通过，直接复用编译内核，把编译开销摊薄到一次次调用里；任一守卫失败，则触发重编译，新形状对应的产物同样被缓存。这也解释了“第一次慢、后面快”的现象。对于动态形状场景，用户可把部分维度标记为符号，守卫从“等于某值”放宽为“满足符号约束”，从而减少重编译频率。

墙上的工序图是编译产物，到货核验是形状守卫。核验通过则照图施工，核验失败则重新画图。形状推理的质量，决定了图纸被复用的概率——推理与守卫，本就是一枚硬币的两面。

局限与挑战

手艺再好的工坊，也会碰上难缠的单子：客户现场改尺寸，或要求“先看料再定方案”。这时图纸标不全，老师傅只能交代工人边做边量、做到哪算哪。不是老师傅不专业，而是这类活从来就没有“图纸上的答案”。

形状推理也有标不出的形状。最常见的来源是控制流：计算图里出现 if/while，两个分支的输出形状不一致，汇合处填什么？还有数据依赖形状：输出形状由输入张量的内容决定，例如“返回所有非零行”。这类情形在静态阶段无法确定，只能标记为 Unknown。

编译器对 Unknown 的应对是回退（fallback）：放弃本段静态优化，退回解释执行，或在运行时做动态形状检查，甚至为每个新形状重新编译。回退带来性能损失，但保住正确性。另一重挑战是复杂度：符号约束越多，求解与化简越贵，编译本身可能变慢。因此工程实现普遍采用“先静态、再符号、后动态”的渐进策略，并以缓存缓解重编译的代价。

标不全的图纸配上一套“现场测量”的备案，订单照样能交付，只是慢一些。形状推理的局限不是失败——它替编译器划清了“哪些该静态做、哪些必须动态做”的边界，让回退变得有条不紊。图纸上的问号，同样是专业的一部分。

AI 应用场景

Python 演示

"""极简形状推理引擎: 按拓扑序遍历计算图,用算子注册的形状规则静态推断中间张量形状."""  class ShapeEnv:    """形状环境: 一张'张量名 -> 形状'的台账"""     def __init__(self):        self.shape_map = {}     def bind(self, name, shape):        self.shape_map[name] = shape     def lookup(self, name):        return self.shape_map.get(name)  def infer_conv2d(out_ch, kh, kw):    """卷积规则: 高宽按 (H - k + 1) 收缩"""     def _rule(in_shape):        n, _, h, w = in_shape        return (n, out_ch, h - kh + 1, w - kw + 1)     return _rule  def infer_reshape(target):    """reshape 规则: -1 表示按总数与其余维度反推"""     def _rule(in_shape):        if target[0] != -1:            return tuple(target)        total = 1        for d in in_shape:            total *= d        rest = 1        for d in target[1:]:            rest *= d        return (total // rest,) + tuple(target[1:])     return _rule  def infer_linear(out_features):    """全连接规则: 保留前导维, 末维变为 out_features"""     def _rule(in_shape):        return tuple(in_shape[:-1]) + (out_features,)     return _rule  class MiniCompiler:    """迷你编译器: 注册规则 + 编译图并执行形状推理"""     def __init__(self):        self.rules = {}     def register(self, op_name, rule_factory):        self.rules[op_name] = rule_factory     def compile_graph(self, graph, input_shape):        env = ShapeEnv()        env.bind(graph["input"], input_shape)        for node in graph["nodes"]:            rule = self.rules[node["op"]](*node["params"])            in_shape = env.lookup(node["inputs"][0])            out_shape = rule(in_shape)            env.bind(node["output"], out_shape)            print(node["name"], "->", out_shape, "已标注")        return env  def main():    compiler = MiniCompiler()    compiler.register("conv2d", infer_conv2d)    compiler.register("reshape", infer_reshape)    compiler.register("linear", infer_linear)    graph = {        "input": "img",        "nodes": [            {"name": "conv1", "op": "conv2d", "params": (16, 3, 3),             "inputs": ["img"], "output": "conv1"},            {"name": "flatten", "op": "reshape", "params": (-1, 16),             "inputs": ["conv1"], "output": "flat"},            {"name": "fc", "op": "linear", "params": (10,),             "inputs": ["flat"], "output": "fc"},        ],    }    print("=== 形状推理: 开工前标注尺寸 ===")    compiler.compile_graph(graph, (1, 3, 32, 32))  if __name__ == "__main__":    main()

1. self.shape_map = {}台账初始化：张量名到形状的映射，起步为空。
2. def register(self, op_name, rule_factory):给算子贴铭牌：把规则注册进编译器。
3. rule = self.rules[node["op"]](*node["params"])查铭牌取规则，并注入卷积核尺寸等参数。
4. in_shape = env.lookup(node["inputs"][0])从台账读出输入形状，作为推导原料。
5. out_shape = rule(in_shape)套用规则算出输出形状，即铭牌上的进料-出料对照。
6. env.bind(node["output"], out_shape)把新形状登记进台账，供下游算子查询。
7. compiler.compile_graph(graph, (1, 3, 32, 32))给出输入形状，启动整张图纸的标注。

常见误区

面试题