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 之前就被优化了。在本节中,我们将优化作为 JIT 的一个阶段。目前,这将为我们提供学习更多关于 ORC 层的动力,但从长远来看,将优化作为 JIT 的一部分将带来一个重要的益处:当我们开始延迟编译代码(即延迟编译每个函数,直到它第一次运行)时,由 JIT 管理的优化将允许我们也延迟优化,而不是必须预先进行所有优化。
为了向我们的 JIT 添加优化支持,我们将采用第 1 章中的 KaleidoscopeJIT,并在其之上组合一个 ORC IRTransformLayer。我们将在下面详细了解 IRTransformLayer 的工作原理,但其接口很简单:该层的构造函数接受对执行会话和下层(所有层都一样)的引用,以及一个它将应用于通过 addModule 添加的每个模块的IR 优化函数。
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 层概念,稍后将详细介绍)。该类的大部分内容都很简单:转换函数的类型定义、初始化成员的构造函数、转换函数值的设置器以及默认的无操作转换。最重要的是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 来链接我们的编译器生成的 obj 文件。
到目前为止,我们已经学习了如何优化和编译我们的 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