3. JITの構築:関数ごとの遅延コンパイル

このチュートリアルは現在開発中です。未完成であり、詳細は頻繁に変更される可能性があります。それにもかかわらず、現状のまま試してみていただき、フィードバックをお待ちしております。

3.1. 第3章 イントロダクション

警告:このテキストは、ORC APIの更新により現在古くなっています。

サンプルコードは更新されており、使用できます。APIの変更が落ち着いたら、テキストも更新されます。

「LLVMにおけるORCベースのJITの構築」チュートリアルの第3章へようこそ。この章では、遅延JITについて説明し、第2章のJITにORC CompileOnDemandレイヤーを追加することで、それを有効にする方法を示します。

3.2. 遅延コンパイル

第2章のKaleidoscopeJITクラスにモジュールを追加すると、IRTransformLayer、IRCompileLayer、RTDyldObjectLinkingLayerによって、それぞれすぐに最適化、コンパイル、リンクされます。このスキームでは、モジュールを実行可能にするためのすべての作業が事前に実行されるため、理解しやすく、パフォーマンス特性についても容易に推論できます。しかし、コンパイルするコードの量が多い場合、起動時間が非常に長くなり、実行時に呼び出されるコンパイル済み関数がわずかしかない場合にも、多くの不要なコンパイルが行われる可能性があります。真に「Just-In-Time」コンパイラは、特定の関数のコンパイルを、その関数が最初に呼び出される瞬間まで遅らせることで、起動時間を短縮し、冗長な作業を排除する必要があります。実際、ORC APIは、LLVM IRを遅延コンパイルするためのレイヤー、CompileOnDemandLayerを提供します。

CompileOnDemandLayerクラスは、第2章で説明されているレイヤーインターフェースに準拠していますが、そのaddModuleメソッドの動作は、これまで見てきたレイヤーとは大きく異なります。事前に作業を行うのではなく、追加されるモジュールをスキャンし、それらの各関数が最初に呼び出されたときにコンパイルされるように配置します。これを行うために、CompileOnDemandLayerは、スキャンする各関数に対して、スタブコンパイルコールバックという2つの小さなユーティリティを作成します。(関数がコンパイルされた後、関数の実装を指すようになる)関数ポインタと、そのポインタを介した間接ジャンプです。プログラムの存続期間中に間接ジャンプのアドレスを固定することで、関数に永続的な「有効アドレス」を与えることができます。これは、関数の実装がコンパイルされない場合、または(たとえば、より高い最適化レベルで関数を再コンパイルすることにより)アドレスが変更される場合でも、間接参照や関数ポインタの比較に安全に使用できるアドレスです。2つ目のユーティリティであるコンパイルコールバックは、プログラムからコンパイラへの再エントリポイントを表し、関数のコンパイルと実行をトリガーします。関数のスタブを関数のコンパイルコールバックを指すように初期化することで、遅延コンパイルを有効にします。関数への最初の呼び出しの試行は、関数ポインタに従い、代わりにコンパイルコールバックをトリガーします。コンパイルコールバックは関数をコンパイルし、スタブの関数ポインタを更新してから、関数を実行します。その後の関数の呼び出しでは、関数ポインタは既にコンパイルされた関数を指しているため、コンパイラからの追加のオーバーヘッドはありません。このプロセスについては、このチュートリアルの次の章で詳しく説明しますが、ここでは、CompileOnDemandLayerがすべてのスタブとコールバックを私たちのために設定することを信頼します。必要なのは、スタックの一番上にCompileOnDemandLayerを追加することだけで、遅延コンパイルの利点が得られます。ソースコードを少し変更するだけです。

...
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
...

...
class KaleidoscopeJIT {
private:
  std::unique_ptr<TargetMachine> TM;
  const DataLayout DL;
  RTDyldObjectLinkingLayer ObjectLayer;
  IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;

  using OptimizeFunction =
      std::function<std::shared_ptr<Module>(std::shared_ptr<Module>)>;

  IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;

  std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager;
  CompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer;

public:
  using ModuleHandle = decltype(CODLayer)::ModuleHandleT;

最初に、CompileOnDemandLayer.hヘッダーを含める必要があります。次に、クラスに2つの新しいメンバを追加します。std::unique_ptrとCompileOnDemandLayerです。CompileCallbackManagerメンバは、CompileOnDemandLayerによって、各関数に必要なコンパイルコールバックを作成するために使用されます。

KaleidoscopeJIT()
    : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
      ObjectLayer([]() { return std::make_shared<SectionMemoryManager>(); }),
      CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
      OptimizeLayer(CompileLayer,
                    [this](std::shared_ptr<Module> M) {
                      return optimizeModule(std::move(M));
                    }),
      CompileCallbackManager(
          orc::createLocalCompileCallbackManager(TM->getTargetTriple(), 0)),
      CODLayer(OptimizeLayer,
               [this](Function &F) { return std::set<Function*>({&F}); },
               *CompileCallbackManager,
               orc::createLocalIndirectStubsManagerBuilder(
                 TM->getTargetTriple())) {
  llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
}

次に、新しいメンバを初期化するためにコンストラクタを更新する必要があります。適切なコンパイルコールバックマネージャーを作成するには、createLocalCompileCallbackManager関数を使用します。この関数は、TargetMachineと、不明な関数のコンパイル要求を受け取った場合に呼び出すExecutorAddrを取ります。この単純なJITでは、この状況が発生する可能性は低いため、ここでは「0」を渡してごまかします。本番環境品質のJITでは、JITコードのスタックをアンワインドするために例外をスローする関数のアドレスを渡すことができます。

これで、CompileOnDemandLayerを構築できます。以前のレイヤーのパターンに従って、スタックの下の次のレイヤーへの参照(OptimizeLayer)を渡して開始します。次に、「パーティショニング関数」を提供する必要があります。まだコンパイルされていない関数が呼び出されると、CompileOnDemandLayerはこの関数を呼び出して、コンパイルするものを尋ねます。最低限、呼び出された関数(パーティショニング関数の引数で指定)をコンパイルする必要がありますが、呼び出された関数から無条件に呼び出される(または呼び出される可能性が高い)他の関数のコンパイルも要求できます。KaleidoscopeJITでは、シンプルに、呼び出された関数のコンパイルのみを要求します。次に、CompileCallbackManagerへの参照を渡します。最後に、「間接スタブマネージャービルダー」を提供する必要があります。これは、間接スタブマネージャーを構築するユーティリティ関数であり、間接スタブマネージャーは、各モジュールの関数のスタブを構築するために使用されます。CompileOnDemandLayerは、addModuleの呼び出しごとに間接スタブマネージャービルダーを1回呼び出し、結果の間接スタブマネージャーを使用して、セット内のすべてのモジュール内のすべての関数のスタブを作成します。モジュールセットがJITから削除されると、間接スタブマネージャーが削除され、スタブに割り当てられたメモリが解放されます。createLocalIndirectStubsManagerBuilderユーティリティを使用してこの関数を提供します。

// ...
        if (auto Sym = CODLayer.findSymbol(Name, false))
// ...
return cantFail(CODLayer.addModule(std::move(Ms),
                                   std::move(Resolver)));
// ...

// ...
return CODLayer.findSymbol(MangledNameStream.str(), true);
// ...

// ...
CODLayer.removeModule(H);
// ...

最後に、addModule、findSymbol、removeModuleメソッド内のOptimizeLayerへの参照を置き換える必要があります。これで、実行準備が整いました。

未完了

**章の結論。**

3.3. 完全なコードリスト

関数の遅延コンパイルを有効にするためにCompileOnDemandレイヤーを追加した実行例に関する完全なコードリストを以下に示します。この例をビルドするには、以下を使用します。

# 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.dokyumento.jp/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/CompileOnDemandLayer.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.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;
  std::unique_ptr<EPCIndirectionUtils> EPCIU;

  DataLayout DL;
  MangleAndInterner Mangle;

  RTDyldObjectLinkingLayer ObjectLayer;
  IRCompileLayer CompileLayer;
  IRTransformLayer OptimizeLayer;
  CompileOnDemandLayer CODLayer;

  JITDylib &MainJD;

  static void handleLazyCallThroughError() {
    errs() << "LazyCallThrough error: Could not find function body";
    exit(1);
  }

public:
  KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
                  std::unique_ptr<EPCIndirectionUtils> EPCIU,
                  JITTargetMachineBuilder JTMB, DataLayout DL)
      : ES(std::move(ES)), EPCIU(std::move(EPCIU)), 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),
        CODLayer(*this->ES, OptimizeLayer,
                 this->EPCIU->getLazyCallThroughManager(),
                 [this] { return this->EPCIU->createIndirectStubsManager(); }),
        MainJD(this->ES->createBareJITDylib("<main>")) {
    MainJD.addGenerator(
        cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
            DL.getGlobalPrefix())));
  }

  ~KaleidoscopeJIT() {
    if (auto Err = ES->endSession())
      ES->reportError(std::move(Err));
    if (auto Err = EPCIU->cleanup())
      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));

    auto EPCIU = EPCIndirectionUtils::Create(*ES);
    if (!EPCIU)
      return EPCIU.takeError();

    (*EPCIU)->createLazyCallThroughManager(
        *ES, ExecutorAddr::fromPtr(&handleLazyCallThroughError));

    if (auto Err = setUpInProcessLCTMReentryViaEPCIU(**EPCIU))
      return std::move(Err);

    JITTargetMachineBuilder JTMB(
        ES->getExecutorProcessControl().getTargetTriple());

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

    return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(*EPCIU),
                                             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 CODLayer.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

次へ:極端な遅延 – コンパイルコールバックを使用してASTから直接JITを使用する