# 阶段 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 上做。