3. ONNX 核心概念

在把 PyTorch 模型送进 MLIR 之前，ONNX 是绕不开的中间站。它既不是训练框架，也不是推理引擎，而是一份“模型交换契约”：用统一的计算图描述、统一的算子语义、统一的张量类型，把不同框架的模型翻译成同一种可被解析、优化、再下发的形式。理解 ONNX 的五个核心概念——计算图、算子集、张量类型、opset 版本、生态工具链——直接决定了后续 MLIR 转换能否顺利落地。

一、计算图：ONNX 的骨架

ONNX 模型本质上是一个有向无环图（DAG）。图中的节点（Node）代表算子调用，边（Edge）代表张量数据流。每个节点包含三要素：算子类型（op_type）、输入名列表（input）、输出名列表（output），以及可选的属性（attribute）。图的入口是 graph.input，出口是 graph.output，中间还可能存在 initializer（权重常量）和 value_info（中间张量类型标注）。

与 PyTorch 动态图不同，ONNX 是静态图：所有张量的形状、类型在导出时就被固定下来（或标注为动态维度）。这正是 MLIR 喜欢的形式——静态图便于做类型推导、形状推断和算子融合。下面是一个最小 ONNX 图的文本表示（简化版）：

graph main {
  %input  = input  float[1,3,224,224]
  %weight = initializer float[64,3,7,7]
  %conv   = Conv(%input, %weight) {kernel_shape=[7,7], strides=[2,2]}
  %relu   = Relu(%conv)
  output %relu
}

可以看到，节点之间的连接完全靠张量名建立，没有隐式状态。这种“纯数据流”结构，使得 ONNX 图可以被逐节点遍历、重写、切分，也方便映射到 MLIR 的 func.func + onnx.* 方言。

二、算子集与张量类型

ONNX 的算子集（Operator Set）是一份标准化清单，定义了每个算子的名称、输入输出数量、属性、类型约束和语义。例如 Conv、Gemm、BatchNormalization、Resize、Softmax 等。算子集不是随意扩张的，而是按 opset 版本逐步演进。

张量类型方面，ONNX 支持基础数值类型（float16/float/double、int8/16/32/64、uint8/16/32/64、bool、string）以及复杂类型（sequence、map、optional）。在 MLIR 转换中，最常遇到的是 float32、int64、int32 和 bool。需要特别注意的是：ONNX 的 int64 常用于索引和形状，而很多后端只支持 int32，转换时往往需要插入类型转换节点。

三、opset 版本：语义的“时间戳”

opset 版本是 ONNX 最容易被忽视、却最容易引发转换失败的概念。每个 ONNX 模型在 metadata 中声明 opset_import，例如 opset 13、opset 17。同一个算子在不同 opset 下语义可能不同：

• Resize 在 opset 10 之前叫 Upsample，参数含义也变了；
• Clip 在 opset 11 之前只接受 min/max 属性，之后改为输入张量；
• Softmax 在 opset 13 调整了 axis 默认行为；
• Split 在 opset 13 之后支持 split 作为输入而非属性。

因此，在 PyTorch 导出 ONNX 时，必须显式指定 opset 版本，并确保目标 MLIR 转换工具链支持该版本。如果 opset 过高，转换器可能不认识新算子；如果过低，又可能丢失某些优化机会。实践中，opset 13～17 是目前兼容性较好的区间。

opset 版本	典型变化	转换影响
opset 9	基础算子稳定，Upsample 仍存在	兼容老工具链，但缺少新算子
opset 11	Clip 改为输入，Resize 增强	需检查转换器是否支持新语义
opset 13	Softmax/Split 语义调整	推荐作为默认导出目标
opset 17	LayerNormalization 等新算子	部分 MLIR 方言可能尚未覆盖

四、ONNX 生态工具链：从导出到验证

ONNX 不是孤立的文件格式，而是一套工具链。围绕它，形成了完整的“导出—检查—优化—转换”流水线：

1. 导出：PyTorch 的 torch.onnx.export，或 TensorFlow 的 tf2onnx；
2. 检查：onnx.checker.check_model 验证图结构合法性；
3. 可视化：Netron 查看节点连接、张量形状和属性；
4. 优化：onnxoptimizer 或 onnx-simplifier 做常量折叠、算子融合；
5. 转换：onnx-mlir 将 ONNX 图翻译为 MLIR 的 onnx 方言，再降级到 linalg、affine 等；
6. 验证：onnxruntime 做数值比对，确保转换前后输出一致。

其中，onnx-mlir 是 PyTorch/ONNX 到 MLIR 转换实战的核心工具。它读取 ONNX 模型后，先构建 MLIR 的 onnx 方言表示，再经过一系列 pass 逐步降低到目标后端。理解 ONNX 的图结构、算子语义和 opset 版本，正是为了在 MLIR 中正确匹配和重写这些算子。

五、常见误区与复盘

误区一：认为 ONNX 图与 PyTorch 图完全等价。实际上，PyTorch 导出时会做算子分解和常量折叠，某些控制流（如 if、loop）可能被展开或丢失，导致 ONNX 图与原始模型行为不完全一致。

误区二：忽略 opset 版本。同一个模型用 opset 11 和 opset 17 导出，节点数量和算子类型可能完全不同。转换前必须确认工具链支持的 opset 范围。

误区三：把 ONNX 当作最终部署格式。ONNX 是中间表示，不是执行引擎。它的价值在于“交换”，真正的高性能执行需要依赖 MLIR 降级后的后端代码或 ONNX Runtime 的图优化。

误区四：不检查张量类型。ONNX 中 int64 索引在 MLIR 中可能被要求转为 int32，若忽略这一点，会在类型推导阶段报错。建议在导出时尽量固定类型，或在转换前插入 Cast 节点。