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),以及對新類別參考的參考:MaterializationResponsibilityMaterializationResponsibility 引數可用於查詢正在轉換的模組的 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.hllvm/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 方法的預設實作。

這兩個操作,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