2. 构建 JIT:添加优化 – ORC 层简介¶
本教程正在积极开发中。它尚不完整,细节可能会频繁更改。 尽管如此,我们邀请您尝试目前的状态,并欢迎任何反馈。
2.1. 第 2 章 简介¶
警告:本教程目前正在更新以适应 ORC API 的更改。只有第 1 章和第 2 章是最新的。
第 3 章到第 5 章的示例代码可以编译和运行,但尚未更新
欢迎来到“在 LLVM 中构建基于 ORC 的 JIT”教程的第 2 章。在本系列的第 1 章中,我们研究了一个基本的 JIT 类 KaleidoscopeJIT,它可以将 LLVM IR 模块作为输入,并在内存中生成可执行代码。KaleidoscopeJIT 能够以相对较少的代码做到这一点,方法是组合两个现成的 ORC 层:IRCompileLayer 和 ObjectLinkingLayer,以完成大部分繁重的工作。
在这一层中,我们将通过使用一个新的层 IRTransformLayer,为 KaleidoscopeJIT 添加 IR 优化支持,从而更多地了解 ORC 层概念。
2.2. 使用 IRTransformLayer 优化模块¶
在“使用 LLVM 实现语言”教程系列的第 4 章中,介绍了 llvm FunctionPassManager 作为优化 LLVM IR 的一种手段。感兴趣的读者可以阅读该章节了解详细信息,但简而言之:要优化一个模块,我们创建一个 llvm::FunctionPassManager 实例,使用一组优化配置它,然后在模块上运行 PassManager,将其变异为(希望)更优化但语义等价的形式。在原始教程系列中,FunctionPassManager 是在 KaleidoscopeJIT 外部创建的,模块在添加到它之前进行了优化。在本章中,我们将使优化成为 JIT 的一个阶段。目前,这将为我们提供学习更多关于 ORC 层的动力,但从长远来看,将优化作为 JIT 的一部分将产生一个重要的好处:当我们开始延迟编译代码(即,将每个函数的编译推迟到第一次运行时)时,由我们的 JIT 管理优化将允许我们也能延迟优化,而不是必须预先完成所有的优化。
为了为我们的 JIT 添加优化支持,我们将采用第 1 章中的 KaleidoscopeJIT,并在其之上组合一个 ORC IRTransformLayer。我们将在下面更详细地了解 IRTransformLayer 的工作原理,但接口很简单:此层的构造函数接受对执行会话和下层(所有层都这样做)的引用,以及一个 IR 优化函数,它将应用于通过 addModule 添加的每个模块
class KaleidoscopeJIT {
private:
ExecutionSession ES;
RTDyldObjectLinkingLayer ObjectLayer;
IRCompileLayer CompileLayer;
IRTransformLayer TransformLayer;
DataLayout DL;
MangleAndInterner Mangle;
ThreadSafeContext Ctx;
public:
KaleidoscopeJIT(JITTargetMachineBuilder JTMB, DataLayout DL)
: ObjectLayer(ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
CompileLayer(ES, ObjectLayer, ConcurrentIRCompiler(std::move(JTMB))),
TransformLayer(ES, CompileLayer, optimizeModule),
DL(std::move(DL)), Mangle(ES, this->DL),
Ctx(std::make_unique<LLVMContext>()) {
ES.getMainJITDylib().addGenerator(
cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(DL.getGlobalPrefix())));
}
我们扩展的 KaleidoscopeJIT 类开始时与第 1 章相同,但在 CompileLayer 之后,我们引入了一个新成员 TransformLayer,它位于我们的 CompileLayer 之上。我们使用对 ExecutionSession 和输出层(层的标准做法)的引用以及一个 转换函数 初始化我们的 OptimizeLayer。对于我们的转换函数,我们提供我们类的 optimizeModule 静态方法。
// ...
return cantFail(OptimizeLayer.addModule(std::move(M),
std::move(Resolver)));
// ...
接下来,我们需要更新我们的 addModule 方法,以将对 CompileLayer::add
的调用替换为对 OptimizeLayer::add
的调用。
static Expected<ThreadSafeModule>
optimizeModule(ThreadSafeModule M, const MaterializationResponsibility &R) {
// Create a function pass manager.
auto FPM = std::make_unique<legacy::FunctionPassManager>(M.get());
// Add some optimizations.
FPM->add(createInstructionCombiningPass());
FPM->add(createReassociatePass());
FPM->add(createGVNPass());
FPM->add(createCFGSimplificationPass());
FPM->doInitialization();
// Run the optimizations over all functions in the module being added to
// the JIT.
for (auto &F : *M)
FPM->run(F);
return M;
}
在我们的 JIT 的底部,我们添加了一个私有方法来执行实际的优化:optimizeModule。此函数将要转换的模块作为输入(作为 ThreadSafeModule),以及对新类 MaterializationResponsibility
的引用的引用。MaterializationResponsibility 参数可用于查询正在转换的模块的 JIT 状态,例如模块中 JIT 代码正在积极尝试调用/访问的定义集。目前,我们将忽略此参数并使用标准优化管道。为此,我们设置一个 FunctionPassManager,向其中添加一些 pass,在模块中的每个函数上运行它,然后返回变异的模块。具体的优化与“使用 LLVM 实现语言”教程系列的第 4 章中使用的相同。读者可以访问该章节以更深入地讨论这些以及一般的 IR 优化。
这就是对 KaleidoscopeJIT 的更改:当通过 addModule 添加模块时,OptimizeLayer 将在将转换后的模块传递到下面的 CompileLayer 之前调用我们的 optimizeModule 函数。当然,我们可以直接在我们的 addModule 函数中调用 optimizeModule,而不用费心使用 IRTransformLayer,但这样做为我们提供了另一个机会来了解层是如何组合的。它也为 层 概念本身提供了一个简洁的入口点,因为 IRTransformLayer 是可以实现的最简单的层之一。
// From IRTransformLayer.h:
class IRTransformLayer : public IRLayer {
public:
using TransformFunction = std::function<Expected<ThreadSafeModule>(
ThreadSafeModule, const MaterializationResponsibility &R)>;
IRTransformLayer(ExecutionSession &ES, IRLayer &BaseLayer,
TransformFunction Transform = identityTransform);
void setTransform(TransformFunction Transform) {
this->Transform = std::move(Transform);
}
static ThreadSafeModule
identityTransform(ThreadSafeModule TSM,
const MaterializationResponsibility &R) {
return TSM;
}
void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;
private:
IRLayer &BaseLayer;
TransformFunction Transform;
};
// From IRTransformLayer.cpp:
IRTransformLayer::IRTransformLayer(ExecutionSession &ES,
IRLayer &BaseLayer,
TransformFunction Transform)
: IRLayer(ES), BaseLayer(BaseLayer), Transform(std::move(Transform)) {}
void IRTransformLayer::emit(MaterializationResponsibility R,
ThreadSafeModule TSM) {
assert(TSM.getModule() && "Module must not be null");
if (auto TransformedTSM = Transform(std::move(TSM), R))
BaseLayer.emit(std::move(R), std::move(*TransformedTSM));
else {
R.failMaterialization();
getExecutionSession().reportError(TransformedTSM.takeError());
}
}
这是 IRTransformLayer 的完整定义,来自 llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h
和 llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
。这个类关注两个非常简单的工作:(1)通过转换函数对象运行通过此层发出的每个 IR 模块,以及(2)实现 ORC IRLayer
接口(它本身符合一般的 ORC 层概念,稍后会详细介绍)。这个类的大部分内容都很简单:转换函数的 typedef,用于初始化成员的构造函数,转换函数值的 setter,以及默认的 no-op 转换。最重要的方法是 emit
,因为这是我们 IRLayer 接口的一半。emit 方法将我们的转换应用于调用它的每个模块,如果转换成功,则将转换后的模块传递到基础层。如果转换失败,我们的 emit 函数会调用 MaterializationResponsibility::failMaterialization
(这让可能在其他线程上等待的 JIT 客户端知道他们等待的代码编译失败),并在退出之前将错误记录到执行会话中。
IRLayer 接口的另一半是我们从 IRLayer 类继承的,未经修改
Error IRLayer::add(JITDylib &JD, ThreadSafeModule TSM, VModuleKey K) {
return JD.define(std::make_unique<BasicIRLayerMaterializationUnit>(
*this, std::move(K), std::move(TSM)));
}
这段代码来自 llvm/lib/ExecutionEngine/Orc/Layer.cpp
,通过将其包装在 MaterializationUnit
(在本例中为 BasicIRLayerMaterializationUnit
)中,将 ThreadSafeModule 添加到给定的 JITDylib。大多数从 IRLayer 派生的层都可以依赖 add
方法的这个默认实现。
这两个操作 add
和 emit
共同构成了层概念:层是一种包装编译器管道的一部分(在本例中是 LLVM 编译器的“opt”阶段)的方式,该管道的 API 对 ORC 是不透明的,但 ORC 可以根据需要调用它。add 方法接受某种输入程序表示形式(在本例中为 LLVM IR 模块)的模块,并将其存储在目标 JITDylib
中,安排在请求该模块定义的任何符号时将其传递回层的 emit 方法。每个层都可以通过调用其基础层的 emit
方法来完成自己的工作。例如,在本教程中,我们的 IRTransformLayer 调用到我们的 IRCompileLayer 以编译转换后的 IR,而我们的 IRCompileLayer 又调用我们的 ObjectLayer 来链接我们的编译器生成的对象文件。
到目前为止,我们已经学习了如何优化和编译我们的 LLVM IR,但我们还没有关注编译何时发生。我们当前的 REPL 在每个函数被任何其他代码引用时立即对其进行优化和编译,无论它是否在运行时被调用过。在下一章中,我们将介绍完全延迟编译,其中函数在首次在运行时被调用之前不会被编译。此时,权衡变得更加有趣:我们越延迟,我们就可以越快开始执行第一个函数,但是我们将越频繁地暂停以编译新遇到的函数。如果我们只延迟代码生成,但积极优化,我们将有更长的启动时间(因为所有内容都在那时优化),但相对较短的暂停,因为每个函数都只是通过代码生成。如果我们既延迟优化又延迟代码生成,我们可以更快地开始执行第一个函数,但是我们将有更长的暂停,因为每个函数都必须在首次执行时进行优化和代码生成。如果我们考虑过程间优化(如内联),情况会变得更加有趣,内联必须积极执行。这些是复杂的权衡,没有一劳永逸的解决方案,但是通过提供可组合的层,我们将决策留给实现 JIT 的人,并使他们可以轻松地尝试不同的配置。
2.3. 完整代码清单¶
这是我们运行示例的完整代码清单,其中添加了 IRTransformLayer 以启用优化。要构建此示例,请使用
# Compile
clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -O3 -o toy
# Run
./toy
这是代码
//===- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope --------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.net.cn/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Contains a simple JIT definition for use in the kaleidoscope tutorials.
//
//===----------------------------------------------------------------------===//
#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#include "llvm/ADT/StringRef.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/InstCombine/InstCombine.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Scalar/GVN.h"
#include <memory>
namespace llvm {
namespace orc {
class KaleidoscopeJIT {
private:
std::unique_ptr<ExecutionSession> ES;
DataLayout DL;
MangleAndInterner Mangle;
RTDyldObjectLinkingLayer ObjectLayer;
IRCompileLayer CompileLayer;
IRTransformLayer OptimizeLayer;
JITDylib &MainJD;
public:
KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
JITTargetMachineBuilder JTMB, DataLayout DL)
: ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
MainJD(this->ES->createBareJITDylib("<main>")) {
MainJD.addGenerator(
cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
DL.getGlobalPrefix())));
}
~KaleidoscopeJIT() {
if (auto Err = ES->endSession())
ES->reportError(std::move(Err));
}
static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
auto EPC = SelfExecutorProcessControl::Create();
if (!EPC)
return EPC.takeError();
auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
JITTargetMachineBuilder JTMB(
ES->getExecutorProcessControl().getTargetTriple());
auto DL = JTMB.getDefaultDataLayoutForTarget();
if (!DL)
return DL.takeError();
return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(JTMB),
std::move(*DL));
}
const DataLayout &getDataLayout() const { return DL; }
JITDylib &getMainJITDylib() { return MainJD; }
Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
if (!RT)
RT = MainJD.getDefaultResourceTracker();
return OptimizeLayer.add(RT, std::move(TSM));
}
Expected<ExecutorSymbolDef> lookup(StringRef Name) {
return ES->lookup({&MainJD}, Mangle(Name.str()));
}
private:
static Expected<ThreadSafeModule>
optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
TSM.withModuleDo([](Module &M) {
// Create a function pass manager.
auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
// Add some optimizations.
FPM->add(createInstructionCombiningPass());
FPM->add(createReassociatePass());
FPM->add(createGVNPass());
FPM->add(createCFGSimplificationPass());
FPM->doInitialization();
// Run the optimizations over all functions in the module being added to
// the JIT.
for (auto &F : M)
FPM->run(F);
});
return std::move(TSM);
}
};
} // end namespace orc
} // end namespace llvm
#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H