维度的透视眼：深入解析 AI 编译器中的 Shape Inference（形状推导）

维度的透视眼深入解析 AI 编译器中的 Shape Inference形状推导在现代 AI 编译器管线中我们已经见证了代码如何从前端被引渡进 MLIRMLIRGen如何通过内联打破边界Inliner以及如何利用声明式规则进行图化简DRR Pattern。但在大模型LLM和生成式 AI 统治的时代AI 编译器面临着一个极其棘手的特征——多变的形状与维度Shapes。用户输入的 Prompt 长度是变长的、KV Cache 是动态膨胀的、分支预测下的 Tensor 维度是时刻在变的。如果编译器在编译期不知道每个算子输出张量的准确形状和数据类型它就无法进行精准的算子融合Op Fusion、无法提前规划显存空间Bufferization更无法向硬件发射极致向量化的指令。为了在编译期给模糊的维度开辟一双透视眼编译器引入了核心的基础设施——Shape Inference形状推导。1. 什么是 Shape Inference它的核心目的Shape Inference形状推导是指在不需要实际运行模型不流过具体张量数据的前提下仅仅通过算子的输入形状、静态属性以及算子的数学语义在编译期计算出图中间及最终输出张量的形状Dimensions和元素类型DataType的机制。[ 初始高层 IR维度模糊 ] [ 形状推导后维度透视 ] %0 tosa.arg0() : tensor2x4xf32 %0 tosa.arg0() : tensor2x4xf32 %1 tosa.arg1() : tensor4x8xf32 %1 tosa.arg1() : tensor4x8xf32 // 此时输出形状未确定 // 编译器在编译期精准推导出了维度 %2 tosa.matmul(%0, %1) %2 tosa.matmul(%0, %1) : tensor2x8xf32 : (tensor2x4xf32, tensor4x8xf32) - tensor?x?xf32形状推导的双重硬通货静态形状Static Shapes所有的维度在编译期都是已知整数如tensor1x4096xf32。这是硬件厂商最喜欢的形态因为能做到极致的静态显存复用。动态形状Dynamic Shapes含有未知维度的形态如tensor?x4096xf32MLIR 中用?表达。形状推导的目标是在线性代数约束下尽可能把?具象化或者建立动态维度之间的符号等价关系如证明输出的?和输入的?长度绝对相等。2. MLIR 核心基础设施ShapeInferenceOpInterface 接口与传统编译器为每个算子硬编码推导逻辑不同MLIR 采用了高内联、面向对象的算子接口OpInterface机制。任何一个方言算子比如自研 NPU 方言中的伴随特征扩展或标准StableHLO算子只要在 ODS.td文件中挂载了DeclareOpInterfaceMethodsShapeInferenceOpInterface就意味着它必须在 C 侧完美实现一套形状推导契约。工业级工程实践一个标准转置算子的 C 形状推导实现当编译器遍历到某个特定算子时会隐式调用该算子的inferShapes方法#includemlir/Interfaces/InferTypeOpInterface.hnamespacemlir{namespacemy_npu{// 显式实现由 ODS 自动生成的接口方法mlir::LogicalResultTransposeOp::inferShapes(){// 1. 获取输入张量的类型信息autoinputTypegetOperand().getType().dyn_castRankedTensorType();if(!inputType)returnmlir::failure();autoinputShapeinputType.getShape();// 例如 [2, 4]// 2. 模拟算子的物理行为创建输出的形状骨架llvm::SmallVectorint64_t,2inferredShape;if(inputShape.size()2){// 二维矩阵转置行与列互换inferredShape.push_back(inputShape[1]);inferredShape.push_back(inputShape[0]);}else{returnemitOpError(目前只支持二维矩阵转置的形状推导);}// 3. 将计算出来的静态/动态形状重新绑定并更新算子输出槽位的类型autoresultTypeRankedTensorType::get(inferredShape,inputType.getElementType());getResult().setType(resultType);returnmlir::success();}}// namespace my_npu}// namespace mlir3. 大模型时代的硬核演进符号形状Symbolic Shapes与 Shape 方言在传统卷积神经网络如 ResNet时代形状推导非常简单因为维度几乎全静态。但在大模型大行其道的今天变长序列引入了高频的动态性如果编译器动辄遇到?就直接放弃优化推理性能将发生雪崩。因此现代 AI 编译器如 OpenXLA, PyTorch TorchInductor, 以及 MLIR 的核心shape方言引入了更高级的符号形状Symbolic Shape Inference。符号推导的心智模型当编译器看到下面这段大模型推理中常见的拼接Concat操作时输入 Atensor?x4096xf32输入 Btensor?x4096xf32形状推导 Pass 不会仅仅把它变成输出tensor?x4096xf32而是会在后台的Shape 方言中发射一段伴随的标量级图表达式Dimoutput[0]DimA[0]DimB[0]\text{Dim}_{\text{output}}[0] \text{Dim}_A[0] \text{Dim}_B[0]Dimoutput​[0]DimA​[0]DimB​[0]MLIR 会在图的最顶层维护一个符号管理器Symbolic Shape Manager。通过符号化表达即使S0S_0S0​是个未知数编译器也能推导出tensorS0x4096和tensorS0x4096进矩阵乘法是完全合法的从而在编译期成功对它们执行跨动态维度的算子融合。4. 形状推导在编译管线中的交响共鸣形状推导绝不是一个孤立存在的流水线节点它是整条 AI 编译优化的基石数据源[ MLIRGen 生成的原始 IR ] (大量未填写的 tensor?x?xf32) │ ▼ [ Shape Inference Pass ] -- 顺藤摸瓜横扫全图灌注维度元数据 │ ┌─────────┴─────────┐ ▼ ▼ [ DRR / 算子融合 Pass ] [ 显式内存规划 Bufferization ] (依据形状大小决定是否融合) (根据精确字节数提前在显存池占座)驱动算子融合Op Fusion优化 Pass如 DRR在尝试将一个Add算子和一个Mul算子融合时会先探查形状推导的结果。如果两者形状不具备广播一致性Broadcast CompatibilityPass 会立马安全退出避免生成错误的硬件 Kernel。驱动显存零拷贝Zero-copy Memory Planning当高层的张量语义向底层内存块MemRef级别降级时底层 Pass 需要调用形状推导产出的数据。有了明确的维度编译器才能在编译期精确算出每一个生命周期阶段需要开辟多少字节Bytes的连续硬件 Buffer从而实现最大程度的宿主与设备显存复用。总结一句话概括Shape Inference 是 AI 编译器在面对大模型动态世界时在编译期精准锚定确定性数学规律的最高法庭。它通过算子接口Interface将多变的维度逻辑规范化配合符号推导系统把模糊的?降维打击成可计算的代数关系。掌握和编写严谨的形状推导机制是异构算力团队完美对接主流深度学习大模型生态、实现极致内存压缩与高性能硬件发射的关键底牌。