FX -> torch-mlir -> Torch Backend MLIR 详细流程

本文聚焦以下阶段：

• 输入：torch.fx.GraphModule（以下简称 FX Graph）
• 过程：通过 torch-mlir 的 FX importer 导入并运行 lowering pipeline
• 输出：Torch Backend Contract 形态的 MLIR 文本

对应主流程代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:108

---

1. 入口与上下文

入口函数是 NpuMlirKernel.codegen_kernel()：

• 文件：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:108
• 关键调用链：

1. scalarize_tensor_ops_on_scalars(self._gm_with_prim_cast)
2. set_model_name("MODEL_NAME")
3. model_name = get_aot_compilation_context() ...
4. mlir_module = stateless_fx_import(...)
5. run_pipeline_with_repro_report(..., "builtin.module(torch-lower-to-backend-contract)", ...)
6. str(mlir_module) 生成 MLIR 字符串

---

2. FX 导入前预处理（稳定 importer 输入）

代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:124

调用：

• scalarize_tensor_ops_on_scalars(...)

定义：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/utils.py:370

主要作用：

• 将 aten.add.Tensor / aten.mul.Tensor 且“第二个参数不是 Node（即常量）”的节点改成 aten.add.Scalar / aten.mul.Scalar
• 降低 FX importer 对 Tensor/Scalar 混合形式的歧义，提高导入稳定性

---

3. 设置 MLIR 函数名（符号名来源）

代码：

• torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:125
• torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:126

流程：

1. set_model_name("MODEL_NAME")
2. 从 get_aot_compilation_context() 取到 model_name
3. 后续传给 importer 作为 func_name

说明：

• 这里的 MODEL_NAME 是模板名，稍后在 define_kernel 阶段会替换成真实 kernel 名。

---

4. FX -> torch-mlir 导入

调用点：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:127

mlir_module = stateless_fx_import(
    self._gm_with_prim_cast,
    model_name=model_name,
    import_symbolic_shape_expressions=False
)

实现位于：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:195

stateless_fx_import 做三件事：

1. 创建 ir.Context 并注册 torch dialect
   代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:207-209
2. 创建 FxImporter，调用 import_stateless_graph(gm.graph, ...) 导入 FX 节点
   代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:210-211
3. 调用 _module_lowering(...) 返回 ir.Module
   代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:212-214

---

5. 为什么要 patch torch-mlir 的 importer

补丁注入位置：

• fx_importer.FxImporter.import_stateless_graph = _patch_import_stateless_graph
torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:192
• fx_importer.sympy_expr_to_semi_affine_expr = _patch_sympy_expr_to_semi_affine_expr
torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:193

5.1 patch_import_stateless_graph

定义：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:51

关键增强：

• 从 placeholder 的 SymInt 提取符号范围约束，写入 self._cc._symbolic_guards
  代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:65-100
• 创建 func_dialect.FuncOp 时，增加属性 torch.assume_strict_symbolic_shapes
  代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:112
• 通过 GraphNodeImporter.import_nodes(...) 导入 FX 节点
  代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:114-122

5.2 patch_sympy_expr_to_semi_affine_expr

定义：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/torch_mlir_patch.py:127

关键增强：

• 对 Add/Mul/Pow/Mod/FloorDiv 等 sympy 表达式提供到 MLIR affine/semi-affine 的翻译
• 减少动态形状场景中表达式无法转换导致的导入失败

---

6. torch-mlir 模块 -> Torch Backend MLIR

代码：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:131

run_pipeline_with_repro_report(
    mlir_module,
    "builtin.module(torch-lower-to-backend-contract)",
    "Lowering TorchFX IR -> Torch Backend IR",
)

这一步是核心 lowering：

• 将 importer 产出的 Torch IR 降到 backend contract 形态
• 输出仍是 MLIR module，但已进入“后端可消费”的 Torch Backend IR 约定层

---

7. 最终导出 MLIR 字符串

代码：

• torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:138-141
• torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/mlir.py:143

步骤：

1. 遍历 module 中 FuncOp，补充 torch.assume_strict_symbolic_shapes 属性
2. str(mlir_module) 生成文本

该文本随后会进入下一阶段（不在本文范围）：

• MLIRProcessor.get_named_op_str(...)
• Torch Backend MLIR -> NamedOp Backend MLIR

---

8. 一句话总结

这段流程本质是：

1. 先把 FX 图做轻量规范化（尤其 Tensor/Scalar 算子形式）
2. 使用带补丁的 torch-mlir FxImporter 导入成 MLIR Module
3. 跑 torch-lower-to-backend-contract pipeline 降到 Torch Backend IR
4. 把结果序列化成 MLIR 文本交给后续编译链