现在开启 MLIR 学习系列。本篇是跟着 Toy 语言学习 MLIR 的第二篇，主要介绍 toy dialect 高层级方言的匹配重写，前述内容请参考【MLIR】跟着 Toy 语言学习 MLIR【1】Toy 语言和 Toy Dialect。
学习 MLIR 最好的方式还是照着官方教程来，但是只看 MLIR 官网教程，Ch1 到 Ch6，没有基础的话又让人学着有点吃力。所以本教程没有直接参照官网 Toy 教程，而是以 MLIR 工程的 toy 相关代码入手。欢迎参考。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程，【编译器】使用 llvm 编译自定义语言【1】构建 AST ，【MLIR】跟着 Toy 语言学习 MLIR【1】Toy 语言和 Toy Dialect。
作为初学者，错误在所难免，还望不吝赐教。
基本简介
官网 Toy 教程 展示了如何将自定义语言 Toy 借助 MLIR 一步步编译为可执行机器码的过程。Toy 是一种简单的自定义语言，为了简便，其所有数据类型定义为 fp64 类型的 Tensor，支持 +/* 操作和 transpose 等有限的操作。
以下是整个编译降级的过程：
Toy txt -> Toy AST -> Toy Dialect -> Affine Dialect -> llvm Dialect -> llvm IR -> 机器码（通过 JIT 编译）
编译过程，下载 github llvm 工程，按照官网编译，值得一提的是，编译整个工具用时太久，我们重点关注 CH* 教程内容，在修改了教程中代码之后，可以只对该部分进行编译：
ninja -j 4 toyc-ch6 |
编译完成之后，可以执行一些指令进行测试：
/your/path/llvm-project/build/bin/toyc-ch6 /your/path/llvm-project/mlir/test/Examples/Toy/Ch5/affine-lowering.mlir -emit=mlir-llvm | |
/your/path/llvm-project/build/bin/toyc-ch6 /your/path/llvm-project/mlir/test/Examples/Toy/Ch5/affine-lowering.mlir -emit=jit | |
/your/path/llvm-project/build/bin/toyc-ch6 /your/path/llvm-project/mlir/test/Examples/Toy/Ch6/codegen.toy -emit=mlir | |
/your/path/llvm-project/build/bin/toyc-ch6 /your/path/llvm-project/mlir/test/Examples/Toy/Ch6/codegen.toy -emit=mlir -opt | |
/hyour/path/llvm-project/build/bin/toyc-ch6 /your/path/llvm-project/mlir/test/Examples/Toy/Ch6/codegen.toy -emit=mlir-affine | |
/your/path/llvm-project/build/bin/toyc-ch6 /your/path/llvm-project/mlir/test/Examples/Toy/Ch6/codegen.toy -emit=mlir-llvm |
前篇博客已经讲解了 toy 语言转换到自定义的 toy dialect 的过程。本篇博客则简单介绍一下后续过程。
主函数调用的 loadAndProcessMLIR 函数完成的，而这个函数里面调用一系列 Pass 实现 dialect 的优化和转换。代码中做了简单的注释。
static int loadAndProcessMLIR(mlir::MLIRContext &context, | |
mlir::OwningOpRef<mlir::ModuleOp> &module) { | |
if (int error = loadMLIR(context, module)) | |
return error; | |
mlir::PassManager pm(module.get()->getName()); | |
if (mlir::failed(mlir::applyPassManagerCLOptions(pm))) | |
return 4; | |
bool isLoweringToAffine = emitAction >= Action::DumpMLIRAffine; | |
bool isLoweringToLLVM = emitAction >= Action::DumpMLIRLLVM; | |
if (enableOpt || isLoweringToAffine) { | |
pm.addPass(mlir::createInlinerPass()); | |
mlir::OpPassManager &optPM = pm.nest<mlir::toy::FuncOp>(); | |
optPM.addPass(mlir::toy::createShapeInferencePass()); | |
optPM.addPass(mlir::createCanonicalizerPass()); | |
optPM.addPass(mlir::createCSEPass()); | |
} | |
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()); | |
} | |
} | |
if (isLoweringToLLVM) { | |
pm.addPass(mlir::toy::createLowerToLLVMPass()); | |
pm.addPass(mlir::LLVM::createDIScopeForLLVMFuncOpPass()); | |
} | |
if (mlir::failed(pm.run(*module))) | |
return 4; | |
return 0; | |
} |
本节内容来自于 教程 Chapter 3: High-level Language-Specific Analysis and Transformation，主要讲解在高层级方言 toy dialect ，如果做 pattern 匹配和重写，用的例子是去除连续的 transpose 操作（连续的两个转置 transpose 可以消除 transpose(transpose(X)) -> X ）。
def transpose_transpose(x) { | |
return transpose(transpose(x)); | |
} |
上述 toy 代码会得到对应的 toy dialect :
toy.func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> { | |
%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64> | |
%1 = toy.transpose(%0 : tensor<*xf64>) to tensor<*xf64> | |
toy.return %1 : tensor<*xf64> | |
} |
从高层级方言进行优化是比较方便的，一旦降级到低层级方言，再进行优化就变得非常困难，例如，当降级到如下层级的时候，再进行优化将非常困难：
#define N 100 | |
#define M 100 | |
void sink(void *); | |
void double_transpose(int A[N][M]) { | |
int B[M][N]; | |
for(int i = 0; i < N; ++i) { | |
for(int j = 0; j < M; ++j) { | |
B[j][i] = A[i][j]; | |
} | |
} | |
for(int i = 0; i < N; ++i) { | |
for(int j = 0; j < M; ++j) { | |
A[i][j] = B[j][i]; | |
} | |
} | |
sink(A); | |
} |
教程将一个自定义的 RewritePattern 注册到 MLIR 的 CanonicalizerPass 中，来实现双转置的匹配和重写。前面的代码中提到过，该 pass 用于 [规范化 IR：消除冗余操作、简化常量表达式、应用简单的重写规则]。
在 ToyCombine.cpp 源码中能够找到对应的自定义 RewritePattern 。
struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp> { | |
SimplifyRedundantTranspose(mlir::MLIRContext *context) | |
: OpRewritePattern<TransposeOp>(context, 1) {} | |
llvm::LogicalResult matchAndRewrite(TransposeOp op, | |
mlir::PatternRewriter &rewriter) const override { | |
mlir::Value transposeInput = op.getOperand(); | |
TransposeOp transposeInputOp = transposeInput.getDefiningOp<TransposeOp>(); | |
if (!transposeInputOp) | |
return failure(); | |
rewriter.replaceOp(op, {transposeInputOp.getOperand()}); | |
return success(); | |
} | |
}; |
matchAndRewrite 是该 RewritePattern 的核心方法，当 MLIR 框架会在遍历 toy dialect 时，对每个 TransposeOp 调用这个方法。核心思路很简单，即匹配到一个 TransposeOp 的时候，会检查前一个节点是否也是 TransposeOp，不是的话匹配失败，是的话怎直接去掉两个 Transpose 操作。
CanonicalizerPass 以贪婪的、迭代的方式应用由操作定义的转换。为了保证 CanonicalizerPass 能够使用到该 pattern ，在 Ops.td 文件定义 Transpose 操作的时候，能够找到 let hasCanonicalizer = 1; 字段，以及需要通过 Canonicalization framework 注册这个转换 pattern：
void TransposeOp::getCanonicalizationPatterns( | |
RewritePatternSet &results, MLIRContext *context) { | |
results.add<SimplifyRedundantTranspose>(context); | |
} |
以上是一种 C++ 的方式实现匹配重写，MLIR 还提供了声明的方式，其提供了一套 rprovides a table-based syntax for pattern-match and rewrite rules，从 td 文件中声明匹配和重写规则。
以下是 toy 教程中针对 ReshapeOp 操作定义的声明式重写规则，其位于 ToyCombine.td 文件中。该文件主要定义了三个针对 toy.reshape 操作的优化，同上述 C++ 风格的匹配重写函数一样，这些优化在 CanonicalizerPass 运行时会被自动调用。
#ifndef TOY_COMBINE | |
#define TOY_COMBINE | |
include "mlir/IR/PatternBase.td" | |
include "toy/Ops.td" | |
def ReshapeReshapeOptPattern : Pat<(ReshapeOp(ReshapeOp $arg)), | |
(ReshapeOp $arg)>; | |
def ReshapeConstant : | |
NativeCodeCall<"$0.reshape(::llvm::cast<ShapedType>($1.getType()))">; | |
def FoldConstantReshapeOptPattern : Pat< | |
(ReshapeOp:$res (ConstantOp $arg)), | |
(ConstantOp (ReshapeConstant $arg, $res))>; | |
def TypesAreIdentical : Constraint<CPred<"$0.getType() == $1.getType()">>; | |
def RedundantReshapeOptPattern : Pat< | |
(ReshapeOp:$res $arg), (replaceWithValue $arg), | |
[(TypesAreIdentical $res, $arg)]>; | |
#endif |
我们来分别看看这三个重写规则：
1. 消除连续的 Reshape
def ReshapeReshapeOptPattern : Pat<(ReshapeOp(ReshapeOp $arg)), | |
(ReshapeOp $arg)>; |
其实当有两个连续的 Reshape 存在的时候，内层的 reshape 就没有作用了，所以将源模式 (ReshapeOp(ReshapeOp $arg)) 替换为新模式 (ReshapeOp $arg) ，这样就完成了 连续 reshape 的替换。而且保留的 reshape 参数是外层的 reshape 参数，这是符合逻辑的。
2. 常量折叠 + reshape
def ReshapeConstant : | |
NativeCodeCall<"$0.reshape(::llvm::cast<ShapedType>($1.getType()))">; | |
def FoldConstantReshapeOptPattern : Pat< | |
(ReshapeOp:$res (ConstantOp $arg)), | |
(ConstantOp (ReshapeConstant $arg, $res))>; |
匹配的是 Reshape (Constant (x)) 结构，匹配之后变成单个 ConstantOp。
NativeCodeCall 调用一段手写 C++ 代码片段，将常量 $arg 重塑为 $res 所要求的形状。
举个例子，假设优化前的源代码如下：
%0 = toy.constant dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]> : tensor<6xf64> | |
%1 = toy.reshape(%0) : tensor<6xf64> -> tensor<2x3xf64> |
之后 CanonicalizerPass 会匹配 reshape + Constant 的结构，然后调用写好的重写规则，将其变成新的：
%0 = toy.constant dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]> : tensor<6xf64> | |
%1 = toy.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64> |
3. 消除多余的 reshape
当输入和 reshape 的输出，他们的数据类型相同，即维度 shape 相同的时候，去掉 reshape 算子。
def TypesAreIdentical : Constraint<CPred<"$0.getType() == $1.getType()">>; | |
def RedundantReshapeOptPattern : Pat< | |
(ReshapeOp:$res $arg), (replaceWithValue $arg), | |
[(TypesAreIdentical $res, $arg)]>; |
举了例子：
%0 = toy.reshape(%arg0) : tensor<2x3xf64> -> tensor<2x3xf64> | |
多 td 文件中声明匹配和重写规则的细节，可以参考 Table-driven Declarative Rewrite Rule (DRR)
上述就是两种匹配重写的方式：声明式和 C++ 手写式，相比较而言，声明式更简洁方便，代码量少，但是调试难度高，不够灵活，适用于简单的匹配场景，而 C++ 方式更灵活，灵活性高，能实现复杂逻辑，但是写起来更复杂。
I have seen things you people wouldn't believe. | |
Attack ships on fire off the shoulder of Orion. | |
I have watched C-beams glitter in the dark near the Tannhauser Gate. | |
All those ... moments will be lost in time, like tears ... in rain. | |
Time to die ... | |
《银翼杀手》 |
本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目 或随便一个项目下提出 issue，并指明哪一篇博客，看到一定及时回复！
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。