160 lines
7.3 KiB
Markdown
160 lines
7.3 KiB
Markdown
# 阶段 3：前端输入（StableHLO / TOSA）✅
|
||
|
||
> 2026-08-15。理解模型如何进入 IREE，以及两个统一的输入方言如何收敛到 linalg。
|
||
> 对应规划阶段 3。源码位置：`src/iree/compiler/plugins/input/`。
|
||
|
||
## 1. 模型怎么进入 IREE
|
||
|
||
IREE 不直接吃框架模型，而是通过**导入工具**先把框架模型转成 MLIR 的「输入方言」，再走统一的编译管线：
|
||
|
||
| 框架 | 导入工具 | 产生的输入方言 |
|
||
|---|---|---|
|
||
| PyTorch | torch-mlir / iree-turbine | Torch → StableHLO 或 TOSA |
|
||
| TensorFlow | iree-import-tf（XLA 导出） | StableHLO |
|
||
| JAX | jax.export / XLA | StableHLO |
|
||
| ONNX | iree-import-onnx | TOSA |
|
||
| TFLite | iree-import-tflite | TOSA |
|
||
|
||
> 当前环境只装了 `iree-import-onnx`（base 包自带）。TF 导入器需要额外的 `iree-tools-tf` 包，
|
||
> torch 前端需要 `iree-turbine`（依赖 torch）。后面阶段 3 深入时按需补装。
|
||
|
||
## 2. 两个统一的输入方言
|
||
|
||
IREE 支持多个输入方言，但**两个主力是 StableHLO 和 TOSA**，它们各自代表一条血统：
|
||
|
||
| | StableHLO | TOSA |
|
||
|---|---|---|
|
||
| 全称 | Stable High-Level Ops | Tensor Operator Set Architecture |
|
||
| 血统 | XLA / HLO（TF、JAX、PyTorch-XLA） | Arm TOSA 规范（ONNX、TFLite） |
|
||
| 风格 | 图级高层 op（dot/transpose/select/broadcast） | 精确 spec 的算子集（max_pool2d/…，带 kernel/stride/pad 属性） |
|
||
| 特点 | 保留框架 `{name=...}` 属性，op 粒度粗 | 每个 op 有精确语义规范，量化友好 |
|
||
|
||
**为什么要两个？** 因为它们分别服务于不同的框架生态（XLA 系 vs ONNX/TFLite 系），
|
||
但**最终都收敛到同一个内部表示（linalg / linalg-ext）**，所以下游编译管线是共享的。
|
||
|
||
## 3. IR 示例
|
||
|
||
### StableHLO（fullyconnected.mlir，一个全连接层）
|
||
```mlir
|
||
func.func @main(%arg0: tensor<1x5xf32>, %arg1: tensor<1x5x3x1xf32>) -> tensor<5x1x5xf32> {
|
||
%0 = stablehlo.transpose %arg0, dims = [1, 0] {name = "transpose.41"} : ...
|
||
%3 = stablehlo.broadcast_in_dim %2, dims = [0, 1] {name = "broadcast.56"} : ...
|
||
%29 = stablehlo.dot %28, %cst_11, precision = [DEFAULT, DEFAULT] {name = "dot.62"} : ...
|
||
%31 = stablehlo.add %29, %30 {name = "add.65"} : ...
|
||
%32 = stablehlo.maximum %25, %31 {name = "maximum.68"} : ...
|
||
%34 = stablehlo.select %23, %24, %33 {name = "select.71"} : ...
|
||
}
|
||
```
|
||
- 高层图级 op，保留 TF 原图的 `{name = "..."}` 属性
|
||
|
||
### TOSA（max_pool.mlir）
|
||
```mlir
|
||
%result = tosa.max_pool2d %0 {kernel = array<i64: 2, 2>, stride = array<i64: 1, 1>,
|
||
pad = array<i64: 0, 0, 0, 0>} : (tensor<1x2x4x1xi8>) -> tensor<1x1x3x1xi8>
|
||
```
|
||
- 算子 + 显式属性（kernel/stride/pad），语义精确
|
||
|
||
## 4. 前端转换：输入方言 → linalg
|
||
|
||
用 `--mlir-print-ir-after-all` 追踪 StableHLO 模型的转换，前端约 19 个 pass，
|
||
其中**最关键的两个**：
|
||
|
||
| pass | 作用 |
|
||
|---|---|
|
||
| `ConvertStableHloToLinalgExt` | StableHLO → LinalgExt（IREE 扩展的 linalg） |
|
||
| `ConvertStableHloToIreeInputDialects` | StableHLO → IREE 输入方言（linalg 等） |
|
||
|
||
转换前后对照（`stablehlo.dot` → `linalg.generic`）：
|
||
|
||
```mlir
|
||
// 之前：stablehlo（高层语义）
|
||
%29 = stablehlo.dot %28, %cst_11, precision = [DEFAULT, DEFAULT] {name = "dot.62"} : ...
|
||
|
||
// 之后：linalg.generic（结构化计算 + 显式迭代空间）
|
||
%1 = linalg.generic {indexing_maps = [#map, #map1],
|
||
iterator_types = ["parallel", "parallel", "parallel"]}
|
||
ins(%expanded : tensor<5x1xf32>) outs(%0 : tensor<5x1x5xf32>) attrs = {name = "broadcast.56"} {
|
||
^bb0(%in: f32, %out: f32):
|
||
linalg.yield %in : f32
|
||
} -> tensor<5x1x5xf32>
|
||
|
||
// stablehlo.compare GT → linalg.generic + arith.cmpf
|
||
^bb0(%in: f32, %out: i1):
|
||
%45 = arith.cmpf ogt, %in, %cst : f32
|
||
linalg.yield %45 : i1
|
||
```
|
||
|
||
**关键点**：`stablehlo.*` 是"命名高层 op"，`linalg.generic` 是"结构化计算"——用
|
||
`indexing_maps`（仿射映射描述数据布局）+ `iterator_types`（parallel/reduction）+
|
||
region（元素级计算）显式描述计算。这就是阶段 2 里看到的 codegen 起点的由来。
|
||
|
||
## 5. 源码结构（plugins/input/）
|
||
|
||
前端转换是**插件式**的（IREE 2024 重构后，输入转换拆成独立 plugin）：
|
||
|
||
```
|
||
src/iree/compiler/plugins/input/
|
||
├── StableHLO/
|
||
│ ├── Conversion/ # StableHLO → linalg/linalg-ext
|
||
│ │ ├── StableHLOToIREEInputDialects.cpp # 主转换
|
||
│ │ ├── StableHLOToLinalgExt.cpp
|
||
│ │ ├── LegalizeControlFlow.cpp / LegalizeShapeComputations.cpp
|
||
│ │ └── VerifyCompilerInputLegality.cpp
|
||
│ └── PluginRegistration.cpp
|
||
├── TOSA/
|
||
│ ├── InputConversion/
|
||
│ │ ├── TosaToLinalgExt.cpp
|
||
│ │ ├── StripSignedness.cpp / Converti48Toi64.cpp # TOSA 特有预处理
|
||
│ │ └── VerifyCompilerTOSAInputLegality.cpp
|
||
│ └── PluginRegistration.cpp
|
||
└── Torch/ # torch-mlir（Torch → StableHLO/TOSA）
|
||
```
|
||
|
||
## 6. 完整输入 pass 序列（StableHLO 模型，去重后前 19 个）
|
||
|
||
```
|
||
CheckVHLOStableHloMixUsage → VhloToVersionPass → VhloLegalizeToStablehloPass # VHLO 版本化 → StableHLO
|
||
→ Canonicalizer → StableHLOCanonicalize → CSE # 化简
|
||
→ LegalizeStableHLOCustomCalls → LegalizeControlFlow → FlattenTuplesInSCF # 控制流/自定义调用
|
||
→ StableHLOToStableHLOPreprocessing → ShapeToShapeLowering → ConvertShapeToStandard # shape 计算
|
||
→ Inliner → LegalizeShapeComputations
|
||
→ ConvertStableHloToLinalgExt → ChloLegalizeToStablehloPass # ★ StableHLO → linalg-ext
|
||
→ ConvertStableHloToIreeInputDialects # ★ → IREE 输入方言
|
||
→ ReconcileUnrealizedCastsPass → VerifyCompilerStableHloInputLegality # 校验
|
||
```
|
||
|
||
> 之后（pass 20+）就接上阶段 2 的 IREEImportPublicPass → WrapEntryPointsPass → … 主管线。
|
||
|
||
## 7. 与 TVM 0.6 对照
|
||
|
||
| IREE 前端 | TVM 0.6 |
|
||
|---|---|
|
||
| iree-import-tf / torch-mlir / iree-import-onnx | 各框架的 relay importer |
|
||
| StableHLO / TOSA | Relay IR（图级） |
|
||
| ConvertStableHloToIreeInputDialects（→linalg） | Relay → TE 的 lowering |
|
||
| linalg.generic | TE 的 compute |
|
||
|
||
## 8. 实操命令
|
||
|
||
```bash
|
||
# 编译 StableHLO 模型，看完整前端转换
|
||
iree-compile --iree-input-type=stablehlo --iree-hal-target-backends=llvm-cpu \
|
||
--mlir-print-ir-after-all src/iree/tests/e2e/stablehlo_models/fullyconnected.mlir -o /tmp/fc.vmfb 2> dump.txt
|
||
|
||
# 只看 StableHLO → linalg 的关键转换
|
||
iree-compile --iree-input-type=stablehlo --iree-hal-target-backends=llvm-cpu \
|
||
--mlir-print-ir-after=ConvertStableHloToIreeInputDialects --mlir-print-ir-module-scope \
|
||
src/iree/tests/e2e/stablehlo_models/fullyconnected.mlir -o /tmp/fc.vmfb 2>&1 | head -60
|
||
|
||
# TOSA 模型（同理，--iree-input-type=tosa）
|
||
# ONNX → TOSA（iree-import-onnx 可用）
|
||
```
|
||
|
||
## 9. 关键结论
|
||
|
||
1. **前端 = 框架方言 → 输入方言 → linalg**，下游（Flow→Stream→HAL→codegen→VM）完全共享。
|
||
2. **StableHLO 和 TOSA 只是两条不同的"入口"**，最终都在 `ConvertStableHloToIreeInputDialects`
|
||
/ `TosaToLinalgExt` 收敛到 linalg/linalg-ext。
|
||
3. **linalg.generic 是前端到 codegen 的分水岭**：之前是"命名高层 op"，之后是"结构化计算"，
|
||
后续所有 tiling/vectorization 都在 linalg 上做。
|