Hands-On Practical MLIR Tutorial

Kexing Zhou（周可行）

北京大学

zhoukexing@pku.edu.cn

• 1. MLIR 简介
• 1.1. MLIR 编译管线
• 1.2. 常见的 Dialect
• 1.3. insight：“及时做优化”
• 1.4. MLIR 的用处
• 1.5. MLIR 的缺点
• 1. MLIR 基本用法
• 2.1. IR 基本结构
• 2.2. MLIR 基本工程模板
  • 2.2.1. 配置 clangd
• 2.3. MLIR 的读入、输出
• 2.4. 用代码生成 MLIR
• 1. MLIR Op 的结构
• 3.1. Attribute 和 Operand
• 3.2. Attribute, Value 和 Type
• 1. MLIR 的类型转换
• 4.1. Op 的类型转换
• 4.2. Type / Attribute 的类型转换
• 1. MLIR 的图结构
• 5.1. MLIR 数据流图结构
• 5.2. MLIR 数据流图的遍历与修改
• 5.3. MLIR 控制流图的遍历与修改
• 1. 基本的 Dialect 工程
• 6.1. TableGen 工程模板
• 6.2. Tablegen Language Server
• 6.3. IR 的默认定义与实现
  • 6.3.1. TableGen 文件
  • 6.3.2. 头文件
  • 6.3.3. 库文件
  • 6.3.4. 程序入口
• 1. TableGen Op 定义详解
• 7.1. Attribute、Type、Constraint
  • 7.1.1. 内置 Attribute
  • 7.1.2. 内置的 Type
  • 7.1.3. 为什么 Attribute 和 Type 都是 Constraint
• 7.2. Verifier：发现IR错误
  • 7.2.1. emitError
  • 7.2.2. LogicalResult
• 7.3. Variadic：可变参数
  • 7.3.1. 多个可变参数：AttrSizedOperandSegments
• 7.4. AssemblyFormat：更易读的输出
  • 7.4.1. 常用关键字
  • 7.4.2. 额外 attr 字典
  • 7.4.3. 输出 type
  • 7.4.4. 可选输出：Optional、UnitAttr
• 7.5. Builder：自定义 create 函数
  • 7.5.1. 默认Builder
  • 7.5.2. 自定义builder
• 7.6. 自定义函数
  • 7.6.1. header target
• 7.7. 使用 Trait
  • 7.7.1. 内存副作用：SideEffectInterfaces
  • 7.7.2. 类型推断：InferTypeOpInterface
• 7.8. 函数：FunctionOpTrait
  • 7.8.1. 定义 Return
  • 7.8.2. 定义 Function
  • 7.8.3. 定义 Call
• 1. 添加 Pass
• 8.1. Pass 工程模板
• 8.2. Pass 定义详解
  • 8.2.1. 指定 Pass 在哪个 Op 上运行
  • 8.2.2. 带参数的 Pass
• 8.3. 简单的 DCE Pass 实现
  • 8.3.1. 定义
  • 8.3.2. 实现
• 1. Pattern Rewrite
• 9.1. Pattern Rewrite
  • 9.1.1. 描述 Pattern
  • 9.1.2. 调用 Pattern
  • 9.1.3. Depedent Dialect & Linking
• 9.2. Dialect Convertion (Type Conversion)
  • 9.2.1. TypeConverter
  • 9.2.2. Conversion Pattern：自动做 Operand 的类型转换
  • 9.2.3. 类型转换的细节与 Debug
  • 9.2.4. 使用自己的 materialization
• 9.3. 使用 MLIR 里已有的 Pattern 做多步转换
• 1. 自定义 Type
• 1. TIPS
• 11.1. 如何找头文件、找想要的函数
• 11.2. 如何找需要连接的库
• 11.3. 如何加快编译速度
• 11.4. 去 MLIR 里抄代码
• 1. MLIR 的批判：C++ v.s. Rust
• 1. Issue & Reply

1. MLIR 简介

1.1. MLIR 编译管线

MLIR 在于设计一套可复用的编译管线，包括可复用的 IR、Pass 和 IO 系统。在 IR 中，多个 Dialect 可以混合存在。MLIR 已经定义好了一套 Dialect Translation Graph：

1.2. 常见的 Dialect

MLIR 的 Dialect 是相对独立的，下面列举一些常见的 dialect：

1. func：处理函数的dialect，包含的函数定义、调用、返回等基本操作
2. arith：处理加减乘除移位等各种运算
   • math：更复杂的运算，如 log, exp, tan 等
3. affine：处理循环嵌套，实现了循环展开、多面体变换等一些算法
4. scf：(structured control flow) 结构化控制流，保留 for，if 等语句
   • cf：无结构控制流，只有条件跳转命令
5. llvm：LLVM IR 的 binding，可以直接翻译给 LLVM 做后续编译

MLIR的编译从高层次的 tensor 到 低层次的 scf,cf，每个阶段都是多个 dialect 的混合体，每次 lowering 往往只针对一个 dialect 进行。

1.3. insight：“及时做优化”

这里简单举例，dialect 是如何混合的。

例如，我是 pytorch，生成了一些神经网络，我想要表示这些运算：

• Tensor 是一块带 shape 的指针：使用 tensor dialect
• 简单的 elementwise 加减乘除：使用 arith dialect
• 复杂的 log、exp 等运算：使用 math dialect
• 矩阵线性代数运算：使用 linalg dialect
• 可能有一些控制流：使用 scf dialect
• 整个网络是一个函数：使用 func dialect

接下来，将其逐渐 lower 到 LLVM：

• 调用 Lowering Pass，把 tensor lowering 到 linalg，而其他的 dialect 不会改变。
• 继续调用 pass，直到把 linalg 转换到 affine -> scf -> cf，其他表示运算的 dialect 保留不变。
• 继续 Lowering，把 memref 转换为裸指针、arith 和 func 转换为 llvm 内置运算。
• 最后，所有非 llvm dialect 都被转换为了 llvm dialect，现在可以导出为 llvm ir 交给 llvm 继续编译。

可见，MLIR 编译有一个特点：不同 dialect 是独立的。

• 例如，在做循环展开等优化的时候，我不需要关心加法和减法可以合并；而在做算数表达式优化的时候，也不需要关心当前在哪个函数里边。

MLIR 可以从各个层次优化 IR：例如：

• 在 affine 层面，可以根据循环大小做展开，向量化
• 在 scf 层面，可以发现循环不变量
• 在 arith 层面，可以用算数恒等式优化代码

MLIR 的 insight 在于“及时做优化”。很明显，linalg 层次，我们很容易发现矩阵被转置了两次，但一旦 lower 到 scf，所有转置操作都变成循环，优化就很难进行了。

1.4. MLIR 的用处

我们使用 MLIR，主要也是想要复用别人已经写好的代码，一般包括：

• 复用已有 dialect 作为 输入，不用自己写前端。
  • 如 Polygeist 能把 C 翻译成 Affine Dialect，这样我们就不用写 C Parser
• 将已有 dialect 混入或作为输出。
  • 如 arith 等 dialect，可以直接集成起来，不需要自己写。
  • 要生成 binary 的时候，可以直接生成 LLVM Dialect，复用后端 LLVM 编译管线
• 复用已有的 Pass。
  • 常见的 Pass 如 CSE，DCE 可以复用
  • Dialect 专用 Pass，如循环展开，也可以复用

1.5. MLIR 的缺点

MLIR 也有缺点：

• 太过笨重，编译、链接时间长（可能会连接出上百M的文件）
• 可以用 lld 来加快链接速度，但依然很慢 见TIPS
• Dialect 定义极不灵活，定义较复杂 Op 时非常麻烦

2. MLIR 基本用法

2.1. IR 基本结构

MLIR 是 树形结构，每个节点是 Operation，Op 可以组成 Block，Block 组成 Region，而 Region 又可以嵌套在 Op 内部。

• Operation 指单个运算，运算内可以嵌套 Region
• Block 指基本块，基本块包含一个或多个 Operation
• Region 指区域，类似于循环体或函数体，包含若干 Block

MLIR 的基本块使用 “基本块参数” 来取代“phi函数”，如下面的例子：

• 块参数：每个基本块都带有参数，块内可以使用
• 终止符：每个基本块的一般为跳转或返回，跳转时需要附上块参数

module {
func.func @foo(%a: i32, %b: i32, %c: i32) -> i32 {
  %cmp = arith.cmpi "sge", %a, %b : i32
  cf.cond_br %cmp, ^add(%a: i32), ^add(%b: i32)
^add(%1: i32):
  %ret = llvm.add %1, %c : i32
  cf.br ^ret
^ret:
  func.return %ret : i32
}
}

module: 默认情况下，mlir 最外层是 builtin.module，作为 IR 的根。

2.2. MLIR 基本工程模板

构建第一个 mlir 项目往往非常困难，下面给一个我常用的工程模板：

mlir-tutorial
├── install       # Install Prefix，把 MLIR 编译后安装到这里
├── llvm-project  # MLIR 项目
└── mlir-toy      # 自己的 MLIR 工程

首先，按照 MLIR getting started 的方法，安装 MLIR。注意，安装的时候要设置 PREFIX 为 install 目录，如下面所示，和 getting start 上的略有区别：

git clone https://github.com/llvm/llvm-project.git
cd llvm-project
git checkout 186a4b3b657878ae2aea23caf684b6e103901162 # 本教程使用的版本
mkdir build && cd build
cmake -G Ninja ../llvm \
  -DCMAKE_INSTALL_PREFIX=/mlir-tutorial/install \
  -DLLVM_ENABLE_PROJECTS=mlir \
  -DLLVM_BUILD_EXAMPLES=ON \
  -DLLVM_TARGETS_TO_BUILD="Native;NVPTX;AMDGPU" \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=ON

在 build 完之后，安装到 prefix

ninja install

现在，mlir 把所有二进制文件、库文件都装到了 install 目录下。现在可以 export PATH=/mlir-tutorial/install/bin:$PATH，方便调用 bin 里面的 mlir-opt 等程序。

install
├── bin
├── examples
├── include
├── lib
└── share

接下来，在 mlir-toy 里面建立一个简单的工程

mlir-toy
├── CMakeLists.txt
└── main.cpp

其中 CMakeLists.txt 文件写法比较固定：

cmake_minimum_required(VERSION 3.13.4)

project(mlir-toy VERSION 0.0.0)

set(CMAKE_EXPORT_COMPILE_COMMANDS ON) # 生成 compile_commands.json 便于代码高亮
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED YES)

find_package(MLIR REQUIRED CONFIG)
list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")

include(TableGen)
include(AddLLVM)
include(AddMLIR)
include(HandleLLVMOptions)

include_directories(${LLVM_INCLUDE_DIRS} ${MLIR_INCLUDE_DIRS})

add_executable(mlir-toy main.cpp)

在 main.cpp 里面写一个 main 函数，然后先 build 一下。注意，必须要写上 CMAKE_INSTALL_PREFIX，这样cmake 可以自动找到 MLIR。

cd build
cmake .. -GNinja -DCMAKE_INSTALL_PREFIX=/mlir-tutorial/install
ninja

2.2.1. 配置 clangd

使用 vscode 默认的 lint 工具跑 mlir 会非常卡，建议使用 clangd。

• 在扩展里安装 clangd 插件
• cmake 的时候加上 -DCMAKE_EXPORT_COMPILE_COMMANDS=ON（上面的CMakeLists.txt 已经加上了）
• 有时候要吧 compile_commands.json 拷贝到工程根目录，或者在 vscode 设置里配置一下
• 一旦发现高亮炸了，vscode里 Ctrl + Shift + P，输入 clangd: restart language server
• 有时候，mlir 的编译选项与 clangd 冲突，在 mlir-toy 目录下建立 .clangd 文件，去掉相关的选项： yaml CompileFlags: Remove: - -fno-lifetime-dse

2.3. MLIR 的读入、输出

测试用 mlir：

func.func @test(%a: i32, %b: i32) -> i32 {
  %c = arith.addi %a, %b : i32
  func.return %c : i32
}

最简单的读入输出：

#include "mlir/IR/AsmState.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "llvm/Support/raw_ostream.h"

using namespace mlir;

int main(int argc, char ** argv) {
  MLIRContext ctx;
  // 首先，注册需要的 dialect
  ctx.loadDialect<func::FuncDialect, arith::ArithDialect>();
  // 读入文件
  auto src = parseSourceFile<ModuleOp>(argv[1], &ctx);
  // 输出dialect，也可以输出到 llvm::errs(), llvm::dbgs()
  src->print(llvm::outs());
  // 简单的输出，在 debug 的时候常用
  src->dump();
  return 0;
}

需要连接上所有依赖的文件：

target_link_libraries(
  ex1-io
  MLIRIR
  MLIRParser
  MLIRFuncDialect
  MLIRArithDialect
)

测试方式：

./ex1-io ../ex1-io/ex1.mlir

2.4. 用代码生成 MLIR

#include "mlir/IR/AsmState.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Support/FileUtilities.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "llvm/Support/raw_ostream.h"

using namespace mlir;

int main(int argc, char ** argv) {
  MLIRContext ctx;
  ctx.loadDialect<func::FuncDialect, arith::ArithDialect>();

  // 创建 OpBuilder
  OpBuilder builder(&ctx);
  auto mod = builder.create<ModuleOp>(builder.getUnknownLoc());

  // 设置插入点
  builder.setInsertionPointToEnd(mod.getBody());

  // 创建 func
  auto i32 = builder.getI32Type();
  auto funcType = builder.getFunctionType({i32, i32}, {i32});
  auto func = builder.create<func::FuncOp>(builder.getUnknownLoc(), "test", funcType);

  // 添加基本块
  auto entry = func.addEntryBlock();
  auto args = entry->getArguments();

  // 设置插入点
  builder.setInsertionPointToEnd(entry);

  // 创建 arith.addi
  auto addi = builder.create<arith::AddIOp>(builder.getUnknownLoc(), args[0], args[1]);

  // 创建 func.return
  builder.create<func::ReturnOp>(builder.getUnknownLoc(), ValueRange({addi}));
  mod->print(llvm::outs());
  return 0;
}

如何寻找 builder.create 的参数：builder.create 内部是调用 Op::build 函数的，你可以 Ctrl + 鼠标点击找到 func::FuncOp 的定义，然后找里面的 build 函数，看参数表。

3. MLIR Op 的结构

MLIR 的一个 Operation 里可以包含下面的一些东西：

• Operand：这个 Op 接受的操作数
• Result：这个 Op 生成的新 Value
• Attribute：可以理解为编译器常量
• Region：这个 Op 内部的 Region

MLIR 中，Attribute 是高度灵活的，允许插入原来不存在的 attr，允许不同 dialect 互相插入 attribute。

3.1. Attribute 和 Operand

Attribute 和 Operand 有一些区别。Attribute 指的编译器已知的量，而 Operand 指只有运行时才能知道的量。

如下面的这个Op，0 是一个 Attribute 而不是一个 Operand

%c0 = arith.constant 0 : i32

3.2. Attribute, Value 和 Type

Value 必然包含 Type，Type 也可以作为 Attribute 附加在 Operation 上。

例如函数 Op，虽然 %a, %b 出现在了参数表里，但它们实际上是函数类型的一部分，算是 Type Attribute。

func.func @test(%a: i32, %b: i32) -> i32 {
  %c = arith.addi %a, %b : i32
  func.return %c : i32
}

• mlir-opt --mlir-print-op-generic 来打印这里的代码，得到下面的代码。参数名被隐去，只有 function_type 作为 attribute 保留了下来。 mlir "builtin.module"() ({ "func.func"() <{function_type = (i32, i32) -> i32, sym_name = "test"}> ({ ^bb0(%arg0: i32, %arg1: i32): %0 = "arith.addi"(%arg0, %arg1) : (i32, i32) -> i32 "func.return"(%0) : (i32) -> () }) : () -> () }) : () -> ()

4. MLIR 的类型转换

4.1. Op 的类型转换

MLIR 的所有 Op 都有一个统一的储存格式，叫 Operation。Operation 里面存了 OpName 和所有的 operands, results, attributes 和其它的东西。

用户定义的 arith.addi 等等 Op，本质上都是 Operation 的指针。但与 Operation* 不同的是，AddIOp 定义了 Operation 里储存的数据的解释方式。如 AddOp，自己是一个 Operation 的指针