ONNX 前端

目录结构

compiler/bindings/python/iree/compiler/tools/import_onnx/ ├── main.py # CLI 入口 (9.7KB) ├── importer_externalization_overrides.py # IREE 特定的 ONNX 导入覆盖 (16.4KB)

核心文件

文件	作用
`__main__.py`	`iree-import-onnx` 命令的入口，解析 CLI 参数，调用 ONNX importer
`importer_externalization_overrides.py`	IREE 自定义的 op 外部化逻辑，处理 ONNX → torch dialect 转换

底层依赖

底层依赖的是 torch-mlir 的 ONNX importer（注意是小写 m）：

third_party/torch-mlir/python/torch_mlir/extras/onnx_importer.py

转换链：ONNX protobuf → onnx_importer.py → IREE overrides → MLIR (torch dialect)

使用方法


python -m iree.compiler.tools.import_onnx model.onnx -o model.mlir

# 指定 opset 版本
python -m iree.compiler.tools.import_onnx model.onnx -o model.mlir --opset-version 17

# 大模型支持 - 参数外部化（推荐 >2GB 的模型）
python -m iree.compiler.tools.import_onnx model.onnx -o model.mlir --large-model --externalize-params

转换流程

ONNX 文件 (.onnx)
      │
      ▼
① onnx.load()              ← 用 onnx Python 包读取 protobuf
                             * 支持 load_external_data 模式（大模型）
                             * 支持 data_dir 指定外部权重目录
      │
② onnx.version_converter   ← 如果指定了 --opset-version，升级 opset（只在新版本时）
      │
③ onnx.shape_inference()   ← 做 shape 推导（填充未标注的维度）
                             * 小模型：in-memory (infer_shapes)
                             * 大模型：file-based (infer_shapes_path)
      │
④ onnx_importer.ModelInfo() ← 解析 ONNX 图结构（节点、输入、输出）
      │
⑤ model_info.create_module() ← 创建 MLIR Module（空的 MLIR 容器）
      │
⑥ NodeImporter.define_function() ← 遍历 ONNX 节点，逐一转为 MLIR op
                             * 普通模式：onnx_importer.NodeImporter
                             * 参数外部化：IREENodeImporter（继承自 NodeImporter）
      │
⑦ imp.import_all()         ← 导入所有节点
      │
⑧ m.verify()               ← 验证 MLIR 模块（可通过 --no-verify 跳过）
      │
⑨ m.get_asm()              ← 把 MLIR 输出为文本格式 (.mlir)
      │
⑩ imp.save_params()        ← 保存外部化参数到 .irpa 文件（如果启用）

关键代码对应

步骤	代码行	做什么
读取 ONNX	`__main__.py:99` `onnx.load(args.input_file)`	读 protobuf
Opset 版本转换	`__main__.py:109` `onnx.version_converter.convert_version()`	升级 opset
Shape 推导（小模型）	`__main__.py:124` `onnx.shape_inference.infer_shapes()`	小模型 in-memory
Shape 推导（大模型）	`__main__.py:139` `onnx.shape_inference.infer_shapes_path()`	大模型 file-based
创建 MLIR	`__main__.py:29-30` `ModelInfo` → `create_module()`	创建空 MLIR module
普通模式转换节点	`__main__.py:66` `NodeImporter.define_function(graph, m)`	ONNX → MLIR
参数外部化转换	`__main__.py:64` `IREENodeImporter.define_function()`	参数外部化模式
全部导入	`__main__.py:68` `imp.import_all()`	遍历所有节点
验证 MLIR	`__main__.py:71` `m.verify()`	验证模块正确性
输出文本	`__main__.py:78/80` `m.get_asm()`	打印 MLIR 文本
保存参数	`__main__.py:83` `imp.save_params()`	保存 `.irpa` 参数文件

重要特性

大模型支持 (`--large-model`)

跳过 onnx.checker 检查
使用文件模式进行 shape inference
避免 2GB protobuf 限制

参数外部化 (`--externalize-params`)

将大权重参数外部化为 .irpa 文件
通过 --num-elements-threshold 控制阈值（默认 100）
通过 --params-scope 设置命名空间（默认 "model"）
通过 --save-params-to 指定输出路径
通过 --param-gb-threshold 控制内存刷新阈值（低内存机器用）

输入外部化 (`--externalize-inputs-threshold`)

将指定索引之后的输入视为外部参数
适用于有多个权重输入的模型