Toy方言源文件下译、MLIRGen、相关编译操作流程技术
//第5章/mlir_ast_toy_transpose.c
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<2×3×f64>
%1 = toy.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]>: tensor<6×f64>
%2 = toy.reshape(%1: tensor<6×f64>) to tensor<2×3×f64>
%3 = toy.generic_call @multiply_transpose(%0, %2): (tensor<2×3×f64>, tensor<2×3×f64>) -> tensor<*xf64>
%4 = toy.generic_call @multiply_transpose(%2, %0): (tensor<2×3×f64>, tensor<2×3×f64>) -> tensor<*xf64>
%5 = toy.generic_call @multiply_transpose(%2, %3): (tensor<2×3×f64>, tensor<*xf64>) -> tensor<*xf64>
%6 = toy.transpose(%0: tensor<2×3×f64>) to tensor<*xf64>
%7 = toy.generic_call @multiply_transpose(%6, %3): (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64>
toy.return
}
}
//第5章/mlir_func_multiply_transpose.c
toy.func @multiply_transpose(%arg0: tensor<*xf64>, %arg1:
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);
}
//第5章/mlir_func_toy.transpose.c
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<2×3×f64>
%1 = toy.reshape(%0: tensor<2×3×f64>) to tensor<2×3×f64>
%2 = toy.generic_call @transpose_transpose(%1): (tensor<2×3×f64>) -> tensor<*xf64>
toy.print %2: tensor<*xf64>
toy.return
}
}
//第5章/mlir_manual_toy_match.c
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>
}
//第5章/mlir_reshape_ddr_match.c
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<2×f64>
%1 = toy.reshape(%0: tensor<2×f64>) to tensor<2×1×f64>
%2 = toy.reshape(%1: tensor<2×1×f64>) to tensor<2×1×f64>
%3 = toy.reshape(%2: tensor<2×1×f64>) to tensor<2×1×f64>
toy.print %3: tensor<2×1×f64>
toy.return
}
}
//第5章/mlir_Toy_combine_format.c
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<2×1×f64>
toy.print %0: tensor<2×1×f64>
toy.return
}
}
//第5章/mlir_inline_shape.c
input（/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<2×3×f64>
%1 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]>: tensor<2×3×f64>
%2 = toy.generic_call @multiply_transpose(%0, %1): (tensor<2×3×f64>, tensor<2×3×f64>) -> tensor<*xf64>
%3 = toy.generic_call @multiply_transpose(%1, %0): (tensor<2×3×f64>, tensor<2×3×f64>) -> tensor<*xf64>
toy.print %3: tensor<*xf64>
toy.return
}
}
//第5章/mlir_inline_toy.c
1）首先，需要一个专属于 Toy 的内联函数接口，MLIR 中已经提供了相应的内联函数接口模板DialectInlinerInterface，只需要在Toy方言中继承这一类来编写 Toy 中的内联函数接口即可。
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
// 返回函数操作中可调用区域
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,
LocationconversionLoc) 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<2×3×f64>
%1 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]>: tensor<2×3×f64>
%2 = toy.cast %1: tensor<2×3×f64> to tensor<*xf64>
%3 = toy.cast %0: tensor<2×3×f64> 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
}
}
//第5章/mlir_ShapeInference.c
1）首先，使用 ODS 来定义 ShapeInference 操作的接口
下面这段代码位于ShapeInferenceInterface.td中，默认位置为
llvm-project/MLIR/examples/toy/Ch4/include/toy/ShapeInferenceInterface.td
def ShapeInferenceOpInterface: OpInterface<"ShapeInference"> {
...
let methods = [
InterfaceMethod<"推断并设置当前操作的输出形状。",
"void", "inferShapes">
];
}
2）其次，在 Toy方言中 ShapeInferenceOp 添加到需要它的 操作（就像实现内联 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）然后，一些操作就获得了形状推理操作(ShapeInferenceOp)的接口，就需要在这些操作中定义对应的形状推断函数，独立定义可以保证 ShapeInferencePass 会独立地作用于该操作。
下面这段代码位于 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<2×3×f64>
%1 = toy.transpose(%0: tensor<2×3×f64>) to tensor<3×2×f64>
%2 = toy.mul %1, %1: tensor<3×2×f64>
toy.print %2: tensor<3×2×f64>
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_AffineDialect.c
A）合法方言（目标方言）
target.addLegalDialect<affine::AffineDialect, arith::ArithDialect>();
将 AffineDialect 与 ArithDialect 添加为合法的目标
B）非法方言（如果未转换则失败）
target.addIllegalDIalect<toy::ToyDialect>();
由于 Toy方言已经转换走了，就将其添加为非法的目标
C）合法和非法操作
target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op);
// 将保留操作添加为合法操作
//第5章/mlir_codegen_toy.c
// 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-下译.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<2×3×f64>
%1 = toy.transpose(%0: tensor<2×3×f64>) to tensor<3×2×f64>
%2 = toy.mul %1, %1: tensor<3×2×f64>
toy.print %2: tensor<3×2×f64>
toy.return
}
1）第一步：定义转换目标（conversion Target）
为了实现进一步优化，将Toy方言中计算密集操作转换为仿射方言与算术方言（这两个都是MLIR内置的 Dialect）的组合，但由于仿射方言中没有输出操作，就需要将 Toy方言中的输出操作保留并重写。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/MLIR/examples/toy/Ch5/MLIR/LowerToAffineLoops.cpp
void ToyToAffine下译Pass::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操作，后续重写
target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op) {
return llvm::none_of(op->getOperandTypes(), [](Type type) { return llvm::isa<tensorType>(type); });
});
...
}
2）第二步：明确转换模式（conversion Patterns）
这一步将使用conversionPattern实现对操作的匹配与重写，把 非法操作转换为合法操作。
下面以转换 ToyDialect 中转换操作为例。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/MLIR/examples/toy/Ch5/MLIR/LowerToAffineLoops.cpp
struct TransposeOp下译: publicconversionPattern {
TransposeOp下译(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）第三步：将第二步定义的转换模式(TransposeOp下译)添加到lower过程中用到的模式列表。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/MLIR/examples/toy/Ch5/MLIR/LowerToAffineLoops.cpp
// 添加一组可以降低toy op的模式
void ToyToAffine下译Pass::runOnOperation() {
...
RewritePatternSet patterns(&getContext());
patterns.add<..., TransposeOp下译>(&getContext());
}
4）第四步：确定部分下译模式
方言转换 框架提供了两种模式的下译，部分方法与全局方法：
1） 部分: 当前方言中某些操作在下译中先进行保留（保留部分之前的信息）
2）全局: 当前方言中全部操作在 下译 中全部去除（类似转换到 LLVM IR）。
由于需要将 Toy方言中的打印操作保留并重写，所以这里使用 Partial Method 执行。
下面这段代码位于 LowerToAffineLoops.cpp 中，默认位置为llvm-project/MLIR/examples/toy/Ch5/MLIR/LowerToAffineLoops.cpp
void ToyToAffine下译Pass::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 (is下译ToAffine) { // 若命令行中有-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-下译.MLIR
故运行./toyc-ch5../../test/Examples/Toy/Ch5/affine-下译.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<3×2×f64>
%alloc_5 = memref.alloc(): memref<2×3×f64>
affine.store %cst_4, %alloc_5[0, 0]: memref<2×3×f64>
affine.store %cst_3, %alloc_5[0, 1]: memref<2×3×f64>
affine.store %cst_2, %alloc_5[0, 2]: memref<2×3×f64>
affine.store %cst_1, %alloc_5[1, 0]: memref<2×3×f64>
affine.store %cst_0, %alloc_5[1, 1]: memref<2×3×f64>
affine.store %cst, %alloc_5[1, 2]: memref<2×3×f64>
affine.for %arg0 = 0 to 3 {
affine.for %arg1 = 0 to 2 {
%0 = affine.load %alloc_5[%arg1, %arg0]: memref<2×3×f64>
%1 = arith.mulf %0, %0: f64
affine.store %1, %alloc[%arg0, %arg1]: memref<3×2×f64>
}
}
toy.print %alloc: memref<3×2×f64>
memref.dealloc %alloc_5: memref<2×3×f64>
memref.dealloc %alloc: memref<3×2×f64>
return
}
}
//第5章/mlir_codegen_toy.c
Affine --
↓
Arithmetic + Func --> LLVM (Dialect)
↑
'toy.print' --> Loop (SCF) --
1）第一步：下译toy.print
已经下译了除toy.print操作之外的所有操作，将此操作下译到一个为每个元素调用printf的非仿射循环嵌套。
方言转换框架支持传递下译(transitive 下译)，不需要直接生成在 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 ToyToLLVM下译Pass::runOnOperation() {
// The first thing to define is theconversion target. This will define the
// final target for this 下译. For this 下译, 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需要独立编写PrintOp下译
// 类似于上一节第二步中的TransposeOp下译
patterns.add<PrintOp下译>(&getContext());
5）第五步：确定全局下译模式
下面这段代码位于 LowerToLLVM.cpp 中，默认位置为llvm-project/MLIR/examples/toy/Ch6/MLIR/LowerToLLVM.cpp
void ToyToLLVM下译Pass::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 (is下译ToLLVM) {
// 完成从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-下译.MLIR -emit=MLIR-llvm，最终会获得的 LLVM方言格式的MLIR表达式如下
// LLVM IR方言形式的 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。
