2. 建置 JIT:新增最佳化 – ORC Layers 簡介¶
本教學仍在積極開發中。內容尚未完整,細節可能會頻繁變更。 然而,我們邀請您試用目前的版本,並歡迎任何回饋意見。
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 的一種方法。有興趣的讀者可以閱讀該章以了解詳細資訊,但簡而言之:要最佳化 Module,我們建立一個 llvm::FunctionPassManager 實例,使用一組最佳化配置它,然後在 Module 上執行 PassManager 以將其變異為(希望)更最佳化但語義上等效的形式。在原始教學系列中,FunctionPassManager 是在 KaleidoscopeJIT 外部建立的,模組在新增到其中之前就已最佳化。在本章中,我們將使最佳化成為 JIT 的一個階段。目前,這將為我們提供學習更多關於 ORC 層的動機,但從長遠來看,將最佳化作為 JIT 的一部分將產生一個重要的好處:當我們開始延遲編譯程式碼時(即,將每個函式的編譯延遲到第一次執行時),讓最佳化由我們的 JIT 管理將使我們也能夠延遲最佳化,而不必預先完成所有最佳化。
為了為我們的 JIT 新增最佳化支援,我們將採用第 1 章中的 KaleidoscopeJIT,並在其之上組合一個 ORC IRTransformLayer。我們將在下面更詳細地了解 IRTransformLayer 的工作原理,但介面很簡單:此層的建構子接受對執行階段工作階段和下層(與所有層一樣)的參考,以及一個 IR 最佳化函式,它將應用於透過 addModule 新增的每個 Module
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 Layer 概念,稍後會詳細介紹)。這個類別的大部分內容都很簡單:轉換函式的 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
,透過將 ThreadSafeModule 包裝在 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.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