1. 建立 JIT:從 KaleidoscopeJIT 開始

1.1. 第 1 章 簡介

警告:本教學目前正在更新中,以反映 ORC API 的變更。目前只有第 1 章和第 2 章是最新的。

第 3 章到第 5 章的範例程式碼可以編譯和執行,但尚未更新

歡迎來到「在 LLVM 中建立基於 ORC 的 JIT」教學的第 1 章。本教學將逐步說明如何使用 LLVM 的隨需編譯 (ORC) API 來實作 JIT 編譯器。它首先介紹在「使用 LLVM 實作語言」教學中使用的 KaleidoscopeJIT 類別的簡化版本,然後介紹新的功能,例如並行編譯、最佳化、延遲編譯和遠端執行。

本教學的目標是向您介紹 LLVM 的 ORC JIT API,展示這些 API 如何與 LLVM 的其他部分互動,並教導您如何重新組合它們以建立適合您使用情境的客製化 JIT。

本教學的結構如下:

  • 第 1 章:探討簡單的 KaleidoscopeJIT 類別。這將介紹 ORC JIT API 的一些基本概念,包括 ORC 的概念。

  • 第 2 章:透過新增一個新的層來擴充基本的 KaleidoscopeJIT,該層將最佳化 IR 和產生的程式碼。

  • 第 3 章:透過新增隨需編譯層來延遲編譯 IR,進一步擴充 JIT。

  • 第 4 章:透過使用 ORC 編譯回呼 API 直接將 IR 產生延遲到函式被呼叫時,以自訂層取代隨需編譯層,來改善 JIT 的延遲性。

  • 第 5 章:使用 JIT 遠端 API 將程式碼 JIT 編譯到具有較低權限的遠端程序中,以新增程序隔離。

為了提供 JIT 的輸入,我們將使用「在 LLVM 中實作語言」教學的第 7 章中稍微修改過的 Kaleidoscope REPL 版本。

最後,關於 API 世代的一些說明:ORC 是 LLVM JIT API 的第三代。它之前是 MCJIT,而 MCJIT 之前是(現在已刪除的)傳統 JIT。這些教學不會假設您有任何使用這些舊 API 的經驗,但熟悉這些 API 的讀者會看到許多熟悉的元素。在適當的情況下,我們會明確說明與舊 API 的關聯,以幫助從舊 API 轉換到 ORC 的使用者。

1.2. JIT API 基礎

JIT 編譯器的目的是在需要時「即時」編譯代碼,而不是像傳統編譯器那樣提前將整個程序編譯到磁盤。為了支持這一目標,我們最初的、精簡的 JIT API 將只有兩個函數

  1. Error addModule(std::unique_ptr<Module> M):使給定的 IR 模組可供執行。

  2. Expected<ExecutorSymbolDef> lookup():搜索已添加到 JIT 的符號(函數或變數)的指標。

此 API 的基本用例是執行模組中的「main」函數,如下所示

JIT J;
J.addModule(buildModule());
auto *Main = J.lookup("main").getAddress().toPtr<int(*)(int, char *[])>();
int Result = Main();

我們在這些教程中構建的 API 都將是對這個簡單主題的變體。在此 API 背後,我們將改進 JIT 的實現,以添加對並發編譯、優化和延遲編譯的支持。最終,我們將擴展 API 本身,以允許將更高級別的程序表示形式(例如 AST)添加到 JIT。

1.3. KaleidoscopeJIT

在上一節中,我們描述了我們的 API,現在我們來研究一下它的簡單實現:使用 LLVM 實現一種語言教程中使用的 KaleidoscopeJIT 類別 [1]。我們將使用該教程第 7 章中的 REPL 代碼為我們的 JIT 提供輸入:每次用戶輸入表達式時,REPL 都會向 JIT 添加一個新的 IR 模組,其中包含該表達式的代碼。如果表達式是頂級表達式,例如「1+1」或「sin(x)」,REPL 還將使用我們的 JIT 類別的查找方法查找並執行該表達式的代碼。在本教程的後續章節中,我們將修改 REPL 以實現與我們的 JIT 類別的新交互,但現在我們將理所當然地使用此設置,並將注意力集中在 JIT 本身的實現上。

我們的 KaleidoscopeJIT 類別在 KaleidoscopeJIT.h 標頭中定義。在通常的 include 防護和 #includes [2] 之後,我們來到了類別的定義

#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/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LLVMContext.h"
#include <memory>

namespace llvm {
namespace orc {

class KaleidoscopeJIT {
private:
  ExecutionSession ES;
  RTDyldObjectLinkingLayer ObjectLayer;
  IRCompileLayer CompileLayer;

  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))),
        DL(std::move(DL)), Mangle(ES, this->DL),
        Ctx(std::make_unique<LLVMContext>()) {
    ES.getMainJITDylib().addGenerator(
        cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(DL.getGlobalPrefix())));
  }

我們的類別以六個成員變數開頭:一個 ExecutionSession 成員,ES,它為我們運行的 JIT 代碼提供上下文(包括字符串池、全局互斥鎖和錯誤報告設施);一個 RTDyldObjectLinkingLayer,ObjectLayer,可用於向我們的 JIT 添加目標文件(儘管我們不會直接使用它);一個 IRCompileLayer,CompileLayer,可用於向我們的 JIT 添加 LLVM 模組(並建立在 ObjectLayer 之上),一個 DataLayout 和 MangleAndInterner,DLMangle,將用於符號修飾(稍後將詳細介紹);最後一個 LLVMContext,客戶端在為 JIT 構建 IR 文件時將使用它。

接下來是我們的類別建構函式,它接受一個 JITTargetMachineBuilder`,它將由我們的 IRCompiler 使用,以及一個 DataLayout,我們將使用它來初始化我們的 DL 成員。建構函式首先初始化我們的 ObjectLayer。ObjectLayer 需要一個對 ExecutionSession 的引用,以及一個函式物件,該物件將為添加的每個模組建立一個 JIT 記憶體管理器(JIT 記憶體管理器管理記憶體分配、記憶體許可權和 JIT 程式碼的例外處理常式註冊)。為此,我們使用了一個 lambda,它返回一個 SectionMemoryManager,這是一個現成的實用程式,提供了本章所需的所有基本記憶體管理功能。接下來,我們初始化我們的 CompileLayer。CompileLayer 需要三樣東西:(1) 對 ExecutionSession 的引用,(2) 對我們的物件層的引用,以及 (3) 一個編譯器實例,用於執行從 IR 到物件檔的實際編譯。我們使用現成的 ConcurrentIRCompiler 實用程式作為我們的編譯器,我們使用這個建構函式的 JITTargetMachineBuilder 參數來建構它。ConcurrentIRCompiler 實用程式將根據需要使用 JITTargetMachineBuilder 來建構 llvm TargetMachines(它們不是執行緒安全的)以進行編譯。之後,我們使用輸入的 DataLayout、ExecutionSession 和 DL 成員,以及一個新的預設建構的 LLVMContext,分別初始化我們的支援成員:DLManglerCtx。現在我們的成員已經被初始化了,所以剩下的唯一一件事就是調整我們將要儲存程式碼的 JITDylib 的配置。我們想要修改這個 dylib,使其不僅包含我們添加到其中的符號,還包含我們 REPL 程序中的符號。我們通過使用 DynamicLibrarySearchGenerator::GetForCurrentProcess 方法附加一個 DynamicLibrarySearchGenerator 實例來實現這一點。

static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
  auto JTMB = JITTargetMachineBuilder::detectHost();

  if (!JTMB)
    return JTMB.takeError();

  auto DL = JTMB->getDefaultDataLayoutForTarget();
  if (!DL)
    return DL.takeError();

  return std::make_unique<KaleidoscopeJIT>(std::move(*JTMB), std::move(*DL));
}

const DataLayout &getDataLayout() const { return DL; }

LLVMContext &getContext() { return *Ctx.getContext(); }

接下來我們有一個命名的建構函式,Create,它將建構一個 KaleidoscopeJIT 實例,該實例被配置為為我們的宿主程序生成程式碼。它首先使用該類別的 detectHost 方法生成一個 JITTargetMachineBuilder 實例,然後使用該實例為目標程序生成一個數據佈局。這些操作都可能失敗,因此每個操作都返回一個包裝在 Expected 值 [3] 中的結果,我們必須在繼續之前檢查錯誤。如果兩個操作都成功,我們可以解開它們的結果(使用解引用運算符),並將它們傳遞給函數最後一行的 KaleidoscopeJIT 的建構函式。

在命名的建構函式之後,我們有 getDataLayout()getContext() 方法。這些方法用於使 JIT 建立和管理的數據結構(尤其是 LLVMContext)可供將建構我們的 IR 模組的 REPL 程式碼使用。

void addModule(std::unique_ptr<Module> M) {
  cantFail(CompileLayer.add(ES.getMainJITDylib(),
                            ThreadSafeModule(std::move(M), Ctx)));
}

Expected<ExecutorSymbolDef> lookup(StringRef Name) {
  return ES.lookup({&ES.getMainJITDylib()}, Mangle(Name.str()));
}

現在我們來看看第一個 JIT API 方法:addModule。這個方法負責將 IR 添加到 JIT 並使其可執行。在我們 JIT 的初始實現中,我們將通過將模組添加到 CompileLayer 來使其「可執行」,而 CompileLayer 又會將模組儲存在主要的 JITDylib 中。這個過程會在 JITDylib 中為模組中的每個定義創建新的符號表條目,並延遲模組的編譯,直到它的任何定義被查找為止。請注意,這不是懶惰編譯:即使從未使用過,僅僅引用一個定義就足以觸發編譯。在後面的章節中,我們將教導我們的 JIT 將函數的編譯延遲到它們被實際調用時才進行。要添加我們的模組,我們必須首先將其包裝在 ThreadSafeModule 實例中,該實例以線程安全的方式管理模組的 LLVMContext(我們的 Ctx 成員)的生命週期。在我們的示例中,所有模組將共享 Ctx 成員,該成員將在 JIT 的整個生命週期中存在。在後面的章節中,當我們切換到並發編譯時,我們將為每個模組使用一個新的上下文。

我們最後一個方法是 lookup,它允許我們根據符號名稱查找添加到 JIT 的函數和變量定義的地址。如上所述,lookup 將隱式地為任何尚未被編譯的符號觸發編譯。我們的 lookup 方法會調用 ExecutionSession::lookup,傳入要搜索的 dylibs 列表(在我們的例子中只有主要的 dylib)以及要搜索的符號名稱,但有一個轉折:我們必須先對要搜索的符號名稱進行*名稱修飾*。ORC JIT 組件在內部使用與靜態編譯器和鏈接器相同的方式使用修飾後的符號,而不是使用普通的 IR 符號名稱。這使得 JIT 編譯的代碼可以輕鬆地與應用程序或共享庫中的預編譯代碼進行交互操作。名稱修飾的種類取決於 DataLayout,而 DataLayout 又取決於目標平台。為了讓我們保持可移植性並根據未經修飾的名稱進行搜索,我們只需使用我們的 Mangle 成員函數對象自己重新生成這個名稱修飾。

我們現在來到了構建 JIT 的第 1 章的末尾。您現在有了一個基本但功能齊全的 JIT 堆棧,您可以使用它來獲取 LLVM IR 並使其在您的 JIT 進程的上下文中可執行。在下一章中,我們將探討如何擴展這個 JIT 以生成更高質量的代碼,並在此過程中更深入地了解 ORC 層的概念。

下一章:擴展 KaleidoscopeJIT

1.4. 完整代碼列表

以下是我們運行示例的完整代碼列表。要構建此示例,請使用

# 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/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 <memory>

namespace llvm {
namespace orc {

class KaleidoscopeJIT {
private:
  std::unique_ptr<ExecutionSession> ES;

  DataLayout DL;
  MangleAndInterner Mangle;

  RTDyldObjectLinkingLayer ObjectLayer;
  IRCompileLayer CompileLayer;

  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))),
        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 CompileLayer.add(RT, std::move(TSM));
  }

  Expected<ExecutorSymbolDef> lookup(StringRef Name) {
    return ES->lookup({&MainJD}, Mangle(Name.str()));
  }
};

} // end namespace orc
} // end namespace llvm

#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H