章节3：MLIR完成一个端到端的流程

系列文章目录

章节1：什么是MLIR

章节2：MLIR的引入

章节3：MLIR完成一个端到端的流程

章节4：编写一个pass

​​​​​​​

---

环境说明

操作系统：

ProductName: macOS

ProductVersion: 26.6.2

BuildVersion: 25G83

软件依赖：

CMake，Ninja，Bazel

1.大体的流程如下：

GraphDef (.pbtxt)
    │ ① tf-mlir-translate
    ▼
MLIR: tf + tf_executor 方言（MLIR 第一次登场）
    │ ② tf-opt -tf-executor-to-functional-conversion
    ▼
MLIR: tf 方言（函数式）
    │ ③ tf-opt --tf-to-hlo-pipeline
    ▼
MLIR: HLO (MHLO) 方言
    │ ④ mlir-hlo-opt --hlo-legalize-to-stablehlo
    ▼
MLIR: StableHLO 方言
    │ ⑤ mlir-hlo-opt --stablehlo-legalize-to-linalg --one-shot-bufferize
    ▼
MLIR: Linalg 方言（memref 语义）
    │
    ├─ 方案一：⑥ mlir-opt lowering + mlir-runner JIT（MLIR 最后一次登场）
    │       ▼
    │   直接执行 → 输出 [5, 5, ..., 5]
    │
    └─ 方案二：⑥ mlir-opt lowering → ⑦ mlir-translate （MLIR 最后一次登场）→ clang
            ▼
        LLVM IR → .o → 链接 C 程序 → 可执行文件 → 输出 5 5 ...

工具	类型	核心作用	在你的流程中的位置
tf-mlir-translate	翻译器（入口）	GraphDef → MLIR	①
tf-opt	Pass 驱动器	tf_executor → tf 函数式；tf → HLO	② ③
mlir-hlo-opt	Pass 驱动器	HLO → StableHLO → Linalg + bufferize	④ ⑤
mlir-opt	Pass 驱动器	Linalg/memref/arith/cf/func → LLVM 方言	⑥（两条方案共用）
mlir-translate	翻译器（出口）	LLVM 方言 MLIR → LLVM IR	⑦（方案二）
mlir-runner	JIT 执行器	LLVM 方言 MLIR → 机器码 → 执行	⑥（方案一）
clang	编译器/链接器	LLVM IR → .o，链接生成可执行文件	方案二最后两步

1.1 tf-mlir-translate：MLIR 的"入口翻译器"

核心作用 ：把 TensorFlow 的 GraphDef（.pbtxt/.pb）翻译成 MLIR。翻译完成后，IR 进入 MLIR 体系，后续所有操作都在 MLIR 方言之间进行。

1.2 tf-opt：TensorFlow MLIR 的 Pass 驱动器

tf-opt 是 TensorFlow 提供的 MLIR pass 驱动工具，专门用来运行 TF 相关的 lowering pass。在流程中它执行了两次：

1.2.1 -tf-executor-to-functional-conversion

把 tf_executor 方言的"图执行语义"提升为函数式 IR。

• 输入是 tf_executor.graph + tf_executor.island，输出是普通的 func.func + tf 算子。

• 这一步是为了去掉 GraphDef 的执行包装，让 IR 变成更纯粹的"函数 + 算子"形式，方便后续做算子级映射。

1.2.2 --tf-to-hlo-pipeline

把 TensorFlow 算子逐个映射到 HLO（MHLO）算子 。这是一个聚合 pipeline，内部包含一系列 pattern，负责把 tf.* 重写成 mhlo.*

1.3 mlir-hlo-opt：HLO/StableHLO/Linalg 的 Pass 驱动器

mlir-hlo-opt 是 XLA/MLIR-HLO 项目提供的 pass 驱动工具，专注于 HLO → StableHLO → Linalg 这条路径。

1.3.1 --hlo-legalize-to-stablehlo

把 XLA 私有的 MHLO 方言转成 StableHLO 方言。StableHLO 是跨框架可移植的标准化方言，这一步让后续 lowering 走一条更通用、更稳定的路径。

1.3.2 --stablehlo-legalize-to-linalg --one-shot-bufferize

这是两个 pass 的组合：

• --stablehlo-legalize-to-linalg ：把 StableHLO 的算子降成 Linalg 方言（结构化循环表示）。

• --one-shot-bufferize ：把 tensor 的值语义 转成 memref 的 buffer 语义 。Linalg 从"返回 tensor"变成"原地写入 outs"，函数参数/返回值也从 tensor 变成 memref。这是从"抽象值"到"真实内存"的关键一步。

1.4 mlir-opt：LLVM 方言的"总 lowering 驱动器"

mlir-opt 是 MLIR 最核心的 pass 驱动工具 ，用来运行通用的 lowering pass。在你的两条方案里，它负责把 Linalg/memref/arith/cf/func 一路降到 LLVM 方言。

用到的 pass 包括：

Pass	作用
-convert-linalg-to-loops	把 linalg.generic 展开成 scf.for 循环
-lower-affine	affine 方言降为 scf/memref/arith
-convert-scf-to-cf	scf 控制流转为 cf 方言
-convert-to-llvm	聚合 pass，把所有已注册方言统一降到 LLVM 方言
-llvm-request-c-wrappers	给函数打上 llvm.emit_c_interface 属性，请求生成 C ABI 包装函数

-convert-to-llvm 是方案一和方案二共用的"总降级"步骤，执行完后 IR 里就只剩 llvm.* 操作了。

1.5 mlir-translate：MLIR 的"出口翻译器"

mlir-translate 负责 MLIR ↔ 外部格式 的翻译。在方案二里，它执行 --mlir-to-llvmir，把 LLVM 方言 MLIR 翻译成 LLVM IR。

• 注意：mlir-translate --mlir-to-llvmir 只接受纯 LLVM 方言的 MLIR ，所以前面必须先通过 mlir-opt -convert-to-llvm 把所有高层方言都消掉。

1.6 mlir-runner：MLIR 的"JIT 执行器"

mlir-runner 是方案一的执行工具。它直接消费 LLVM 方言的 MLIR，在进程内 JIT 编译成机器码并立即执行，不落地任何中间文件。

详细流编译程如下：

---

一、 构建工具链

1.构建llvm-project

#下载
git clone https://github.com/llvm/llvm-project.git  

#配置参数（根据需要调整Native为其他需要支持的平台）
cd llvm-project
mkdir build
cd build
cmake -G Ninja ../llvm \
     -DLLVM_ENABLE_PROJECTS=mlir \
     -DLLVM_BUILD_EXAMPLES=ON \
     -DLLVM_TARGETS_TO_BUILD="Native;NVPTX;AMDGPU" \
     -DCMAKE_BUILD_TYPE=Release \
     -DLLVM_ENABLE_ASSERTIONS=ON \

#构建
cmake --build . --target check-mlir

#构建结束后，工具链在llvm-project/build/bin路径下

2.构建tensorflow

#下载
git clone https://github.com/tensorflow/tensorflow.git

#先通过指令获取tensorflow获取对应的LLVM对应版本
grep -r "LLVM_COMMIT" tensorflow/ --include="*.bzl" --include="*.bazel" --include="*.py"

#根据获取到的commit id，切换LLVM的对应版本
cd llvm-project
git checkout d5efcd0ec41d0f425076e3dd738a68d333ad610f

LLVM_SRC=<使用自己的llvm-project地址>
touch ${LLVM_SRC}/BUILD.bazel ${LLVM_SRC}/WORKSPACE


#下面指令指定了python版本，如果已经满足，可以不指定
cd tensorflow
bazel build --override_repository="llvm-raw=${LLVM_SRC}" --repo_env=HERMETIC_PYTHON_VERSION=3.12 -c opt tensorflow/compiler/mlir:tf-opt

bazel build --override_repository="llvm-raw=${LLVM_SRC}" --repo_env=HERMETIC_PYTHON_VERSION=3.12 -c opt tensorflow/compiler/mlir:tf-mlir-translate

#生成的文件在
ls bazel-bin/tensorflow/compiler/mlir

3.构建mlir-hlo

#下载
git clone https://github.com/openxla/xla.git
cd xla 

#按照上面的版本号，先将llvm-project切到对应版本，再执行构建流程
grep -r "LLVM_COMMIT" ./ --include="*.bzl" --include="*.bazel" --include="*.py"

#构建
bazel build --override_repository="llvm-raw=${LLVM_SRC}" --repo_env=HERMETIC_PYTHON_VERSION=3.12 -c opt xla/mlir_hlo:mlir-hlo-opt

#生成的文件在
bazel-bin/xla/mlir_hlo/mlir-hlo-opt --help

二、编译通路

2.pbtxt翻译为MLIR

输入的文件为：https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/add.pbtxt

node {
  name: "Add"
  op: "Add"
  input: "input0"
  input: "input1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "input0"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "input1"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
versions {
  producer: 27
}

#pbtxt翻译为MLIR
#注册了 graphdef-to-mlir 这个翻译规则，由 tf-mlir-translate -graphdef-to-mlir 触发
#-tf-input-arrays、-tf-input-shapes、-tf-output-arrays 这些参数，都是在告诉翻译器如何把 GraphDef 的 Placeholder 和输出节点映射成 MLIR 函数的参数和返回值
/Users/lizutuan/Desktop/code/tensorflow/bazel-bin/tensorflow/compiler/mlir/tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false add.pbtxt -tf-input-arrays=input0,input1 -tf-input-data-types=DT_INT32,DT_INT32 -tf-input-shapes=10:10 -tf-output-arrays=Add -o add.mlir

#输出内容为：
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 27 : i32}} {
  func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<*xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "input0,input1", outputs = "Add"}} {
    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Add"(%arg0, %arg1) {device = ""} : (tensor<10xi32>, tensor<10xi32>) -> tensor<*xi32>
      tf_executor.fetch %outputs : tensor<*xi32>
    }
    return %0 : tensor<*xi32>
  }
}

3.从"tf_executor Dialect"一路下降到"linalg Dialect"

#1.tf_executor Dialect -> tf Dialect
#-tf-executor-to-functional-conversion把 tf_executor 方言的"图执行语义"提升为函数式 IR。
#输入是 tf_executor.graph + tf_executor.island，输出是普通的 func.func + tf 算子。
#这一步是为了去掉 GraphDef 的执行包装，让 IR 变成更纯粹的"函数 + 算子"形式，方便后续做算子级映射。
/Users/lizutuan/Desktop/code/tensorflow/bazel-bin/tensorflow/compiler/mlir/tf-opt -tf-executor-to-functional-conversion add.mlir -o add-func.mlir

#输出内容为
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 27 : i32}} {
  func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<*xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "input0,input1", outputs = "Add"}} {
    %0 = "tf.Add"(%arg0, %arg1) {device = ""} : (tensor<10xi32>, tensor<10xi32>) -> tensor<*xi32>
    return %0 : tensor<*xi32>
  }
}

#2.tf Dialect -> hlo Dialect
#-tf-to-hlo-pipeline 把 TensorFlow 算子逐个映射到 HLO（MHLO）算子。这是一个聚合 pipeline，内部包含一系列 pattern，负责把 tf.* 重写成 mhlo.*
/Users/lizutuan/Desktop/code/tensorflow/bazel-bin/tensorflow/compiler/mlir/tf-opt --tf-to-hlo-pipeline add-func.mlir -o add-mhlo.mlir

#输出内容
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 27 : i32}} {
  func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "input0,input1", outputs = "Add"}} {
    %0 = mhlo.add %arg0, %arg1 : tensor<10xi32>
    return %0 : tensor<10xi32>
  }
}

#3.hlo Dialect -> stablehlo Dialect
/Users/lizutuan/Desktop/code/xla/bazel-bin/xla/mlir_hlo/mlir-hlo-opt add-mhlo.mlir --hlo-legalize-to-stablehlo -o add-stablehlo.mlir

#输出内容
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 27 : i32}} {
  func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "input0,input1", outputs = "Add"}} {
    %0 = stablehlo.add %arg0, %arg1 : tensor<10xi32>
    return %0 : tensor<10xi32>
  }
}

#4.stablehlo Dialect -> linalg Dialect
/Users/lizutuan/Desktop/code/xla/bazel-bin/xla/mlir_hlo/mlir-hlo-opt add-stablehlo.mlir --stablehlo-legalize-to-linalg --one-shot-bufferize="bufferize-function-boundaries" --canonicalize --cse -o add-linalg.mlir

#输出内容
#map = affine_map<(d0) -> (d0)>
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 27 : i32}} {
  func.func @main(%arg0: memref<10xi32, strided<[?], offset: ?>>, %arg1: memref<10xi32, strided<[?], offset: ?>>) -> memref<10xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "input0,input1", outputs = "Add"}} {
    %alloc = memref.alloc() alignment = 64 : memref<10xi32>
    linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : memref<10xi32, strided<[?], offset: ?>>, memref<10xi32, strided<[?], offset: ?>>) outs(%alloc : memref<10xi32>) {
    ^bb0(%in: i32, %in_0: i32, %out: i32):
      %0 = arith.addi %in, %in_0 : i32
      linalg.yield %0 : i32
    }
    return %alloc : memref<10xi32>
  }
}

上述的部分是有tensorflow和xla社区维护，后续的是由MLIR社区维护

4.修改

手动修改代码的原因并不复杂：我们最初拿到的是计算图，它只描述中间的计算过程，并不包含输入数据。即使 lowering 到 linalg Dialect，得到的仍然只是计算图的函数式表达；要让计算图真正运行起来，必须为它提供输入数据，因此需要手动调整代码。

先看 lowering 到 linalg Dialect 后的形态；

可以看到，计算图已经变成一个 main 函数，输入输出均为 memref<10xi32>，核心计算通过 addi 完成两个输入相加。

接下来有两种运行方式：

1. 使用 MLIR 的 JIT 即时编译引擎（mlir-cpu-runner）直接执行；

2. 用 C 程序调用 MLIR 编译出的函数，生成可执行文件。

这两种方式对代码修改的要求不同，下面分别说明。

4.1JIT 即时编译引擎

修改后的文件名：add-linalg-jit.mlir，内容如下：

#map = affine_map<(d0) -> (d0)>
func.func private @printMemrefI32(memref<*xi32>)

func.func @alloc_filled_i32(%arg0: i32) -> memref<10xi32> {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0 = memref.alloc() : memref<10xi32>
  scf.for %idx = %c0 to %c10 step %c1 {
    memref.store %arg0, %0[%idx] : memref<10xi32>
  }
  return %0 : memref<10xi32>
}

func.func @main() {   
  %0 = memref.alloc() : memref<10xi32>
  %i2 = arith.constant 2 : i32
  %i3 = arith.constant 3 : i32
  %arg0 = call @alloc_filled_i32(%i2) : (i32) -> memref<10xi32>
  %arg1 = call @alloc_filled_i32(%i3) : (i32) -> memref<10xi32>

  linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : memref<10xi32>, memref<10xi32>) outs(%0 : memref<10xi32>) {
  ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):  // no predecessors
    %1 = arith.addi %arg2, %arg3 : i32
    linalg.yield %1 : i32
  }

  %U = memref.cast %0 : memref<10xi32> to memref<*xi32>
  call @printMemrefI32(%U) : (memref<*xi32>) -> ()

  memref.dealloc %arg0 : memref<10xi32>
  memref.dealloc %arg1 : memref<10xi32>
  memref.dealloc %0 : memref<10xi32>
  return
}

运行指令如下：

#-e main：指定入口函数名。
#-entry-point-result=void：声明入口函数没有返回值（你 main 内部用 printMemrefI32 自己打印）
#-shared-libs=.../libmlir_runner_utils.dylib：加载运行时库，让 JIT 代码能解析 printMemrefI32 等符号


/Users/lizutuan/Desktop/code/llvm-project/build/bin/mlir-opt add-linalg-jit.mlir -convert-linalg-to-loops -lower-affine -convert-scf-to-cf -convert-to-llvm | /Users/lizutuan/Desktop/code/llvm-project/build/bin/mlir-runner -e main -entry-point-result=void -shared-libs=/Users/lizutuan/Desktop/code/llvm-project/build/lib/libmlir_runner_utils.dylib
Unranked Memref base@ = 0x7990cc390 rank = 1 offset = 0 sizes = [10] strides = [1] data = 
[5,  5,  5,  5,  5,  5,  5,  5,  5,  5]

4.2使用C程序调用MLIR的函数编译出可执行文件

除了前面介绍的 JIT 即时编译方式，还有另一条路径：用 C 程序调用 MLIR 函数，最终链接生成可执行文件，为了避免跟c程序的main函数重名冲突，需要将add-linalg.mlir中的main改为tf_add，再执行以下操作。

# 1. 降级并生成 C 包装器
/Users/lizutuan/Desktop/code/llvm-project/build/bin/mlir-opt add-linalg.mlir -expand-strided-metadata -convert-linalg-to-loops -lower-affine -convert-scf-to-cf -convert-cf-to-llvm -convert-arith-to-llvm -finalize-memref-to-llvm -convert-func-to-llvm -llvm-request-c-wrappers -reconcile-unrealized-casts -o add-linalg-llvm.mlir

# 2. 生成 LLVM IR
/Users/lizutuan/Desktop/code/llvm-project/build/bin/mlir-translate --mlir-to-llvmir add-linalg-llvm.mlir -o add-linalg.ll

# 3. 编译为目标文件
clang -c add-linalg.ll -o add-linalg.o

对应c程序如下：

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>

// 1. 定义与 MLIR 一致的 MemRef 描述符
typedef struct {
    int32_t *allocated;
    int32_t *aligned;
    intptr_t offset;
    intptr_t sizes[1];    // rank = 1
    intptr_t strides[1];  // rank = 1, 动态 stride
} MemRef_descriptor;

// 2. 声明生成的 C 接口
// 注意：返回的 memref 变成了第一个参数（指针）
extern void _mlir_ciface_tf_add(
    MemRef_descriptor *result,
    MemRef_descriptor *arg0,
    MemRef_descriptor *arg1
);

int main() {
    // 3. 准备数据
    int32_t data0[10], data1[10];
    for (int i = 0; i < 10; ++i) {
        data0[i] = 2;
        data1[i] = 3;
    }

    // 4. 构造输入 MemRef 描述符
    //    对于动态 strided memref，需要显式设置 offset 和 strides
    MemRef_descriptor arg0 = {
        .allocated = data0,
        .aligned = data0,
        .offset = 0,
        .sizes = {10},
        .strides = {1}   // 连续数组，stride 为 1
    };
    MemRef_descriptor arg1 = {
        .allocated = data1,
        .aligned = data1,
        .offset = 0,
        .sizes = {10},
        .strides = {1}
    };

    // 5. 准备输出 MemRef
    //    MLIR 函数内部会 allocate 结果，C 侧只需提供一个描述符来接收它
    int32_t out_data[10];
    MemRef_descriptor result = {
        .allocated = out_data,
        .aligned = out_data,
        .offset = 0,
        .sizes = {10},
        .strides = {1}
    };

    // 6. 调用 C 接口
    _mlir_ciface_tf_add(&result, &arg0, &arg1);

    // 7. 打印结果
    for (int i = 0; i < 10; ++i) {
        printf("%d ", result.aligned[i]);
    }
    printf("\n");

    return 0;
}

# 编译 C 主程序
clang -c main.c -o main.o

# 链接
clang main.o add-linalg.o -o add-main

#运行
./add-main 
5 5 5 5 5 5 5 5 5 5 

---

参考

https://zhuanlan.zhihu.com/p/328559928

https://zhuanlan.zhihu.com/p/328993481