现在开启 MLIR 学习系列。本篇是跟着 Toy 语言学习 MLIR 的第四篇，主要介绍高层级 toy dialect 降级到低层方言的过程，前述内容请参考【MLIR】跟着 Toy 语言学习 MLIR【1】Toy 语言和 Toy Dialect，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换。
学习 MLIR 最好的方式还是照着官方教程来，但是只看 MLIR 官网教程，Ch1 到 Ch6，没有基础的话又让人学着有点吃力。现发现好多细节，不看教程认识不到，所以当前内容跟随官网教程。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程，【编译器】使用 llvm 编译自定义语言【1】构建 AST ，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换。
作为初学者，错误在所难免，还望不吝赐教。
官网 Toy 教程 展示了如何将自定义语言 Toy 借助 MLIR 一步步编译为可执行机器码的过程。Toy 是一种简单的自定义语言，为了简便，其所有数据类型定义为 fp64 类型的 Tensor，支持 +/* 操作和 transpose 等有限的操作。
以下是整个编译降级的过程：
Toy txt -> Toy AST -> Toy Dialect -> Affine Dialect -> llvm Dialect -> llvm IR -> 机器码（通过 JIT 编译）
前述章节中我们已经将 toy 语言转换成 toy dialect，甚至还添加了 Operation 匹配重写、inline 和 inferShape 接口等，对 toy dialect 进行优化，现在我们迫切地将 toy dialect 转换成实际的代码，最终转换成 llvm，让其投入使用。但本章不会直接将其转换为 llvm，而是将其进行逐步降级 (渐进式降级)，这样我们就可以重用 MLIR 提供的针对中间层级方言的优化技术。
Affine 调方言专为程序中计算密集的部分设计，其功能有限：例如，它不支持表示 Toy 中的 print 内置函数，也不应支持！因此我们可以将 toy dialect 密集计算部分降级到 affine dialect 并进行优化，而 print 内置函数保留原来的 toy dialect。
MLIR 拥有多种不同的方言，因此有必要建立一个统一的框架来实现它们之间的转换。这时，DialectConversion 框架就派上了用场。DialectConversion 框架提供了一个标准化的、可扩展的 “翻译工具箱”，用于将一个方言（Dialect）的操作降级（Lowering）或转换到另一个方言。
使用这个框架需要提供两个必要内容和一个可选内容。
Conversion Target 转换目标：ConversionTarget 定义了 “什么是合法的”，什么是非法的，然后必须要将非法操作替换为合法操作。
Rewrite Patterns 一组翻译规则 ：RewritePattern 定义了具体的转换逻辑。比如它告诉框架：“当你遇到一个 toy.add 操作时，请按照以下方式将其替换为一组 arith 和 affine 操作。”
Type Converter 跨方言的 “类型翻译” （可选）： 当转换涉及类型变化时（例如从 tensor<...> 到 memref<...>），就需要 TypeConverter 。
我们希望将计算密集型的 Toy 操作转换为 Affine、Arith、Func 和 MemRef 方言的组合操作，以便进一步优化。所以首先定义转换目标：
void ToyToAffineLoweringPass::runOnOperation() { | |
mlir::ConversionTarget target(getContext()); | |
target.addLegalDialect<affine::AffineDialect, arith::ArithDialect, | |
func::FuncDialect, memref::MemRefDialect>(); | |
target.addIllegalDialect<ToyDialect>(); | |
target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op) { | |
return llvm::none_of(op->getOperandTypes(), llvm::IsaPred<TensorType>); | |
}); | |
... | |
} |
我们通过 Pass 来完成方言转换操作， ToyToAffineLoweringPass 是一个 Pass，它重写了 runOnOperation 方法。在前述章节中我们提到过 Pass 需要重写的核心方法就是 runOnOperation 。
在这个方法中指明了哪些方言合法，哪些方言非法，并特别指定了 toy.print 这个操作合法，因为我们没法将其转换到 Affine。
前述章节方言内的优化（如 transpose (transpose (x)) -> x）的时候，我们使用过 RewritePattern ，它是方言内部的匹配和重写，现在我们需要跨方言的转换 ConversionPattern 。
和 RewritePattern 不同的是， ConversionPattern 接收一个额外的参数（OpAdaptor），该参数包含已重新映射或替换的操作数。
意思就是说，我们将操作（例如 toy.transpose）降级到 Affine 之后，需要将对应的操作数也从 TensorType 降级到 MemRefType 类型，因为 affine 方言的所有内存操作（如 affine.load、affine.store）都要求操作数是 memref 类型。而 OpAdaptor 就是用来提供操作数的类型映射。
struct TransposeOpLowering : public OpConversionPattern<toy::TransposeOp> { | |
using OpConversionPattern<toy::TransposeOp>::OpConversionPattern; | |
LogicalResult matchAndRewrite(toy::TransposeOp op, OpAdaptor adaptor, | |
ConversionPatternRewriter &rewriter) const final { | |
auto loc = op->getLoc(); | |
lowerOpToLoops(op, rewriter, | |
[&](OpBuilder &builder, ValueRange loopIvs) { | |
Value input = adaptor.getInput(); | |
SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs)); | |
return affine::AffineLoadOp::create(builder, loc, input, | |
reverseIvs); | |
}); | |
return success(); | |
} | |
}; |
这个模式匹配重写继承的是 OpConversionPattern ，而非之前 transpose(transpose(x)) -> x 时的 OpRewritePattern ，同样是重写 matchAndRewrite 函数，但它多了 适配器 OpAdaptor 。而且我们发现，这个 OpAdaptor 是作为参数传入的，并不需要我们手动实现，实际上， DialectConversion 框架在调用你的 matchAndRewrite 方法时，会自动创建一个 OpAdaptor 实例，并传递当前操作的实际操作数。
接下来就是将这些操作转换匹配全部放在刚才的 Pass 函数里面：
void ToyToAffineLoweringPass::runOnOperation() { | |
... | |
mlir::RewritePatternSet patterns(&getContext()); | |
patterns.add<..., TransposeOpLowering>(&getContext()); | |
... |
部分降级。 DialectConversion 框架提供了多种不同的降级模式，这里选择部分降级，因为还有部分（toy.print）不会降级：
void ToyToAffineLoweringPass::runOnOperation() { | |
... | |
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, patterns))) | |
signalPassFailure(); | |
} |
一个 toy dialect 降级的例子，降级之前：
toy.func @main() { | |
%0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64> | |
%2 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64> | |
%3 = toy.mul %2, %2 : tensor<3x2xf64> | |
toy.print %3 : tensor<3x2xf64> | |
toy.return | |
} |
降级之后：
func.func @main() { | |
%cst = arith.constant 1.000000e+00 : f64 | |
%cst_0 = arith.constant 2.000000e+00 : f64 | |
%cst_1 = arith.constant 3.000000e+00 : f64 | |
%cst_2 = arith.constant 4.000000e+00 : f64 | |
%cst_3 = arith.constant 5.000000e+00 : f64 | |
%cst_4 = arith.constant 6.000000e+00 : f64 | |
%0 = memref.alloc() : memref<3x2xf64> | |
%1 = memref.alloc() : memref<3x2xf64> | |
%2 = memref.alloc() : memref<2x3xf64> | |
affine.store %cst, %2[0, 0] : memref<2x3xf64> | |
affine.store %cst_0, %2[0, 1] : memref<2x3xf64> | |
affine.store %cst_1, %2[0, 2] : memref<2x3xf64> | |
affine.store %cst_2, %2[1, 0] : memref<2x3xf64> | |
affine.store %cst_3, %2[1, 1] : memref<2x3xf64> | |
affine.store %cst_4, %2[1, 2] : memref<2x3xf64> | |
affine.for %arg0 = 0 to 3 { | |
affine.for %arg1 = 0 to 2 { | |
%3 = affine.load %2[%arg1, %arg0] : memref<2x3xf64> | |
affine.store %3, %1[%arg0, %arg1] : memref<3x2xf64> | |
} | |
} | |
affine.for %arg0 = 0 to 3 { | |
affine.for %arg1 = 0 to 2 { | |
%3 = affine.load %1[%arg0, %arg1] : memref<3x2xf64> | |
%4 = affine.load %1[%arg0, %arg1] : memref<3x2xf64> | |
%5 = arith.mulf %3, %4 : f64 | |
affine.store %5, %0[%arg0, %arg1] : memref<3x2xf64> | |
} | |
} | |
toy.print %0 : memref<3x2xf64> | |
memref.dealloc %2 : memref<2x3xf64> | |
memref.dealloc %1 : memref<3x2xf64> | |
memref.dealloc %0 : memref<3x2xf64> | |
return | |
} |
当一个操作（如 toy.print）没有被完全降级时，如何让它同时支持高层的 tensor 和低层的 memref？
在降级过程中，大部分 toy 操作（如 add、transpose）被转换成了 affine 循环和 memref。但 toy.print 没有被降级，它仍然是一个 toy 方言的操作，此时它的输入出现了矛盾：toy.print 的定义只接受 tensor，但降级后的数据是 memref。
教程提出了三种方案：
1. 从 memref 生成 tensor（隐式复制）
%memref = ... : memref<2x3xf64> | |
%tensor = "toy.materialize"(%memref) : memref<2x3xf64> -> tensor<2x3xf64> | |
toy.print %tensor : tensor<2x3xf64> | |
2. 创建 toy.print_memref（新增操作）
%memref = ... : memref<2x3xf64> | |
toy.print_memref %memref : memref<2x3xf64> |
3. 让 toy.print 同时支持 tensor 和 memref
%memref = ... : memref<2x3xf64> | |
toy.print %memref : memref<2x3xf64> | |
教程选择了第三种方案，即定义 Print 操作的时候，让其支持两种类型的操作数。
def PrintOp : Toy_Op<"print"> { | |
... | |
let arguments = (ins AnyTypeOf<[F64Tensor, F64MemRef]>:$input); | |
} |
AnyTypeOf<[F64Tensor, F64MemRef]> 表示 PrintOp 的输入可以是 F64Tensor 或 F64MemRef 。
这样，toy.print 在降级前可以接受 tensor，在降级后可以接受 memref，不需要任何额外操作。
上述已经将 toy dialect 降级到 Affine dialect，它是正确的，但在效率方面仍有很大提升空间。例如，toy.mul 的降级产生了一些冗余的加载操作。而 MLIR 为已有的 Affine dialect 提供了好用的优化 Pass。我们只需要将 LoopFusion 和 AffineScalarReplacement 通过 passes 添加到流水线后，就能对当前的方言进行优化：
if (isLoweringToAffine) { | |
pm.addPass(mlir::toy::createLowerToAffinePass()); | |
mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>(); | |
optPM.addPass(mlir::createCanonicalizerPass()); | |
optPM.addPass(mlir::createCSEPass()); | |
if (enableOpt) { | |
optPM.addPass(mlir::affine::createLoopFusionPass()); | |
optPM.addPass(mlir::affine::createAffineScalarReplacementPass()); | |
} | |
} |
LoopFusion 是循环融合 pass，能够将相邻的、有相同循环边界的循环融合为一个循环，减少循环开销，提升数据局部性。
AffineScalarReplacement 检测临时分配的 memref 是否可以被标量（寄存器）替换，从而减少内存访问。
func.func @main() { | |
%cst = arith.constant 1.000000e+00 : f64 | |
%cst_0 = arith.constant 2.000000e+00 : f64 | |
%cst_1 = arith.constant 3.000000e+00 : f64 | |
%cst_2 = arith.constant 4.000000e+00 : f64 | |
%cst_3 = arith.constant 5.000000e+00 : f64 | |
%cst_4 = arith.constant 6.000000e+00 : f64 | |
%0 = memref.alloc() : memref<3x2xf64> | |
%1 = memref.alloc() : memref<2x3xf64> | |
affine.store %cst, %1[0, 0] : memref<2x3xf64> | |
affine.store %cst_0, %1[0, 1] : memref<2x3xf64> | |
affine.store %cst_1, %1[0, 2] : memref<2x3xf64> | |
affine.store %cst_2, %1[1, 0] : memref<2x3xf64> | |
affine.store %cst_3, %1[1, 1] : memref<2x3xf64> | |
affine.store %cst_4, %1[1, 2] : memref<2x3xf64> | |
affine.for %arg0 = 0 to 3 { | |
affine.for %arg1 = 0 to 2 { | |
%2 = affine.load %1[%arg1, %arg0] : memref<2x3xf64> | |
%3 = arith.mulf %2, %2 : f64 | |
affine.store %3, %0[%arg0, %arg1] : memref<3x2xf64> | |
} | |
} | |
toy.print %0 : memref<3x2xf64> | |
memref.dealloc %1 : memref<2x3xf64> | |
memref.dealloc %0 : memref<3x2xf64> | |
return | |
} |
可以通过以下指令，以及是否添加 -opt 来检查优化前后的区别。
toyc-ch5 test/Examples/Toy/Ch5/affine-lowering.mlir -emit=mlir-affine |
I know someone in the world is waiting for me, although I've no idea of who he is. But I feel happy every day for this. | |
- 《返老还童》 |
本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目 或随便一个项目下提出 issue，并指明哪一篇博客，看到一定及时回复！
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。