从LLVM方言到LLVM IR，再到CodeGen
从LLVM方言到LLVM IR，再到CodeGen
在中使用 toy 语言接入 MLIR，最终转化为 LLVM IR，具体的流程如下：
.toy 源文件 —> AST —> MLIRGen(遍历AST生成MLIR表达式) —> Transformation(变形消除冗余) —> 下译 —> LLVM IR/JIT编译引擎
0. 编译相关工具链
1）克隆仓库
$ git clone <https://github.com/llvm/llvm-project.git>
2）配置
$ mkdir llvm-project/build
$ cd llvm-project/build
$ cmake -G Ninja ../llvm \\
-DLLVM_ENABLE_PROJECTS=mlir \\
-DLLVM_BUILD_EXAMPLES=ON \\
-DLLVM_TARGETS_TO_BUILD="X86;NVPTX;AMDGPU" \\
-DCMAKE_BUILD_TYPE=Release \\
-DLLVM_ENABLE_ASSERTIONS=ON \\
3）构建
$ cmake --build . --target check-mlir
构建结束后，工具链在llvm-project/build/bin路径下。
进入工具链所在文件夹cd llvm-project/build/bin
为了方便索引函数，可以在cmake配置中加上 -DCMAKE_EXPORT_COMPILE_COMMANDS=ON
生成的compile_commands.json在llvm-project/build目录下，复制到llvm-project目录即可。
配置vscode的clangd插件
ctrl + p 输入 clangd，先点击下载language server；然后加 settings.json, ctrl + p → '> 打开工作区设置json`粘贴以下。
{
"clangd.arguments": [
"--header-insertion=never",
"--compile-commands-dir=${workspaceFolder}/",
"--query-driver=**",
]
}
1. 输入代码转为ast
输入ast.toy
def multiply_transpose(a, b){
return transpose(a) * transpose(b);
}
def main() {
var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
var b<2, 3> = [1, 2, 3, 4, 5, 6];
var c = multiply_transpose(a, b);
print(c);
}
运行 ./toyc-ch1 ../../mlir/test/Examples/Toy/Ch1/ast.toy --emit=ast
Module:
Function
Proto 'multiply_transpose' @test/Examples/Toy/Ch1/ast.toy:4:1'
Params: [a, b]
Block {
Return
BinOp: * @test/Examples/Toy/Ch1/ast.toy:5:25
Call 'transpose' [ @test/Examples/Toy/Ch1/ast.toy:5:10
var: a @test/Examples/Toy/Ch1/ast.toy:5:20
]
Call 'transpose' [ @test/Examples/Toy/Ch1/ast.toy:5:25
var: b @test/Examples/Toy/Ch1/ast.toy:5:35
]
} // Block
... // main函数的ast未写出
2. 生产MLIR表达式
MLIRGen 模块会遍历 AST，递归调用子函数，构建operation，一个 dialect 中可以有很多的 operation。
运行./toyc-ch2 ../../mlir/test/Examples/Toy/Ch2/ast.toy --emit=mlir
module {
toy.func @multiply_transpose(%arg0: tensor<*xf64>, %arg1: tensor<*xf64>) -> tensor<*xf64> {
%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64>
%1 = toy.transpose(%arg1 : tensor<*xf64>) to tensor<*xf64>
%2 = toy.mul %0, %1 : tensor<*xf64>
toy.return %2 : tensor<*xf64>
}
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>
%1 = toy.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>
%2 = toy.reshape(%1 : tensor<6xf64>) to tensor<2x3xf64>
%3 = toy.generic_call @multiply_transpose(%0, %2) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
%4 = toy.generic_call @multiply_transpose(%2, %0) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
%5 = toy.generic_call @multiply_transpose(%2, %3) : (tensor<2x3xf64>, tensor<*xf64>) -> tensor<*xf64>
%6 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<*xf64>
%7 = toy.generic_call @multiply_transpose(%6, %3) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64>
toy.return
}
}
3.优化MLIR表达式
发现生成的 MLIR 表达式往往存在冗余的操作，如下输入
def transpose_transpose(x) {
return transpose(transpose(x));
}
def main() {
var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
var b = transpose_transpose(a);
print(b);
}
转化为MLIR表达式如下。可以看到，其实transpose_transpose函数输出恒等于输出，但该函数转化成的MLIR表达式还是调用了两次 toy.transpose
module {
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>
}
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>
%1 = toy.reshape(%0 : tensor<2x3xf64>) to tensor<2x3xf64>
%2 = toy.generic_call @transpose_transpose(%1) : (tensor<2x3xf64>) -> tensor<*xf64>
toy.print %2 : tensor<*xf64>
toy.return
}
}
为了提升程序性能就需要对表达式进行转换变形。MLIR 提供以下两种方式进行模式匹配转换：
1）使用 C++手动编写代码进行表达式的匹配与重写；
2）使用基于规则的模式匹配与重写的声明式重写规则（DRR）进行，但该方法要求使用ODS定义操作。
3.1 手动匹配重写
1）第一步：直接使用 C++ 写出匹配与重写的代码
下面这段代码位于在 ToyCombine.cpp 中，默认位置在 llvm-project/mlir/examples/toy/Ch3/mlir/ToyCombine.cpp
struct struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp> {
// 匹配该IR中所有的toy.transpose，根据benefit的值来决定优先级
SimplifyRedundantTranspose(mlir::MLIRContext *context)
: OpRewritePattern<TransposeOp>(context, /*benefit=*/1) {}
// 尝试匹配并重写
mlir::LogicalResult matchAndRewrite(TransposeOp op, mlir::PatternRewriter &rewriter)const override{
// 获取当前transpose的操作数
mlir::Value transposeInput = op.getOperand();
// 获取该操作数对应的op
TransposeOp transposeInputOp = transposeInput.getDefiningOp<TransposeOp>();
// 如果没有对应的op，返回失败
if(!transposeInputOp) return failure();
// 反之替换
rewriter.replaceOp(op, {TransposeOp.getOperand()});
return success();
}
}
2）第二步：将自定义的匹配与重写模式登记为规范化模式，使得后续可以使用它
下面这段代码位于 toyc.cpp 中，默认位置为 llvm-project/mlir/examples/toy/Ch3/mlir/ToyCombine.cpp
// 将自定义的匹配与重写模式登记为规范化模式
void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context){
results.add<SimplifyRedundantTranspose>(context);
}
3）第三步：在Ops.td中设置相应选项
下面这段代码位于Ops.td中，默认位置为llvm-project/mlir/examples/toy/Ch3/include/toy/Ops.td
def TransposeOp : Toy_Op<"transpose", [Pure]> {
// MLIR 在优化代码时较为保守，可能会保留一些无效操作
// 设置[Pure] 可解决这一问题
...
// 确保启用规范化框架，应用 canonicalization pass
let hasCanonicalizer = 1;
...
4）第四步：更新主文件以添加 optimization pipeline
下面这段代码位于 toyc.cpp 中，默认位置在 ../mlir/examples/toy/Ch3/toyc.cpp
if (enableOpt) {// enableOpt 是从命令行输入的编译选项
// 使用 PassManger 模块添加优化一道优化工序
if (mlir::failed(mlir::applyPassManagerCLOptions(pm)))
return 4;
// createCanonicalizerPass 创建并使用规范化框架
pm.addNestedPass<mlir::toy::FuncOp>(mlir::createCanonicalizerPass());
// 运行定义好的 canonicalizer 来优化 MLIR 表达式
if (mlir::failed(pm.run(*module)))
return 4;
}
使用-opt开启优化，对比transpose_transpose函数转化为的MLIR表达式变化
./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/transpose_transpose.toy --emit=mlir -opt
toy.func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {
toy.return %arg0 : tensor<*xf64>
}
对比 ./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/transpose_transpose.toy --emit=mlir
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>
}
3.2.采用 DDR 自动生成匹配与重写函数
DRR(Declarative, rule-based pattern-match and rewrite)是一种基于 DAG(Directed acyclic graph) 的声明性重写器，提供 table-base 的模式匹配与规则重写的句法。类似于 ODS 框架，只需要使用一定的声明性描述，就可以自动生成匹配与规则重写程序。
生成的 MLIR 表达式存在许多冗余的 reshape 操作，以消除冗余的 reshape 操作为例，说明第二种模式匹配转换方法。
输入
def main() {
var a<2,1> = [1, 2];
var b<2,1> = a;
var c<2,1> = b;
print(c);
}
转换为没经过优化的MLIR表达式
module {
toy.func @main() {
%0 = toy.constant dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>
%1 = toy.reshape(%0 : tensor<2xf64>) to tensor<2x1xf64>
%2 = toy.reshape(%1 : tensor<2x1xf64>) to tensor<2x1xf64>
%3 = toy.reshape(%2 : tensor<2x1xf64>) to tensor<2x1xf64>
toy.print %3 : tensor<2x1xf64>
toy.return
}
}
三种重写格式
下面这段代码位于ToyCombine.td中，默认位置为llvm-project/mlir/examples/toy/Ch3/mlir/ToyCombine.td
1）基础方法
// Reshape(Reshape(x)) = Reshape(x)
def ReshapeReshapeOptPattern : Pat<(ReshapeOp(ReshapeOp $arg)), (ReshapeOp $arg)>;
2）使用 NativeCodeCall （该方法可以通过调用 C++ helper function 或使用 inline C++ 进行更复杂的转换。）
// Reshape(Constant(x)) = x'
def ReshapeConstant : NativeCodeCall<"$0.reshape(::llvm::cast<ShapedType>($1.getType()))">;
def FoldConstantReshapeOptPattern : Pat<(ReshapeOp:$res (ConstantOp $arg)), (ConstantOp (ReshapeConstant $arg, $res))>;
3）添加参数约束的方法
DDR 提供了一种添加参数约束的方法，以应对当改写只发生在某些特定条件下的情况。
// 当输入形状与输出形状相同时，才消除该 reshape 操作
def TypesAreIdentical : Constraint<CPred<"$0.getType() == $1.getType()">>;
def RedundantReshapeOptPattern : Pat<(ReshapeOp:$res $arg), (replaceWithValue $arg), [(TypesAreIdentical $res, $arg)]>;
再使用下面的语句自动生成c++代码（虽然在本地运行会报错。）
${build_root}/bin/mlir-tblgen --gen-rewriters ${mlir_src_root}/examples/toy/Ch3/mlir/ToyCombine.td -I ${mlir_src_root}/include/
# 此时在llvm-project/build/bin/文件夹下
./mlir-tblgen --gen-rewriters ../../mlir/Examples/Toy/Ch3/mlir/ToyCombine.td -I ../../mlir/include/
最后执行 ./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/trivial_reshape.toy --emit=mlir -opt，得到优化后的MLIR表达式如下
module {
toy.func @main() {
%0 = toy.constant dense<[[1.000000e+00], [2.000000e+00]]> : tensor<2x1xf64>
toy.print %0 : tensor<2x1xf64>
toy.return
}
}
4.通用的转换接口
通过使用 Dialect，MLIR 可以表示多种不同等级的抽象。尽管这些不同的 Dialect 表示不同的抽象，但某些操作的算法机制十分相似，为了减少代码重复，MLIR 提供了一组通用的转换与分析。（让每个dialect创建时都可以复用一些操作）
1)为了代码执行速度更快，将函数进行内联(inline)操作；
2)为了代码生成阶段更方便，需要进行形状推断，确定所有 tensor 的 shape。
输入（/mlir/test/Examples/Toy/Ch3/codegen.toy）
def multiply_transpose(a, b) {
return transpose(a) * transpose(b);
}
def main() {
var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
var b<2, 3> = [1, 2, 3, 4, 5, 6];
var c = multiply_transpose(a, b);
var d = multiply_transpose(b, a);
print(d);
}
先输出经过reshape优化后的MLIR表达式（ ./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/codegen.toy --emit=mlir -opt）
module {
toy.func @multiply_transpose(%arg0: tensor<*xf64>, %arg1: tensor<*xf64>) -> tensor<*xf64> {
%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64>
%1 = toy.transpose(%arg1 : tensor<*xf64>) to tensor<*xf64>
%2 = toy.mul %0, %1 : tensor<*xf64>
toy.return %2 : tensor<*xf64>
}
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>
%1 = 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.generic_call @multiply_transpose(%0, %1) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
%3 = toy.generic_call @multiply_transpose(%1, %0) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
toy.print %3 : tensor<*xf64>
toy.return
}
}
4.1 函数内联
内联（inline）会将函数展开，把函数的代码复制到每一个调用处，以解决一些频繁调用的函数大量消耗栈空间（栈内存）的问题。使用该优化方法后，编译器会将简单函数内嵌到调用处，以储存空间为代价换取运行速度。
1）首先，需要一个专属于 Toy 的内联函数接口，MLIR 中已经提供了相应的内联函数接口模板DialectInlinerInterface，只需要在Toy方言中继承这一类来编写 Toy 中的内联函数接口即可。
下面这段代码位于 Dialect.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/Dialect.cpp
struct ToyInlinerInterface : public DialectInlinerInterface {
using DialectInlinerInterface::DialectInlinerInterface;
//===--------------------------------------------------------------------===//
// 分析挂钩
//===--------------------------------------------------------------------===//
/// toy dialect中所有callop、op、functions操作能够inline
bool isLegalToInline(Operation *call, Operation *callable, bool wouldBeCloned) const final {
return true;
}
bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
return true;
}
bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
return true;
}
//===--------------------------------------------------------------------===//
// 转换挂钩
//===--------------------------------------------------------------------===//
void handleTerminator(Operation *op, ArrayRef<Value> valuesToRepl) const final {
// 此处只对 toy.return 进行处理
auto returnOp = cast<ReturnOp>(op);
assert(returnOp.getNumOperands() == valuesToRepl.size());
for (const auto &it : llvm::enumerate(returnOp.getOperands()))
valuesToRepl[it.index()].replaceAllUsesWith(it.value());
}
};
2）该内联操作会忽略未被引用的函数操作，此外，还需要规定其处理的函数范围（除了主函数）
下面这段代码位于 MLIRGen.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/MLIRGen.cpp
mlir::toy::FuncOp mlirGen(FunctionAST &funcAST) {
...
// 如果此函数不是main函数，则将可见性设置为private。
if (funcAST.getProto()->getName() != "main")
function.setPrivate();
return function;
}
3）其次，需要在Toy方言中注册内联接口
下面这段代码位于 Dialect.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/Dialect.cpp
void ToyDialect::initialize() {
addInterfaces<ToyInlinerInterface>();
}
4）然后，需要定位函数调用的位置。
由于内联操作都是对调用函数进行操作，所以需要让内联器（inliner）知道 Toy Dialect IR 中toy.generic_call代表调用。这里需要实现将 CallOpInterface 添加到 GenericCallOp。
下面这段代码位于 Ops.td 中，默认位置为llvm-project/mlir/examples/toy/Ch4/include/toy/Ops.td
// 使用 CallOpInterface 可以将操作标记为调用
include "mlir/Interfaces/CallInterfaces.td"
...
// 将其加入到 GenericCallOp 的 traits 列表中
def FuncOp : Toy_Op<"func",
[DeclareOpInterfaceMethods<CallableOpInterface>]> {
...
}
def GenericCallOp : Toy_Op<"generic_call",
[DeclareOpInterfaceMethods<CallOpInterface>]> {
...
}
上面的程序中，使用DeclareOpInterfaceMethods指令，在GenericCallOp类中声明接口的使用方法。还需要提供 GenericCallOp 的定义。
下面这段代码位于 Dialect.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/Dialect.cpp
// 返回function operation中可调用区域
Region *FuncOp::getCallableRegion() { return &getBody(); }
// 返回结果类型
ArrayRef<Type> FuncOp::getCallableResults() { return getType().getResults(); }
// 返回所有参数的属性，如果没有则返回 null。
ArrayAttr FuncOp::getCallableArgAttrs() {
return getArgAttrs().value_or(nullptr);
}
// 返回所有结果的属性，如果没有则返回 null。
ArrayAttr FuncOp::getCallableResAttrs() {
return getResAttrs().value_or(nullptr);
}
// 返回被调用者
CallInterfaceCallable GenericCallOp::getCallableForCallee() {
return getAttrOfType<SymbolRefAttr>("callee");
}
// 设置generic call op的被调用者
void GenericCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
(*this)->setAttr("callee", callee.get<SymbolRefAttr>());
}
// 获得被调用函数的操作数
Operation::operand_range GenericCallOp::getArgOperands() { return getInputs(); }
5）在调用时实参与形参的类型可能不同，所以需要添加一个显式的类型转换(explicit cast)，因此需要在Toy方言中添加 cast 操作并设置调用接口。
下面这段代码位于 Ops.td 中，默认位置为llvm-project/mlir/examples/toy/Ch4/include/toy/Ops.td
def CastOp : Toy_Op<"cast", [ DeclareOpInterfaceMethods<CastOpInterface>,
NoMemoryEffect, SameOperandsAndResultShape]
> {
let summary = "shape cast operation";
let description = [{
强制转换操作将张量从一种类型转换为等效类型，而不更改任何数据元素。源类型和目标类型必须都是具有相同元素类型的张量类型。如果两者都进行了排名，则需要形状匹配。如果转换为不匹配的常量维度，则该操作无效。
}];
let arguments = (ins F64Tensor:$input);
let results = (outs F64Tensor:$output);
let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
}
上述代码将 CastOpInterface 加入了 traits 列表中，还需要使用areCastCompatible来定义进入此接口的方法(hook into this interface)
下面这段代码位于 Dialect.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/Dialect.cpp
// 该程序限定了能够进行 explicit cast 的条件
bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
if (inputs.size() != 1 || outputs.size() != 1)
return false;
TensorType input = llvm::dyn_cast<TensorType>(inputs.front());
TensorType output = llvm::dyn_cast<TensorType>(outputs.front());
if (!input || !output || input.getElementType() != output.getElementType())
return false;
return !input.hasRank() || !output.hasRank() || input == output;
}
然后在（1）中定义好的ToyInlinerInterface中增加显式强制转换的内容，以保证内联操作顺利执行。
下面这段代码位于 Dialect.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/Dialect.cpp
// 定义 Toy 内联函数接口
struct ToyInlinerInterface : public DialectInlinerInterface {
...
// 是否在调用中启用 explicit cast
Operation *materializeCallConversion(OpBuilder &builder, Value input,
Type resultType,
Location conversionLoc) const final {
return builder.create<CastOp>(conversionLoc, resultType, input);
}
};
6）最后将内联优化添加到优化管道中
下面这段代码位于 toyc.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/toyc.cpp
if (enableOpt) {
mlir::PassManager pm(module->getName());
if (mlir::failed(mlir::applyPassManagerCLOptions(pm)))
return 4;
...
// 将内联优化应用于所有function，然后会删去其他的function，只剩下一个// main
pm.addPass(mlir::createInlinerPass());
...
}
运行后可得到经过内联(inline) Pass后的Toy方言，只剩下了一个function(main函数)。
module {
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>
%1 = 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.cast %1 : tensor<2x3xf64> to tensor<*xf64>
%3 = toy.cast %0 : tensor<2x3xf64> to tensor<*xf64>
%4 = toy.transpose(%2 : tensor<*xf64>) to tensor<*xf64>
%5 = toy.transpose(%3 : tensor<*xf64>) to tensor<*xf64>
%6 = toy.mul %4, %5 : tensor<*xf64>
toy.print %6 : tensor<*xf64>
toy.return
}
}
4.2.形状推理
目前主函数中存在动态与静态形状的混合，提前确定所有张量的形状能够使最终生成的代码更加简洁。
1）首先，使用 ODS 来定义 ShapeInference 操作的接口
下面这段代码位于ShapeInferenceInterface.td中，默认位置为llvm-project/mlir/examples/toy/Ch4/include/toy/ShapeInferenceInterface.td
def ShapeInferenceOpInterface : OpInterface<"ShapeInference"> {
...
let methods = [
InterfaceMethod<"Infer and set the output shape for the current operation.",
"void", "inferShapes">
];
}
2）其次，在 Toy方言中 ShapeInferenceOp 添加到需要它的 operation（就像实现内联 Pass中的第三步：把CallOpInterface添加到GenericCallOp）。
下面这段代码位于Ops.td中，默认位置为llvm-project/mlir/examples/toy/Ch4/include/toy/Ops.td
// 下面是将形状推断操作(ShapeInferenceOp)添加到乘法操作(MulOp)中
// 也可以添加到其他操作中
def MulOp : Toy_Op<"mul",
[..., DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
...
}
3）然后，一些 operation 就获得了形状推理操作(ShapeInferenceOp)的接口，就需要在这些 operation 中定义对应的形状推断函数，独立定义可以保证 ShapeInferencePass 会独立地作用于该 operation。
下面这段代码位于 Dialect.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/Dialect.cpp
void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
// 所有在上一步中添加了ShapeInferenceOpInterface的op，这一步都需要 定义对应的形状推断函数
4）最后，将形状推断优化添加到优化管道中
下面这段代码位于toyc.cpp中，默认位置为llvm-project/mlir/examples/toy/Ch4/toyc.cpp
if (enableOpt) {
mlir::PassManager pm(module->getName());
if (mlir::failed(mlir::applyPassManagerCLOptions(pm)))
return 4;
// 将内联优化应用于所有function，然后会删去其他的function，只剩下一个main
pm.addPass(mlir::createInlinerPass());
// 现在只剩下一个function(main)，可以推断出operations的shape
mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
// 形状推断优化
optPM.addPass(mlir::toy::createShapeInferencePass());
// 规范化框架优化(直接调用就行)
optPM.addPass(mlir::createCanonicalizerPass());
// 公共子表达式消除(直接调用就行)
optPM.addPass(mlir::createCSEPass());
if (mlir::failed(pm.run(*module)))
return 4;
}
经过内联(inline) Pass与形状(shape)推断 Pass，得到优化后的main函数 MLIR 表达式如下。
运行 ./toyc-ch4 ../../mlir/test/Examples/Toy/Ch4/codegen.toy --emit=mlir -opt
module {
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>
%1 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
%2 = toy.mul %1, %1 : tensor<3x2xf64>
toy.print %2 : tensor<3x2xf64>
toy.return
}
}
运行 ./toyc-ch4 ../../mlir/test/Examples/Toy/Ch4/codegen.toy --emit=mlir -opt --mlir-print-ir-after-all 观察每一个pass后的MLIR表达式改变
可选：将优化运用在所有操作operation中
上面编写好的 ShapeInferencePass 会针对每一个函数进行操作，独立地优化每一个 函数。如果想将优化操作泛化到全局（run on any isolated operation），则可以使用 MLIR 的OperationPass接口。
下面这段代码位于 ShapeInferencePass.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch4/mlir/ShapeInferencePass.cpp
// 需要实现全局的pass都要继承并重写mlir::OperationPass<FuncOp>runOnOperation()
struct ShapeInferencePass
: public mlir::PassWrapper<ShapeInferencePass, OperationPass<toy::FuncOp>> {
void runOnOperation() override {
auto f = getOperation();
...
// 算法流程:
// 1.将所有需要进行形状推断的operation加入一个worklist
// 2.遍历这个worklist，对于每个operation，从参数类型推断其输出的形状
// 3.直到worklist为空
}
};
// 通过函数实例化 ShapeInferencePass
std::unique_ptr<mlir::Pass> mlir::toy::createShapeInferencePass() {
return std::make_unique<ShapeInferencePass>();
}
5.从 MLIR 表达式进行部分下译
5.1 背景知识（下译与 DialectConversion）
在编译器一系列转换程序的过程中，越来越多的高层次的简明信息被打散，转换为低层次的细碎指令，这个过程被称为代码表示递降下译，与之相反的过程被称为代码表示递升raising。raising远比下译困难，因为需要在庞杂的细节中找出宏观脉络。
1） 下译过程中越晚执行的转换越有结构劣势，因为缺乏高层次信息。
2）下译主要是为了更贴近硬件做代码生成与做硬件相关的优化。
每次转换遍历(pass) 都需要保持原子性，在其内部可能会临时违反源程序语义，但在每个转换遍历之后，中间表示应该是正确的。编译器依赖每个遍历之后的中间表示验证 (validation) 来保证正确性。
MLIR 中有许多不同的方言，下译过程其实就是在各种方言之间转化，而 MLIR 提供了一套统一的DialectConversion框架来实现不同方言之间的转化。
要使用 DialectConversion 框架需要 Three Components：
1）Conversion Target（转换目标）
对转换目标 Dialect 进行合法化(legal)，对当前的 Dialect 进行非法化(illegal)。主要完成以下三件事：
A）合法方言（目标方言）
target.addLegalDialect<affine::AffineDialect, arith::ArithDialect>();
将 AffineDialect 与 ArithDialect 添加为合法的目标
B）非法方言（如果未转换则失败）
target.addIllegalDIalect<toy::ToyDialect>();
由于 Toy Dialect 已经转换走了，就将其添加为非法的目标
C）合法和非法操作
target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op);
// 将保留操作添加为合法操作
2）转换模式（或者称为重写模式）
上一步相当于转换了命名空间，但并没有对其中的 operation 进行转换，需要对 operation 进行匹配与重写，将非法操作转换为合法操作。（实现operation从源dialect到目标dialect的映射）
3） Type Conversion（类型转换器）
当前 dialect 中若存在某些特定的数据类型，则需要转换到目标 dialect 中相应的数据类型。
5.1 部分下译
当前dialect中若存在某些特定的数据类型，则需要转换到目标dialect中相应的数据类型。
从一个高抽象级别的 Dialect 到一个低抽象级别的 Dialect 过程中，可以只下译其中一部分 operation，剩下的 operation 只需要升级与其他 operation 共存。现在以对 转换后的 MLIR 表达式进行下译为例：
变换起点是上一步优化好的MLIR表达式，首先由codegen.toy获得优化后的MLIR表达式（下面代码中下侧部分代码）。
// codegen.toy 源码
def multiply_transpose(a, b) {
return transpose(a) * transpose(b);
}
def main() {
var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
var b<2, 3> = [1, 2, 3, 4, 5, 6];
var c = multiply_transpose(a, b);
var d = multiply_transpose(b, a);
print(d);
}
// transformation(内联与形状推断) 后的 MLIR 表达式
// ./toyc-ch4 ../../mlir/test/Examples/Toy/Ch4/codegen.toy --emit=mlir -opt
// 或 ./toyc-ch5 ../../mlir/test/Examples/Toy/Ch5/codegen.toy --emit=mlir -opt
// 与 Ch5/affine-lowering.mlir 相同
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>
%1 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
%2 = toy.mul %1, %1 : tensor<3x2xf64>
toy.print %2 : tensor<3x2xf64>
toy.return
}
1）第一步：定义转换目标（Conversion Target）
为了实现进一步优化，将Toy方言中计算密集操作转换为仿射方言与算术方言（这两个都是MLIR内置的 Dialect）的组合，但由于仿射方言中没有输出操作，就需要将 Toy Dialect 中的输出操作保留并重写。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
void ToyToAffineLoweringPass::runOnOperation() {
ConversionTarget target(getContext());
// 将`Affine`, `Arith`, `Func`, and `MemRef` Dialect添加为合法的目标
target.addLegalDialect<affine::AffineDialect, BuiltinDialect,
arith::ArithDialect, func::FuncDialect,
memref::MemRefDialect>();
// 将 `Toy` Dialect添加为非法的目标
target.addIllegalDialect<toy::ToyDialect>();
// 保留ToyDialect中的print operation，后续重写
target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op) {
return llvm::none_of(op->getOperandTypes(), [](Type type) { return llvm::isa<TensorType>(type); });
});
...
}
2）第二步：明确转换模式（Conversion Patterns）
这一步将使用 ConversionPattern实现对 operation 的匹配与重写，把 非法操作转换为合法操作。
下面以转换 ToyDialect 中转换操作为例。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
struct TransposeOpLowering : public ConversionPattern {
TransposeOpLowering(MLIRContext *ctx)
: ConversionPattern(toy::TransposeOp::getOperationName(), 1, ctx) {}
// 匹配与重写函数
LogicalResult
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const final {
auto loc = op->getLoc();
// 实现将当前的操作lower到一组仿射循环
// memRef是AffineDialect的操作数类型，类似于缓存区
lowerOpToLoops(op, operands, rewriter,
[loc](OpBuilder &builder, ValueRange memRefOperands,
ValueRange loopIvs) {
// TransposeOpAdaptor 是在ODS框架执行后自动生成的
toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
Value input = transposeAdaptor.getInput();
SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
return builder.create<affine::AffineLoadOp>(loc, input,
reverseIvs);
});
return success();
}
};
3）第三步：将第二步定义的转换模式(TransposeOpLowering)添加到lower过程中用到的模式列表。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
// 添加一组可以降低toy op的模式
void ToyToAffineLoweringPass::runOnOperation() {
...
RewritePatternSet patterns(&getContext());
patterns.add<..., TransposeOpLowering>(&getContext());
}
4）第四步：确定部分下译模式
DialectConversion 框架提供了两种模式的下译，部分方法与全局方法：
1） 部分: 当前 Dialect 中某些 operation 在 lowering 中先进行保留（保留部分之前的信息）
2）全局: 当前 Dialect 中全部 operation 在 lowering 中全部去除（类似转换到 LLVM IR）。
由于需要将 Toy Dialect 中的 print operation 保留并重写，所以这里使用 Partial Method 执行。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
void ToyToAffineLoweringPass::runOnOperation() {
...
if (failed(applyPartialConversion(getOperation(), target, std::move(patterns))))
signalPassFailure();
}
5）第五步：将保留的 toy.print 进行重写，以匹配数据格式。在这一步需要将保留的 toy.print 的输出格式，并且增加新数据格式支持。
有三种实现方法：
A）从bufferload生成操作
保持操作的定义不变，但涉及完整的复制。
B）生成一个在下译的typetoy.print上操作的新版本
对优化器没有隐藏的、不必要的复制，但需要另一个操作定义。
C）更新以允许在下译的typetoy.print上操作
需要在dialect中混合抽象级别
为了简单使用的是第三种实现方法
下面这段代码位于 Ops.td 中，默认位置为llvm-project/mlir/examples/toy/Ch5/include/toy/Ops.td
def PrintOp : Toy_Op<"print"> {
...
// 之前是 let arguments = (ins F64Tensor:$input);
// 添加对 F64MemRef 类型的输出支持
let arguments = (ins AnyTypeOf<[F64Tensor, F64MemRef]>:$input);
}
6）第六步：将定义好的下译添加到优化管道中
下面这段代码位于 toyc.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch5/toyc.cpp
// 使用 PassManger 模块添加优化工序
if (isLoweringToAffine) { // 若命令行中有-emit=mlir-affine，则为真
// 对toy dialect进行 部分lower
pm.addPass(mlir::toy::createLowerToAffinePass());
// LowerToAffine优化，规范化框架优化，公共子表达式消除优化
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());
}
}
执行 ./toyc-ch5 ../../mlir/test/Examples/Toy/Ch5/codegen.toy --emit=mlir-affine -opt，得到下译后的结果。
codegen.toy --emit=mlir -opt的结果就是affine-lowering.mlir 故运行./toyc-ch5 ../../test/Examples/Toy/Ch5/affine-lowering.mlir -emit=mlir-affine -opt得到的结果等效。
module {
func.func @main() {
%cst = arith.constant 6.000000e+00 : f64
%cst_0 = arith.constant 5.000000e+00 : f64
%cst_1 = arith.constant 4.000000e+00 : f64
%cst_2 = arith.constant 3.000000e+00 : f64
%cst_3 = arith.constant 2.000000e+00 : f64
%cst_4 = arith.constant 1.000000e+00 : f64
%alloc = memref.alloc() : memref<3x2xf64>
%alloc_5 = memref.alloc() : memref<2x3xf64>
affine.store %cst_4, %alloc_5[0, 0] : memref<2x3xf64>
affine.store %cst_3, %alloc_5[0, 1] : memref<2x3xf64>
affine.store %cst_2, %alloc_5[0, 2] : memref<2x3xf64>
affine.store %cst_1, %alloc_5[1, 0] : memref<2x3xf64>
affine.store %cst_0, %alloc_5[1, 1] : memref<2x3xf64>
affine.store %cst, %alloc_5[1, 2] : memref<2x3xf64>
affine.for %arg0 = 0 to 3 {
affine.for %arg1 = 0 to 2 {
%0 = affine.load %alloc_5[%arg1, %arg0] : memref<2x3xf64>
%1 = arith.mulf %0, %0 : f64
affine.store %1, %alloc[%arg0, %arg1] : memref<3x2xf64>
}
}
toy.print %alloc : memref<3x2xf64>
memref.dealloc %alloc_5 : memref<2x3xf64>
memref.dealloc %alloc : memref<3x2xf64>
return
}
}
6.混合方言表达式下译到 LLVM IR
已经将Toy方言转换为仿射方言、算术方言以及包含Toy方言中的输出操作的混合操作，需要全部下译到LLVM方言，再下译到 LLVM IR 接入到 LLVM 后端进行CodeGen。（LLVM方言属于MLIR的方言，LLVM IR是LLVM自己的IR）
Affine --
↓
Arithmetic + Func --> LLVM (Dialect)
↑
'toy.print' --> Loop (SCF) --
1）第一步：下译toy.print
已经下译了除toy.print操作之外的所有操作，将此操作下译到一个为每个元素调用printf的非仿射循环嵌套。
方言转换框架支持传递下译(transitive lowering)，不需要直接生成在 LLVM 方言。通过传递降低，可以应用多种模式来使操作完全合法化。
下面这段代码位于 LowerToLLVM.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp
static FlatSymbolRefAttr getOrInsertPrintf(PatternRewriter &rewriter, ModuleOp module) {
auto *context = module.getContext();
if (module.lookupSymbol<LLVM::LLVMFuncOp>("printf"))
return SymbolRefAttr::get(context, "printf");
// 为printf创建函数声明 signature * `i32 (i8*, ...)`
auto llvmI32Ty = IntegerType::get(context, 32);
auto llvmI8PtrTy = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
auto llvmFnType = LLVM::LLVMFunctionType::get(llvmI32Ty, llvmI8PtrTy, /*isVarArg=*/true);
// 将printf函数插入父模块的主体中
PatternRewriter::InsertionGuard insertGuard(rewriter);
rewriter.setInsertionPointToStart(module.getBody());
rewriter.create<LLVM::LLVMFuncOp>(module.getLoc(), "printf", llvmFnType);
return SymbolRefAttr::get(context, "printf");
}
2）第二步：转换目标
除了顶层模块（ the top-level module）之外，需要将所有内容都下降到LLVM方言。
下面这段代码位于 LowerToLLVM.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp
void ToyToLLVMLoweringPass::runOnOperation() {
// The first thing to define is the conversion target. This will define the
// final target for this lowering. For this lowering, we are only targeting
// the LLVM dialect.
LLVMConversionTarget target(getContext());
target.addLegalOp<ModuleOp>();
...
}
3）第三步：类型转换
接下来的下译过程还需将当前所使用的 MemRef 类型转换为 LLVM 中的表示形式，MLIR 中已经定义好很多 typeConverter 用于复用。
下面这段代码位于 LowerToLLVM.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp
LLVMTypeConverter typeConverter(&getContext());
4）第四步：转换模式
下面这段代码位于 LowerToLLVM.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp
// affine，arith、 与std方言已经提供了将它们转换为 LLVM 方言所需的模式集
RewritePatternSet patterns(&getContext());
populateAffineToStdConversionPatterns(patterns);
populateSCFToControlFlowConversionPatterns(patterns);
mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);
populateFinalizeMemRefToLLVMConversionPatterns(typeConverter, patterns);
cf::populateControlFlowToLLVMConversionPatterns(typeConverter, patterns);
populateFuncToLLVMConversionPatterns(typeConverter, patterns);
// Toy方言中仅存的toy.print需要独立编写PrintOpLowering
// 类似于上一节第二步中的TransposeOpLowering
patterns.add<PrintOpLowering>(&getContext());
5）第五步：确定全局下译模式
下面这段代码位于 LowerToLLVM.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp
void ToyToLLVMLoweringPass::runOnFunction() {
...
auto module = getOperation();
if (failed(applyFullConversion(module, target, std::move(patterns))))
signalPassFailure();
}
6）第六步：将定义好的下译添加到优化管道中
下面这段代码位于 toyc.cpp 中，默认位置为llvm-project/mlir/examples/toy/Ch6/toyc.cpp
if (isLoweringToLLVM) {
// 完成从toy ir到llvm dialect的下降
pm.addPass(mlir::toy::createLowerToLLVMPass());
// 添加pass
pm.addNestedPass<mlir::LLVM::LLVMFuncOp>(
mlir::LLVM::createDIScopeForLLVMFuncOpPass());
}
执行./toyc-ch6 ../../test/Examples/Toy/Ch6/llvm-lowering.mlir -emit=mlir-llvm，最终会获得的 LLVM方言格式的MLIR表达式如下
// LLVM IR Dialect 形式的 MLIR表达式
module{
llvm.func @free(!11vm<"i8*">)
llvm.mlir.global internal constant @nl("\\0A\\00")
llvm.mlir.global internal constant @frmt_spec("%f\\00")
llvm.func @printf(!llvm<"i8*">, ...) -> !llvm.i32
llvm.func @malloc(!llvm.i64) -> !llvm<"i8*"> llvm.func @main(){
%0=llvm.mlir.constant(1.000000e+00 : f64) : !llvm.double
%1=llvm.mlir.constant(2.000000e+00 : f64) : !llvm.double
%2=llvm.mlir.constant(3.000000e+00 : f64) : !llvm.double
...
}
7）第七步：从LLVM方言到LLVM IR，再到CodeGen
现在已经转换到 LLVM方言，最终需要下译到 LLVM IR，使用 MLIR 内置的转换函数translateModuleToLLVMIR即可。然后利用 LLVM工具链即可完成多后端 CodeGen。
人工智能芯片与自动驾驶
