🧐 创建一个Dialect

老实讲, 我在这一篇涉及到的内容上踩的坑还是很多的. 也不是什么很大的坑, 但解决完会有点红温, 为什么会把时间浪费在这上面.

🫠 文件目录结构

MLIR项目的文件结构目录稍微有点复杂, 下面是我使用的一个简化版本

mlir-toy

include

CMakeLists.txt
toy

CMakeLists.txt
ToyDialect.h
ToyDialect.td
ToyOps.h
ToyOps.td
Toy.td

ToyDialect.cpp
CMakeLists.txt

main.cpp
CMakeLists.txt

☝️ 创建Dialect

MLIR提供了TableGen工具, 可以自动生成一些框架代码, 我们优先采用这种方式.

include/toy/ToyDialect.td

#ifndef TOY_DIALECT_TD // include guard
#define TOY_DIALECT_TD

include "mlir/IR/OpBase.td"

def Toy_Dialect : Dialect {
	let name = "toy";
	let cppNamespace = "toy";
	let summary = "Toy dialect";
	let description = [{
		"This is a toy dialect for demonstration purposes."
	}];
}

#endif // TOY_DIALECT_TD

include/toy/ToyOps.td

#ifndef TOY_OPS_TD
#define TOY_OPS_TD

include "mlir/IR/OpBase.td"
include "toy/ToyDialect.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

class Toy_Op<string mnemonic, list<Trait> traits = []> : Op<Toy_Op, mnemonic, traits>;

def AddOp : Toy_Op<"add", [Pure]> {  // 注意AddOp不能写成Add_Op, 否则生成的代码会有名称问题
	let summary = "Add two integers";
	let description = [{
		"This operation adds two integers together."
	}];

	let arguments = (ins AnyInteger:$lhs, AnyInteger:$rhs);
	let results = (outs AnyInteger:$result);
}

#endif // TOY_OPS_TD

然后在Toy.td文件中添加这两个文件的include语句

include/toy/Toy.td

include "toy/ToyDialect.td"
include "toy/ToyOps.td"

在这一级的CMakeLists.txt中添加生成目标

include/toy/CMakeLists.txt

set(LLVM_TARGET_DEFINITIONS Toy.td)
mlir_tablegen(ToyOps.h.inc -gen-op-decls)
mlir_tablegen(ToyOps.cpp.inc -gen-op-defs)
mlir_tablegen(ToyDialect.h.inc -gen-dialect-decls)
mlir_tablegen(ToyDialect.cpp.inc -gen-dialect-defs)
add_public_tablegen_target(MLIRToyIncGen)

在上一级的CMakeLists.txt中添加这个目录

include/CMakeLists.txt

add_subdirectory(toy)

然后在顶级CMakeLists.txt中添加这个目标作为顶层可执行文件的生成依赖

CMakeLists.txt

# 其他内容...
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) 
include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) 

add_subdirectory(include) 
# add_executable(MLIR-Toy main.cpp) 在这里
add_dependencies(MLIR_Toy MLIRToyIncGen)

接着可以先构建程序, 让Tablegen去生成代码. 这个时候主文件main.cpp编译成不成功都无所谓. 生成代码后在include/toy下面加两个我们自己写的头文件, 用来包含它自动生成的头文件.

include/toy/ToyDialect.h

#ifndef TOY_DIALECT_H
#define TOY_DIALECT_H

#include "mlir/IR/Dialect.h"
#include "mlir/IR/BuiltinTypes.h"
#include "toy/ToyDialect.h.inc"

#endif // TOY_DIALECT_H

include/toy/ToyOps.h

#ifndef TOY_OPS_H
#define TOY_OPS_H

#include "toy/ToyDialect.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
// 这里要将ToyOps.td中引入的头文件都要引进来

#define GET_OP_CLASSES
#include "toy/ToyOps.h.inc"

#endif // TOY_OPS_H

这里需要注意的是两行.h.inc文件的引入. .h.inc文件本身不是一个完备的编译单元, 它必须依赖于一个上下文才能编译成功. 因此这两个.h.inc文件必须要注意引入的位置, 也必须要引入, 否则会报各种符号未定义或者链接错误.

接下来新增一个文件夹lib, 用来储存方言和各种操作的实现代码. 这里我们先创建一个ToyDialect.cpp即可, 暂时不需要太多代码.

lib/ToyDialect.cpp

#include "toy/ToyDialect.h"
#include "toy/ToyDialect.cpp.inc"
#include "toy/ToyOps.h"

#define GET_OP_CLASSES
#include "toy/ToyOps.cpp.inc"

using namespace mlir::toy;

void ToyDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "toy/ToyOps.cpp.inc"
  >();
}

值得注意的是MLIR中非常常见的用法, 先#define一个宏然后利用这个宏从inc文件中读取一部分出来. 另外注意一下在ToyDialect.cpp和我们自己写的两个头文件中包含inc文件的方式. mlir-tablegen会生成.h.inc文件和.cpp.inc文件两种. 前者只应该在我们自己编写的.h文件中引入, 后者只应该在我们自己编写的.cpp文件中引用. 在我们自己编写的.cpp文件中除了我们自己写的.h文件外, 只应该还有.cpp.inc文件. 如果随意引用.h.inc文件, 会导致编译器无法找到符号或者重定义之类的错误.

另外为了减少链接出错的可能, 在编译之前先检查一下所有自动生成的.h.inc和.cpp.inc是不是都被包含过了. 如果有生成的文件没有被包含在任何一个项目文件中的话大概率是要未定义引用的.

在lib目录下创建一个CMakeLists.txt, 用来编译这个文件

lib/CMakeLists.txt

add_mlir_library(MLIRToy ToyDialect.cpp DEPENDS MLIRToyIncGen) # 生成的库名称`MLIRToy`, 后面要用到

在顶级CMakeLists.txt中添加这个目录

CMakeLists.txt

add_subdirectory(lib)

如果一切顺利, 你可以先在main.cpp中写一个空的main函数, 然后启动编译. 这个时候应该能编译并且链接成功. 如果链接失败了或者报出未定义符号, 检查一下你的包含头文件方式有没有问题(漏包含或者顺序错了).

👉 简单的测试

我们现在可以手写一段MLIR代码, 然后用MLIR提供的mlir-opt程序入口读取这段代码并做一些内置的优化.

test.mlir

func.func @test(%a: i32, %b: i32) -> i32 {
	%c = "toy.add"(%a, %b): (i32, i32) -> i32
	%d = "toy.add"(%a, %b): (i32, i32) -> i32
	%e = "toy.add"(%c, %d): (i32, i32) -> i32
	%f = "toy.add"(%d, %c): (i32, i32) -> i32
	func.return %f : i32
}

可以看出来这段代码有可以优化的地方, 返回值%c和%d是重复计算的, 只要计算一次就足够了. %e和%f也是重复计算的, 但是它们两个参数的位置相反. 接下来我们使用mlir-opt提供的可复用程序入口, 对这段代码调用两个内置的优化.

main.cpp

#include "mlir/Tools/mlir-opt/MlirOptMain.h"
#include "mlir/Transform/Passes.h"
#include "toy/ToyDialect.h"
#include "toy/ToyOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"

using namespace mlir;

int main(int argc, char **argv) {
	DialectRegistry registry;
	registry.addDialect<toy::ToyDialect, func:FuncDialect>();
	registerCSEPass();
	registerCanonicalizerPass();
	return asMainReturnCode(MlirOptMain(argc, argv, "mlir-opt", registry));
}

在顶级CMakeLists.txt中添加要链接的库

CMakeLists.txt

target_compile_options(MLIR-Toy PRIVATE -fno-rtti) # 如果你在构建MLIR的时候没有开启LLVM_RTTI

target_link_libraries(MLIR-Toy
        PRIVATE
        MLIRToy # 这是`lib`中生成的库
        MLIRIR
        MLIRParser
        MLIRSupport 
        MLIROptLib 
        MLIRTransforms 
        MLIRArithDialect
        MLIRFuncDialect 
        MLIRDialect 
		-fsanitize=address # 如果你在构建MLIR的时候开启了`LLVM_USE_SANITIZER=Address`
)

然后应该可以顺利编译成功了. 编译出来的可执行文件可以通过--help参数看用法.

./build/MLIR-Toy --help
./build/MLIR-Toy /path/to/test.mlir
./build/MLIR-Toy /path/to/test.mlir -canonicalize # 启用正规化
./build/MLIR-Toy /path/to/test.mlir -cse # 启用CSE优化

可以看到输出的结果. 启用正规化后, %e和%f被合并为同一个操作. 启用CSE(公共子表达式消除)后, test.mlir中所有重复的计算都被合并, 最后的优化结果是

module {
  func.func @test(%arg0: i32, %arg1: i32) -> i32 {
    %0 = "toy.add"(%arg0, %arg1) : (i32, i32) -> i32
    %1 = "toy.add"(%0, %0) : (i32, i32) -> i32
    return %1 : i32
  }
}

符合我们的预期. 关于CSE和正规化不多提. 在这里我们借用了func.func这个定义函数的操作, 它来自Func方言. 在MLIR中, 定义一个函数本身也是一个操作, 在Toy方言中我们也可以定义一个toy.func实现类似的功能, 在ToyOps.td中添加func操作并继承相关的接口即可. 类似的, func.return也是一个操作.

我们能直接对add操作使用CSE和正规化的前提是, 指定了add具有Pure属性. 这个属性表明操作的结果只依赖输入, 也就是说只要输入相同, 每一次操作的输出都是一样的. 可以试着在ToyOps.td中定义add的时候去掉这个属性, 然后重新优化test.mlir. 会发现CSE和正规化不会对代码做任何修改.