现在开启 MLIR 学习系列。本篇是跟着 Toy 语言学习 MLIR 的第五篇，前述章节已经介绍到 toy dialect 降级到 Affine dialect，本章节主要介绍将其降级到 LLVM 及 Codegen 生成可执行文件的过程。前述内容请参考【MLIR】跟着 Toy 语言学习 MLIR【1】Toy 语言和 Toy Dialect，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换，【MLIR】跟着 Toy 语言学习 MLIR【4】部分降级到低层方言。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程，【编译器】使用 llvm 编译自定义语言【1】构建 AST ，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换，【MLIR】跟着 Toy 语言学习 MLIR【4】部分降级到低层方言。
作为初学者，错误在所难免，还望不吝赐教。
官网 Toy 教程 展示了如何将自定义语言 Toy 借助 MLIR 一步步编译为可执行机器码的过程。Toy 是一种简单的自定义语言，为了简便，其所有数据类型定义为 fp64 类型的 Tensor，支持 +/* 操作和 transpose 等有限的操作。
以下是整个编译降级的过程：
Toy txt -> Toy AST -> Toy Dialect -> Affine Dialect -> llvm Dialect -> llvm IR -> 机器码（通过 JIT 编译）
在之前的降级中，我们已经降级的大部分的操作到 Affine dialect ，现在方言的状态是：
func.func @main() { | |
toy.print %memref : memref<2x3xf64> | |
return | |
} |
接下来是对 toy.print 进行降级，教程中说我们不需要直接将 toy.print 直接降级到 llvm，这样太复杂。因为 DialectConversion 框架支持传递性降级，不需要直接生成 LLVM 方言的操作。可以先生成一个结构化的循环嵌套（循环嵌套是为了不断打印输出信息），只要我们有从循环操作到 LLVM 的降级规则，整个降级就会成功。
教程给出构建 printf 的声明：
static FlatSymbolRefAttr getOrInsertPrintf(PatternRewriter &rewriter, | |
ModuleOp module, | |
LLVM::LLVMDialect *llvmDialect) { | |
auto *context = module.getContext(); | |
if (module.lookupSymbol<LLVM::LLVMFuncOp>("printf")) | |
return SymbolRefAttr::get(context, "printf"); | |
auto llvmI32Ty = IntegerType::get(context, 32); | |
auto llvmI8PtrTy = | |
LLVM::LLVMPointerType::get(IntegerType::get(context, 8)); | |
auto llvmFnType = LLVM::LLVMFunctionType::get(llvmI32Ty, llvmI8PtrTy, | |
true); | |
PatternRewriter::InsertionGuard insertGuard(rewriter); | |
rewriter.setInsertionPointToStart(module.getBody()); | |
LLVM::LLVMFuncOp::create(rewriter, module.getLoc(), "printf", llvmFnType); | |
return SymbolRefAttr::get(context, "printf"); | |
} |
这是一个 getOrInsertPrintf ，它只是一个普通的辅助函数，它的作用是 ：在模块中获取或创建 printf 函数的声明，并返回它的符号引用。
以下是 toy.print 操作降级 Pattern。
class PrintOpLowering : public OpConversionPattern<toy::PrintOp> { | |
public: | |
using OpConversionPattern<toy::PrintOp>::OpConversionPattern; | |
LogicalResult | |
matchAndRewrite(toy::PrintOp op, OpAdaptor adaptor, | |
ConversionPatternRewriter &rewriter) const override { | |
auto *context = rewriter.getContext(); | |
auto memRefType = llvm::cast<MemRefType>((*op->operand_type_begin())); | |
auto memRefShape = memRefType.getShape(); | |
auto loc = op->getLoc(); | |
ModuleOp parentModule = op->getParentOfType<ModuleOp>(); | |
auto printfRef = getOrInsertPrintf(rewriter, parentModule); | |
Value formatSpecifierCst = getOrCreateGlobalString( | |
loc, rewriter, "frmt_spec", StringRef("%f \0", 4), parentModule); | |
Value newLineCst = getOrCreateGlobalString( | |
loc, rewriter, "nl", StringRef("\n\0", 2), parentModule); | |
SmallVector<Value, 4> loopIvs; | |
for (unsigned i = 0, e = memRefShape.size(); i != e; ++i) { | |
auto lowerBound = arith::ConstantIndexOp::create(rewriter, loc, 0); | |
auto upperBound = | |
arith::ConstantIndexOp::create(rewriter, loc, memRefShape[i]); | |
auto step = arith::ConstantIndexOp::create(rewriter, loc, 1); | |
auto loop = | |
scf::ForOp::create(rewriter, loc, lowerBound, upperBound, step); | |
for (Operation &nested : make_early_inc_range(*loop.getBody())) | |
rewriter.eraseOp(&nested); | |
loopIvs.push_back(loop.getInductionVar()); | |
rewriter.setInsertionPointToEnd(loop.getBody()); | |
if (i != e - 1) | |
LLVM::CallOp::create(rewriter, loc, getPrintfType(context), printfRef, | |
newLineCst); | |
scf::YieldOp::create(rewriter, loc); | |
rewriter.setInsertionPointToStart(loop.getBody()); | |
} | |
auto elementLoad = | |
memref::LoadOp::create(rewriter, loc, op.getInput(), loopIvs); | |
LLVM::CallOp::create(rewriter, loc, getPrintfType(context), printfRef, | |
ArrayRef<Value>({formatSpecifierCst, elementLoad})); | |
rewriter.eraseOp(op); | |
return success(); | |
} | |
private: | |
static LLVM::LLVMFunctionType getPrintfType(MLIRContext *context) { | |
auto llvmI32Ty = IntegerType::get(context, 32); | |
auto llvmPtrTy = LLVM::LLVMPointerType::get(context); | |
auto llvmFnType = LLVM::LLVMFunctionType::get(llvmI32Ty, llvmPtrTy, | |
true); | |
return llvmFnType; | |
} | |
static FlatSymbolRefAttr getOrInsertPrintf(PatternRewriter &rewriter, | |
ModuleOp module) { | |
auto *context = module.getContext(); | |
if (module.lookupSymbol<LLVM::LLVMFuncOp>("printf")) | |
return SymbolRefAttr::get(context, "printf"); | |
PatternRewriter::InsertionGuard insertGuard(rewriter); | |
rewriter.setInsertionPointToStart(module.getBody()); | |
LLVM::LLVMFuncOp::create(rewriter, module.getLoc(), "printf", | |
getPrintfType(context)); | |
return SymbolRefAttr::get(context, "printf"); | |
} | |
static Value getOrCreateGlobalString(Location loc, OpBuilder &builder, | |
StringRef name, StringRef value, | |
ModuleOp module) { | |
LLVM::GlobalOp global; | |
if (!(global = module.lookupSymbol<LLVM::GlobalOp>(name))) { | |
OpBuilder::InsertionGuard insertGuard(builder); | |
builder.setInsertionPointToStart(module.getBody()); | |
auto type = LLVM::LLVMArrayType::get( | |
IntegerType::get(builder.getContext(), 8), value.size()); | |
global = LLVM::GlobalOp::create(builder, loc, type, true, | |
LLVM::Linkage::Internal, name, | |
builder.getStringAttr(value), | |
0); | |
} | |
Value globalPtr = LLVM::AddressOfOp::create(builder, loc, global); | |
Value cst0 = LLVM::ConstantOp::create(builder, loc, builder.getI64Type(), | |
builder.getIndexAttr(0)); | |
return LLVM::GEPOp::create( | |
builder, loc, LLVM::LLVMPointerType::get(builder.getContext()), | |
global.getType(), globalPtr, ArrayRef<Value>({cst0, cst0})); | |
} | |
}; | |
} |
以下 Pass 包装了将 当前 Affine dialect 和 toy.print 降级到 llvm dialect 的 pattern，其中 PrintOpLowering 是自定义的 print 降级 pattern。当前降级操作还要将把当前正在处理的 MemRef 类型转换为 LLVM 中的表示形式。为了完成此转换，我们使用 TypeConverter 作为降级过程的一部分。该转换器用于指定一种类型如何映射到另一种类型。
void ToyToLLVMLoweringPass::runOnOperation() { | |
LLVMConversionTarget target(getContext()); | |
target.addLegalOp<ModuleOp>(); | |
LLVMTypeConverter typeConverter(&getContext()); | |
RewritePatternSet patterns(&getContext()); | |
populateAffineToStdConversionPatterns(patterns); | |
populateSCFToControlFlowConversionPatterns(patterns); | |
mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns); | |
populateFinalizeMemRefToLLVMConversionPatterns(typeConverter, patterns); | |
cf::populateControlFlowToLLVMConversionPatterns(typeConverter, patterns); | |
populateFuncToLLVMConversionPatterns(typeConverter, patterns); | |
patterns.add<PrintOpLowering>(&getContext()); | |
auto module = getOperation(); | |
if (failed(applyFullConversion(module, target, std::move(patterns)))) | |
signalPassFailure(); | |
} | |
std::unique_ptr<mlir::Pass> mlir::toy::createLowerToLLVMPass() { | |
return std::make_unique<ToyToLLVMLoweringPass>(); | |
} |
应用这个 Pass：
if (isLoweringToLLVM) { | |
pm.addPass(mlir::toy::createLowerToLLVMPass()); | |
pm.addPass(mlir::LLVM::createDIScopeForLLVMFuncOpPass()); | |
} |
比较一下降级前后的差异性。降级之前的 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 | |
} |
降级到 llvm dialect :
llvm.func @free(!llvm<"i8*">) | |
llvm.func @printf(!llvm<"i8*">, ...) -> i32 | |
llvm.func @malloc(i64) -> !llvm<"i8*"> | |
llvm.func @main() { | |
%0 = llvm.mlir.constant(1.000000e+00 : f64) : f64 | |
%1 = llvm.mlir.constant(2.000000e+00 : f64) : f64 | |
... | |
^bb16: | |
%221 = llvm.extractvalue %25[0] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }"> | |
%222 = llvm.mlir.constant(0 : index) : i64 | |
%223 = llvm.mlir.constant(2 : index) : i64 | |
%224 = llvm.mul %214, %223 : i64 | |
%225 = llvm.add %222, %224 : i64 | |
%226 = llvm.mlir.constant(1 : index) : i64 | |
%227 = llvm.mul %219, %226 : i64 | |
%228 = llvm.add %225, %227 : i64 | |
%229 = llvm.getelementptr %221[%228] : (!llvm."double*">, i64) -> !llvm<"f64*"> | |
%230 = llvm.load %229 : !llvm<"double*"> | |
%231 = llvm.call @printf(%207, %230) : (!llvm<"i8*">, f64) -> i32 | |
%232 = llvm.add %219, %218 : i64 | |
llvm.br ^bb15(%232 : i64) | |
... | |
^bb18: | |
%235 = llvm.extractvalue %65[0] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }"> | |
%236 = llvm.bitcast %235 : !llvm<"double*"> to !llvm<"i8*"> | |
llvm.call @free(%236) : (!llvm<"i8*">) -> () | |
%237 = llvm.extractvalue %45[0] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }"> | |
%238 = llvm.bitcast %237 : !llvm<"double*"> to !llvm<"i8*"> | |
llvm.call @free(%238) : (!llvm<"i8*">) -> () | |
%239 = llvm.extractvalue %25[0] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }"> | |
%240 = llvm.bitcast %239 : !llvm<"double*"> to !llvm<"i8*"> | |
llvm.call @free(%240) : (!llvm<"i8*">) -> () | |
llvm.return | |
} |
此时我们已处于代码生成的临界点。我们可以使用 LLVM 语法生成代码，现在只需将其导出为 LLVM IR，并设置一个即时编译器（JIT）来运行即可。那么将 llvm dialect 转为 llvm IR :
int dumpLLVMIR(mlir::ModuleOp module) { | |
llvm::LLVMContext llvmContext; | |
auto llvmModule = mlir::translateModuleToLLVMIR(module, llvmContext); | |
if (!llvmModule) { | |
llvm::errs() << "Failed to emit LLVM IR\n"; | |
return -1; | |
} | |
llvm::InitializeNativeTarget(); | |
llvm::InitializeNativeTargetAsmPrinter(); | |
mlir::ExecutionEngine::setupTargetTriple(llvmModule.get()); | |
auto optPipeline = mlir::makeOptimizingTransformer( | |
EnableOpt ? 3 : 0, 0, | |
nullptr); | |
if (auto err = optPipeline(llvmModule.get())) { | |
llvm::errs() << "Failed to optimize LLVM IR " << err << "\n"; | |
return -1; | |
} | |
llvm::errs() << *llvmModule << "\n"; | |
return 0; | |
} |
转到 llvm ir 之后，打印部分内容：
define void @main() { | |
... | |
102: | |
%103 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %8, 0 | |
%104 = mul i64 %96, 2 | |
%105 = add i64 0, %104 | |
%106 = mul i64 %100, 1 | |
%107 = add i64 %105, %106 | |
%108 = getelementptr double, double* %103, i64 %107 | |
%109 = memref.load double, double* %108 | |
%110 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double %109) | |
%111 = add i64 %100, 1 | |
cf.br label %99 | |
... | |
115: | |
%116 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %24, 0 | |
%117 = bitcast double* %116 to i8* | |
call void @free(i8* %117) | |
%118 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %16, 0 | |
%119 = bitcast double* %118 to i8* | |
call void @free(i8* %119) | |
%120 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %8, 0 | |
%121 = bitcast double* %120 to i8* | |
call void @free(i8* %121) | |
ret void | |
} |
还可以对 llvm ir 启用优化，能够降低体积：
define void @main() | |
%0 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00) | |
%1 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.600000e+01) | |
%putchar = tail call i32 @putchar(i32 10) | |
%2 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 4.000000e+00) | |
%3 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 2.500000e+01) | |
%putchar.1 = tail call i32 @putchar(i32 10) | |
%4 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 9.000000e+00) | |
%5 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 3.600000e+01) | |
%putchar.2 = tail call i32 @putchar(i32 10) | |
ret void | |
} |
int runJit(mlir::ModuleOp module) { | |
llvm::InitializeNativeTarget(); | |
llvm::InitializeNativeTargetAsmPrinter(); | |
auto optPipeline = mlir::makeOptimizingTransformer( | |
EnableOpt ? 3 : 0, 0, | |
nullptr); | |
auto maybeEngine = mlir::ExecutionEngine::create(module, | |
nullptr, optPipeline); | |
assert(maybeEngine && "failed to construct an execution engine"); | |
auto &engine = maybeEngine.get(); | |
auto invocationResult = engine->invoke("main"); | |
if (invocationResult) { | |
llvm::errs() << "JIT invocation failed\n"; | |
return -1; | |
} | |
return 0; | |
} |
执行能够得到一下结果：
$ echo 'def main() { print([[1, 2], [3, 4]]); }' | ./bin/toyc-ch6 -emit=jit | |
1.000000 2.000000 | |
3.000000 4.000000 |
While you're spending all this time on your own, building computers or practicing your cello, what you're really doing, is becoming interesting . | |
And when people finally do notice you, they'er gonna find someone a lot cooler than they thought. |
本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目 或随便一个项目下提出 issue，并指明哪一篇博客，看到一定及时回复！
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。