2. 建立 JIT:新增最佳化 - ORC Layers 簡介¶
本教學仍在積極開發中。它尚未完工,細節可能會經常變更。儘管如此,我們仍然邀請您試用它,並歡迎任何回饋。
2.1. 第 2 章 簡介¶
警告:本教學目前正在更新中,以反映 ORC API 的變更。目前只有第 1 章和第 2 章是最新的。
第 3 章到第 5 章的範例程式碼可以編譯和執行,但尚未更新
歡迎來到「在 LLVM 中建構基於 ORC 的 JIT」教學的第 2 章。在本系列的第 1 章中,我們檢視了一個基本的 JIT 類別 KaleidoscopeJIT,它可以將 LLVM IR 模組作為輸入,並在記憶體中產生可執行程式碼。KaleidoscopeJIT 透過組合兩個現成的 *ORC Layers*:IRCompileLayer 和 ObjectLinkingLayer,來完成大部分的繁重工作,因此只需相對較少的程式碼即可完成此操作。
在本層中,我們將透過使用新的 Layer IRTransformLayer 來進一步瞭解 ORC Layer 的概念,以將 IR 最佳化支援新增至 KaleidoscopeJIT。
2.2. 使用 IRTransformLayer 最佳化模組¶
在「使用 LLVM 實作語言」教學系列的第 4 章中,介紹了 llvm *FunctionPassManager* 作為最佳化 LLVM IR 的方法。有興趣的讀者可以閱讀該章以瞭解詳細資訊,但簡而言之:為了最佳化模組,我們建立一個 llvm::FunctionPassManager 執行個體,使用一組最佳化來配置它,然後在模組上執行 PassManager,將其變異為(希望)更最佳化但語義上等效的形式。在原始的教學系列中,FunctionPassManager 是在 KaleidoscopeJIT 之外建立的,並且在新增至 KaleidoscopeJIT 之前會先對模組進行最佳化。在本章中,我們將使最佳化成為 JIT 的一個階段。目前,這將為我們提供學習更多關於 ORC Layers 的動機,但從長遠來看,將最佳化作為 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,但這樣做讓我們有另一個機會了解層是如何組成的。它也為 layer 概念本身提供了一個簡潔的切入點,因為 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 Layer 概念,稍後會詳細介紹)。這個類別的大部分內容都很簡單:轉換函式的 typedef、用於初始化成員的建構函式、用於設置轉換函式值的設置器,以及預設的無操作轉換。最重要的
我們從 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
,透過將 ThreadSafeModule 封裝在 MaterializationUnit
(在這種情況下為 BasicIRLayerMaterializationUnit
)中,將其添加到指定的 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.dev.org.tw/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