MLIR-tutorial学习笔记

最近在学习 MLIR，一开始学的是 MLIR 官方的 Toy Tutorial，但总感觉不得要领。后面在 Github上看到了北京大学周可行写的 MLIR-tutorial，非常清晰易懂，成功跟着做了一遍。在这里记录一下我在跟教程的过程中遇到的一些问题。顺便一提，大模型真的是一位好老师，之前我学习的是大模型没怎么学过的昇腾编程，所以用不上大模型的强大能力，如今在学 MLIR 的时候才体会到大模型的强大助学能力，对代码有什么不理解或者疑问，直接问它，它都能够解释得非常清楚。回忆起当年学习操作系统时面对浩如烟海的 OSWiki 时的无所适从，不禁感叹一句：科学的力量真伟大。

我使用的是 GitHub llvm-project 2025 年 6 月 9 日的 main 分支，commit 编号为 66911b7546f9afcf5f2b842ef0b9a39788dfef39

内容

引用 TableGen 生成的文件

首先是关于 TableGen 生成的头文件和代码文件的引用问题。MLIR 一个比较有特色的写法是通过宏定义开关来控制引用的内容，例如

#define GET_OP_CLASSES
#include "Toy.h.inc"

这段代码用来引用 Toy.h.inc 文件里关于 Toy 方言定义的内容。但有一点要注意，每次打开宏定义开关引用完文件后开关会自动关闭，也就是说，如果如果要引用两个具有相同宏定义开关的文件，就得写两次宏定义，不能省略，比如要同时引用 Toy.h.inc 和 Toy.cpp.inc 里的内容，就需要这样写：

#define GET_OP_CLASSES
#include "Toy.h.inc"
#define GET_OP_CLASSES
#include "Toy.cpp.inc"

#ifdef GET_OP_CLASSES
#undef GET_OP_CLASSES
// …
#endif

另外，这些头文件和代码文件的引用也有顺序，.h.inc 文件需要在 .cpp.inc 文件前引用，类型和方言定义相关的头文件需要在算子定义的头文件之前引用。

ConstantOp 相关

%a = "toy.const"() { value = 0 } : () -> i32

如果不加 assemblyFormat，就必须按照原始写法。当按照教程里定义了 ConstantOp 的 assemblyFormat 以及自动类型推断后就可以这样写：

%a = toy.const 0

这个时候 %a 默认类型推断结果是 i64 类型，如果要让其被推断 i32 类型，可在 MLIR 程序里这样写：

%a = toy.const 0: i32

这里 0 : i32 被当成了一个整体，表示 value 的值是 i32 类型的 0，所以可以正确解析。

在第十章里面，教程自定义了一个ToyInteger类型，并且将 toy::ConstantOp 的结果类型设置为该类型。这个时候就不能写 %a = toy.const 0 或者 %a = toy.const 0: !toy.int<32> 了，因为 MLIR 无法将 0 解释成 ToyInteger。正确方法有两种，一是使用原始写法：

%a = "toy.const"() { value = 0 } : () -> !toy.int<32>

相当于强制指定了 %a 的类型为 ToyInteger，不用自动推断。另一种方法是直接把自动推断关了，即将 toy::ConstantOp 的定义改为：

def ConstantOp : ToyOp<"const", [Pure]> {
  let summary = "const operation";
  let arguments = (ins APIntAttr:$value);
  let results = (outs ToyInteger:$result);
  let assemblyFormat = "$value attr-dict `:` type($result)";
}

%a = toy.const 0: i32: !toy.int<32>

CallOp 定义

def CallOp : ToyOp<"call", [CallOpInterface]> {
  let summary = "call operation";
  let arguments = (ins SymbolRefAttr:$callee, Variadic<AnyType>:$arg_operands);
  let results = (outs AnyType:$result);
  let assemblyFormat = "$callee `(` $arg_operands `)` attr-dict `:` functional-type($arg_operands, results)";
  let extraClassDeclaration = [{
    mlir::CallInterfaceCallable getCallableForCallee() {
      return getCalleeAttr();
    }
    void setCalleeFromCallable(mlir::CallInterfaceCallable callee) {
      setCalleeAttr(callee.get<mlir::SymbolRefAttr>());
    }
  }];
}

原因是在这个 commit 里，为了提高 CallOpInterface 和 CallableOpInterface（FuncOp 依赖的Interface）的一致性，维护者给 CallOpInterface 也添加了 $arg_attrs 和 $res_attrs 这两个属性，所以如果用了这个 Interface 的自定义 Op 没有定义这两个属性，就会报错。

要解决这个问题，在定义 toy::CallOp 的参数时把这俩属性加上去就可以了：

let arguments = (ins SymbolRefAttr:$callee,
    OptionalAttr<DictArrayAttr>:$arg_attrs,
    OptionalAttr<DictArrayAttr>:$res_attrs,
    Variadic<AnyType>:$arg_operands);

在添加完这两个属性后，toy::CallOp 能正常编译，但第九章做 Op 替换的时候又会出现问题了，原因是现在 toy::CallOp 多了俩参数，因此 TableGen 生成的 build 方法也会多俩参数，教程里在调用 replaceOpWithNewOp 没有加上那两个参数，转发到 build 方法后就会报参数不匹配的错误了。

直观的想法是在调用 replaceOpWithNewOp 方法前用 getArgAttrs 和 getResAttr 获取一下这两个参数，但要注意的是，因为我 $arg_attrs 和 $res_attrs 的类型是 OptionalAttr，所以 getArgAttrs 和 getResAttr 返回的类型是 std::optional<mlir::ArrayAttr>，然而，build 方法的对应参数却是没有 std::optional 包装的普通 mlir::ArrayAttr，即：

void FuncOp::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::StringAttr sym_name, ::mlir::TypeAttr function_type, /*optional*/::mlir::ArrayAttr arg_attrs, /*optional*/::mlir::ArrayAttr res_attrs) {

注释里都说了是 optional，但却不用 std::optional 包装起来，不是很能理解。不过改起来也不麻烦，调用 replaceOpWithNewOp 时传 getArgAttrs().value_or(nullptr) 和 getResAttrs().value_or(nullptr) 即可。

Pass 中删除 Op

for(auto v: reverse(opToRemove)) {
    v->erase();
}

对于教程中的实例 MLIR 程序，如果不反向删除也不会有什么问题，但如果 Op 里包含子块，比如 If、For 之类，由于获得opToRemove列表的walk函数是先序遍历，且删除父 Op的时候会同时删除内部的子 Op，因此如果正向删除的话可能会出现先删父 Op，此时子 Op 也被自动删除，再删子 Op，最终子 Op 被重复删除的情况。这可能会导致程序崩溃等问题。

RootUpdate 系列方法

LogicalResult matchAndRewrite(ReturnOp op, ReturnOpAdaptor adaptor, ConversionPatternRewriter & rewriter) const {
    auto data = adaptor.getData();
    rewriter.startRootUpdate(op);
    op.getDataMutable().assign(data);
    rewriter.finalizeRootUpdate(op);
    return success();
}

然而，MLIR 从版本 18.x 开始删掉了 RootUpdate 系列方法，改成了 OpModification 系列方法，不过用法没什么区别。其实这个地方并不一定需要使用这些方法对 Op 进行就地变换，和其他变换函数一样用 replaceOpWithNewOp 也是可以的。

Op 转换的具体流程

在第九章 ConvertToyToArithPass 类的 runOnOperation 方法定义了如何对 Op 进行转换，教程解释得比较粗糙，我查看了 MLIR 的官方文档，并自己做了下实验，大致看懂了代码。

首先，要转换一个 Op，有两个关键条件，一个是用来转换的 Pattern，另一个是用来判断转换后是否合法的 Pattern，如果这两个 Pattern 缺少一个，就不能转换。在 applyPartialConversion 的情况下就是放弃转换，继续处理下一个 Op，在 applyFullConversion 的情况下，如果有判断是否合法的 Pattern，并且判定不转换也合法，那就继续处理下一个 Op，否则就是直接报错。

用来转换的 Pattern 在教程中通过 RewritePatternSet 定义，然后代码里添加了两种 Pattern：

patterns.add<AddOpPat, SubOpPat, ConstantOpPat, ReturnOpPat, CallOpPat>(converter, &getContext());
populateFunctionOpInterfaceTypeConversionPattern<toy::FuncOp>(patterns, converter);

第一行往 Pattern 中添加了除了 toy::FuncOp 的 Op 转换泛式，第二行则专门针对 toy::FuncOp 添加转换泛式，原因是其他 Op 转换只需要转换操作数就行，而 FuncOp 还得考虑返回类型以及 Region 内部用到其参数的各个 Op，所以需要专门处理。

target.addLegalDialect<arith::ArithDialect>();
target.addDynamicallyLegalOp<toy::ReturnOp, toy::CallOp>([](Operation* f) {
        return llvm::all_of(f->getOperandTypes(),
                [](Type t) {return !isa<toy::ToyIntegerType>(t);});
        });
target.addDynamicallyLegalOp<toy::FuncOp>([](toy::FuncOp f) {
        return llvm::all_of(f.getArgumentTypes(),
                [](Type t) {return !isa<toy::ToyIntegerType>(t);});
        });

首先，利用 addLegalDialect 将 Arith 方言的 Op 全部定义为合法，这样转换成 Arith Op 的 AddOp、SubOp、ConstantOp 就有了判断是否合法的 Pattern。这里也可以用 addLegalOp 分别定义哪些 Op 为合法。对于 ReturnOp、CallOp 和 FuncOp，负责转换的 Pattern 只转换了操作数的类型，因此这里调用 addDynamicallyLegalOp，通过函数判断转换后的 Op 是否合法。

这时，如果我忘了定义判断 FuncOp 是否合法的 Pattern，FuncOp 不会被转换，但是后面的 Op 的操作数被转换了，例如

toy.func @add(%a: !toy.int<32>, %b: !toy.int<32>) -> !toy.int<32> {
  %c = toy.add %a, %b : !toy.int<32>
  toy.return %c : !toy.int<32>
}

toy.func @add(%a: !toy.int<32>, %b: !toy.int<32>) -> !toy.int<32> {
  %c = arith.addi %a, %b : i32
  toy.return %c : i32
}

转换结束后解释器发现 %a 和 %b 的类型前后不一致，就会报错

converter.addTargetMaterialization([](OpBuilder& builder, Type resultType, ValueRange inputs, Location loc)
        return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs).getResult(0);
        });

自动插入类型转换，保证解释通过。注意这里其实隐含了一个错误，就是 toy.return 返回了一个 i32 类型的值，但函数要求的返回值类型是 !toy.int<32>。之所以解释能够通过，应该是因为教程中没有给 toy::ReturnOp 添加专门的检查函数，如果把 toy.return 换成 func.return，解释器是能够发现并报错的。

除了 TargetMaterialization，还有另一种自动类型转换，用SourceMaterialization 添加，前者适用于某个指令的操作数是从其他 Op 或函数参数来的，但该指令里操作数的类型被改变了而其他 Op 或参数的类型没有改变的情况，比如上面的就是函数的参数类型没有改变，但 AddOp 的操作数类型改变了。后者则适用于指令里操作数类型没被改变但其他 Op 或参数的类型改变的情况，比如

toy.func @add(%a: i32, %b: i32) -> i32 {
  %c = toy.add %a, %b : !toy.int<32>
  toy.return %c : !toy.int<32>
}

注意这里在调试的时候有个坑点，就是没加 addTargetMaterialization 的时候，如果用 --debug 命令参数看解释器的转换过程，会发现转换的中间输出里面有自动添加的 builtin.unrealized_conversion_cast 命令，但结果还是会报错。并且，只有在需要的类型转换是 TargetMaterialization 的情况下才会出现这样的中间输出，感觉很奇怪，不知道这个中间输出和报错有什么关系。

总结

个人认为在“程序员三大浪漫”（图形学、操作系统、编译器）里，编译器门槛最高。这不是说编译器最难，而是说图形学和操作系统都是普通人能够接触到并由浅入深入门的。例如图形学，普通人可能会从打游戏、画画、做MMD到学习制作游戏、3D 建模、渲染，再到研究着色器、游戏引擎、建模软件，从而踏入图形学的大门；操作系统普通人也可以从玩电子电路、机器人到玩单片机、树莓派，再到自己修改、裁剪编译 Linux，从而踏入操作系统的大门。只有编译器，如果不是非常硬核的程序员，一般也不会想到去自己研究和实现，我虽然本科的时候上过编译原理的课程，但实际上对编译器也是知之甚少。

然而，随着当下大家开始卷 AI 加速器，各种 GPU、NPU、TPU 层出不穷，AI 编译器也获得了越来越多的关注。MLIR 是编写 AI 编译器的强大工具，目前也是处在不断发展的阶段，个人感觉其通过“方言”这个概念在不同阶段下进行组合和转换的思路挺有意思的，但 MLIR 及其所属的 LLVM 这个项目非常庞大，如果要深入学习也不是一件容易的事。

不过，MLIR-tutorial 这个教程确实是非常通俗易懂，让我对 MLIR 有了个基础的了解，所以我也将学习这个教程的心得分享给大家，希望能给其他 MLIR 学习者一点启发，也算是对教程作者的感谢。

概述