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.hllvm/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 方法的這個預設實現。

這兩個操作,addemit,共同構成了層的概念:層是一種將編譯器管道的一部分(在這種情況下是 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