维度透视眼的幕后功臣：深度解构 MLIR 核心虚函数 `inferShapes`

维度透视眼的幕后功臣深度解构 MLIR 核心虚函数inferShapes在上一篇文章中我们聊到了RankedTensorType是如何作为几何铁律规范大模型数据的也提及了Shape Inference形状推导这一不可或缺的基础设施。在 MLIR 的工业级工程落地中形状推导绝不是凭空产生的它高度依赖于各个方言算子Ops内部硬核实现的**核心 C 虚函数——inferShapes**。无论是官方维护的StableHLO、TOSA还是各大算力巨头为自家大模型推理芯片定制的私有方言inferShapes都是算子元数据校验与图优化的“第一质检员”。本文将直接带你把手术刀切入 C 源码层级深度拆解这个关键函数的编写范式、核心逻辑以及大模型时代下的硬核技巧。1.inferShapes来自哪里接口的契约在 MLIR 中如果你希望一个算子具备在编译期推导输出维度的能力首先需要在 ODS.td文件中让该算子继承一个内置的接口OpInterfaceInferTypeOpInterface。// 在 ODS 中声明算子时挂载接口 def MyNPU_MatMulOp : MyNPU_Opmatmul, [ DeclareOpInterfaceMethodsInferTypeOpInterface // 核心承诺实现形状与类型推导 ] { // ... 算子的输入输出定义 }一旦在 ODS 里做出了这个承诺TableGen 就会在编译期自动为你喷涌出该算子的 C 虚函数声明。此时你就必须在 C 文件通常是Ops.cpp中提供它的具体实现也就是inferShapes或者在标准接口里更通用的inferReturnTypes模式。2. 工业级 C 范式一个标准算子的inferShapes实现编写inferShapes的核心心智模型非常纯粹在不触碰真实矩阵数据的前提下仅凭输入的形状Shape、数据类型DataType以及算子的数学规则在编译期精准算出并更新输出的类型。我们来看一个大模型推理框架中极高频的原子算子——二维矩阵乘法MatMulOp的工业级 C 实现#includemlir/Interfaces/InferTypeOpInterface.h#includeMyNPU/MyNPUOps.husingnamespacemlir;// 显式实现该算子的形状推导虚函数LogicalResult my_npu::MatMulOp::inferShapes(){// 1. 捞出左操作数LHS和右操作数RHS的强类型信息autolhsTypegetLhs().getType().dyn_castRankedTensorType();autorhsTypegetRhs().getType().dyn_castRankedTensorType();// 稳健性防线如果前驱算子还没推导出输入类型则本算子安全退出if(!lhsType||!rhsType)returnsuccess();// 获取几何形状例如 LHS 是 [M, K]RHS 是 [K, N]ArrayRefint64_tlhsShapelhsType.getShape();ArrayRefint64_trhsShaperhsType.getShape();// 2. 严谨的数学与拓扑契约校验Verifier 功能if(lhsShape.size()!2||rhsShape.size()!2){returnemitOpError(仅支持二维矩阵乘法的形状推导);}int64_tmlhsShape[0];int64_tkLhslhsShape[1];int64_tkRhsrhsShape[0];int64_tnrhsShape[1];// 检查矩阵乘法的内轴K 轴在静态情况下是否对齐if(!ShapedType::isDynamic(kLhs)!ShapedType::isDynamic(kRhs)){if(kLhs!kRhs){returnemitOpError()矩阵乘法内轴冲突LHS K轴为 kLhs而 RHS K轴为 kRhs;}}// 3. 核心推演输出的几何空间 [M, N]// 如果输入轴是动态的ShapedType::kDynamic即 ?推导结果会自动继承这个动态标志int64_toutMm;int64_toutNn;SmallVectorint64_t,2inferredShape({outM,outN});// 4. 元素类型对齐矩阵乘法输出通常与输入保持相同的 Float 精度如 BF16Type resElementTypelhsType.getElementType();// 5. 惊艳的一击将推导出的 RankedTensorType 强行灌注并更新给算子的输出槽位autoresultTypeRankedTensorType::get(inferredShape,resElementType);getResult().setType(resultType);returnsuccess();}3. 大模型时代下inferShapes的硬核对抗动态轴的交叉感染在大模型LLM变长 Token 推理场景下inferShapes的编写难度呈指数级上升。因为你面临的输入常常不是干净的数字而是满墙的?ShapedType::kDynamic。一个合格的工业级inferShapes必须具备以下对抗动态性的能力① 动态轴的符号传递假设我们在为BroadcastOp广播算子编写推导输入张量tensor?xf32变长 Prompt目标形状tensor1x?xf32在inferShapes内部你不能粗暴地直接给输出写个tensor?x?xf32。你需要建立符号联动如果左边的?代表的是BatchSize那么输出对应的那个轴的?也必须精准承袭自该变量。② 算子退化时的动态逃生在大模型中诸如Reshape或Flatten算子当前端传来的新形状里包含-1代表由编译器自动推导该轴大小时// 意图将一个 3D 张量展平但 Batch 轴动态 %0 my_npu.reshape(%arg0) {new_shape [-1, 4096]} : (tensor?x32x128xf32) - tensor?x?xf32此时的inferShapes内部必须手写一套代数方程求解器OutputDim[1]4096\text{OutputDim}[1] 4096OutputDim[1]4096OutputDim[0]InputDim[0]×32×1284096InputDim[0]\text{OutputDim}[0] \frac{\text{InputDim}[0] \times 32 \times 128}{4096} \text{InputDim}[0]OutputDim[0]4096InputDim[0]×32×128​InputDim[0]只有在inferShapes内部通过数学公式把-1算成与输入等价的?后续的内存分配 Pass 才能安全进行。4.inferShapes在编译器生命周期中的滚动执行写完了inferShapes虚函数后是谁在什么时候去呼叫它呢这主要发生在一个通用的全局优化通道——ShapeInferencePass中。它通常采用类似拓扑排序Topological Sort的贪心算法扫图[ 刚进门的原始 IR ] (大部分中间算子输出全是未填写的 tensor?x?xf32) │ ▼ [ ShapeInferencePass 启动 ] │ ├──► 扫描到算子 %1 my_npu.matmul(%0) │ │ │ └──► 隐式调用 %1 内部手写的 C inferShapes() │ │ │ └──► 算子 %1 的输出类型被正式改写为强类型 │ ▼ [ 驱动后续算子 ] ───► 下游算子感知到 %1 的新维度继续链式触发自身的 inferShapes这种“生而验证、链式推导”的滚动机制确保了只要图的输入边界如模型 Inputs 节点是确定的整张图的所有中间血管和末梢神经都将在极短的编译期时间内被充盈上精准的几何维度信息。总结一句话概括inferShapes是将 AI 算子的物理数学语义转化为编译器可利用元数据的离散加工厂。它要求编译器工程师不仅要对算子的底层的硬件行为明察秋毫还要在 C 层面对各种边界、动态?建立严谨的代数映射。正是因为有了成百上千个类似inferShapes的虚函数在底层默默工作大模型那原本充满迷雾的动态计算图才能在进入芯片晶体管发射前被规整得如同阅兵方阵般严丝合缝。