维度透视眼的幕后功臣：深度解构 MLIR 核心虚函数 `inferShapes`-开发者社区

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

在上一篇文章中，我们聊到了RankedTensorType是如何作为几何铁律规范大模型数据的，也提及了Shape Inference（形状推导）这一不可或缺的基础设施。

在 MLIR 的工业级工程落地中，形状推导绝不是凭空产生的，它高度依赖于各个方言算子（Ops）内部硬核实现的**核心 C++ 虚函数——inferShapes**。

无论是官方维护的StableHLO、TOSA，还是各大算力巨头为自家大模型推理芯片定制的私有方言，inferShapes都是算子元数据校验与图优化的“第一质检员”。本文将直接带你把手术刀切入 C++ 源码层级，深度拆解这个关键函数的编写范式、核心逻辑以及大模型时代下的硬核技巧。

1.`inferShapes`来自哪里？接口的契约

在 MLIR 中，如果你希望一个算子具备在编译期推导输出维度的能力，首先需要在 ODS（.td文件）中让该算子继承一个内置的接口（OpInterface）：InferTypeOpInterface。

// 在 ODS 中声明算子时挂载接口 def MyNPU_MatMulOp : MyNPU_Op<"matmul", [ DeclareOpInterfaceMethods<InferTypeOpInterface> // 核心：承诺实现形状与类型推导 ]> { // ... 算子的输入输出定义 }

一旦在 ODS 里做出了这个承诺，TableGen 就会在编译期自动为你喷涌出该算子的 C++ 虚函数声明。此时，你就必须在 C++ 文件（通常是Ops.cpp）中提供它的具体实现，也就是inferShapes（或者在标准接口里更通用的inferReturnTypes模式）。

2. 工业级 C++ 范式：一个标准算子的`inferShapes`实现

编写inferShapes的核心心智模型非常纯粹：在不触碰真实矩阵数据的前提下，仅凭输入的形状（Shape）、数据类型（DataType）以及算子的数学规则，在编译期精准算出并更新输出的类型。

我们来看一个大模型推理框架中极高频的原子算子——二维矩阵乘法（MatMulOp）的工业级 C++ 实现：

#include"mlir/Interfaces/InferTypeOpInterface.h"#include"MyNPU/MyNPUOps.h"usingnamespacemlir;// 显式实现该算子的形状推导虚函数LogicalResult my_npu::MatMulOp::inferShapes(){// 1. 捞出左操作数（LHS）和右操作数（RHS）的强类型信息autolhsType=getLhs().getType().dyn_cast<RankedTensorType>();autorhsType=getRhs().getType().dyn_cast<RankedTensorType>();// 稳健性防线：如果前驱算子还没推导出输入类型，则本算子安全退出if(!lhsType||!rhsType)returnsuccess();// 获取几何形状，例如 LHS 是 [M, K]，RHS 是 [K, N]ArrayRef<int64_t>lhsShape=lhsType.getShape();ArrayRef<int64_t>rhsShape=rhsType.getShape();// 2. 严谨的数学与拓扑契约校验（Verifier 功能）if(lhsShape.size()!=2||rhsShape.size()!=2){returnemitOpError("仅支持二维矩阵乘法的形状推导");}int64_tm=lhsShape[0];int64_tkLhs=lhsShape[1];int64_tkRhs=rhsShape[0];int64_tn=rhsShape[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_toutM=m;int64_toutN=n;SmallVector<int64_t,2>inferredShape({outM,outN});// 4. 元素类型对齐：矩阵乘法输出通常与输入保持相同的 Float 精度（如 BF16）Type resElementType=lhsType.getElementType();// 5. 惊艳的一击：将推导出的 RankedTensorType 强行灌注并更新给算子的输出槽位autoresultType=RankedTensorType::get(inferredShape,resElementType);getResult().setType(resultType);returnsuccess();}

3. 大模型时代下`inferShapes`的硬核对抗：动态轴的交叉感染

在大模型（LLM）变长 Token 推理场景下，inferShapes的编写难度呈指数级上升。因为你面临的输入常常不是干净的数字，而是满墙的?（ShapedType::kDynamic）。

一个合格的工业级inferShapes必须具备以下对抗动态性的能力：

① 动态轴的符号传递

假设我们在为BroadcastOp（广播算子）编写推导：

输入张量：tensor<?xf32>（变长 Prompt）
目标形状：tensor<1x?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]=4096

OutputDim[0]=InputDim[0]×32×1284096=InputDim[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的虚函数在底层默默工作，大模型那原本充满迷雾的动态计算图，才能在进入芯片晶体管发射前，被规整得如同阅兵方阵般严丝合缝。

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