阶段 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，一个全连接层）

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）

%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）：

// 之前：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. 实操命令

# 编译 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. 关键结论

前端 = 框架方言 → 输入方言 → linalg，下游（Flow→Stream→HAL→codegen→VM）完全共享。
StableHLO 和 TOSA 只是两条不同的"入口"，最终都在 ConvertStableHloToIreeInputDialects / TosaToLinalgExt 收敛到 linalg/linalg-ext。
linalg.generic 是前端到 codegen 的分水岭：之前是"命名高层 op"，之后是"结构化计算"，后续所有 tiling/vectorization 都在 linalg 上做。

7.3 KiB Raw Permalink Blame History Unescape Escape