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 *レイヤー*の概念を含む、ORC JIT API の基本的な概念の一部が紹介されます。

  • 第 2 章: IR と生成されたコードを最適化する新しいレイヤーを追加して、基本的な KaleidoscopeJIT を拡張します。

  • 第 3 章: オンデマンドコンパイルレイヤーを追加して、JIT をさらに拡張し、IR を遅延コンパイルします。

  • 第 4 章: 関数が呼び出されるまで IR 生成を遅延させるために、ORC コンパイルコールバック API を直接使用するカスタムレイヤーでオンデマンドコンパイルレイヤーを置き換えることで、JIT の遅延性を向上させます。

  • 第 5 章: JIT リモート API を使用して権限を減らしたリモートプロセスにコードを JIT することにより、プロセス分離を追加します。

JIT の入力として、「LLVM チュートリアルで言語を実装する」第 7 章からの Kaleidoscope REPL のわずかに変更されたバージョンを使用します。

最後に、API の世代について一言。ORC は LLVM JIT API の第 3 世代です。その前は MCJIT で、その前は (現在削除されている) レガシー JIT でした。これらのチュートリアルでは、以前のこれらの API の経験を前提としていませんが、それらに精通している読者は多くの類似した要素を目にするでしょう。適切な場合は、以前の API から ORC に移行する人を支援するために、以前の API とのこの接続を明示的に行います。

1.2. JIT API の基本

JIT コンパイラの目的は、従来のコンパイラのようにプログラム全体を事前にディスクにコンパイルするのではなく、必要に応じて「オンザフライ」でコードをコンパイルすることです。その目的をサポートするために、最初の、必要最低限の JIT API には 2 つの関数しかありません。

  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 はその式のコードを含む新しい IR モジュールを JIT に追加します。式が「1 + 1」や「sin(x)」のようなトップレベルの式である場合、REPL は JIT クラスの lookup メソッドを使用して、式のコードを見つけて実行します。このチュートリアルの後の章では、JIT クラスとの新しいインタラクションを有効にするために REPL を変更しますが、今のところ、この設定を当然のこととして受け入れ、JIT 自体の実装に注目します。

KaleidoscopeJIT クラスは、KaleidoscopeJIT.h ヘッダーで定義されています。通常のインクルードガードと #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())));
  }

クラスは、6 つのメンバ変数で始まります。ExecutionSession メンバ ES。実行中の JIT コードのコンテキスト (文字列プール、グローバルミューテックス、エラー報告機能を含む) を提供します。RTDyldObjectLinkingLayer ObjectLayer。オブジェクトファイルを JIT に追加するために使用できます (ただし、直接は使用しません)。IRCompileLayer CompileLayer。LLVM モジュールを JIT に追加するために使用できる (ObjectLayer を基盤とする) もの。DataLayout と MangleAndInterner DLMangle。シンボルマングリングに使用されます (詳細は後述)。最後に、クライアントが JIT の IR ファイルを構築するときに使用する LLVMContext です。

次に、クラスコンストラクタがあります。これは、IRCompiler で使用される JITTargetMachineBuilder` と、DL メンバを初期化するために使用される DataLayout を受け取ります。コンストラクタは、ObjectLayer の初期化から始まります。ObjectLayer には、ExecutionSession への参照と、追加された各モジュールの JIT メモリマネージャを構築する関数オブジェクトが必要です (JIT メモリマネージャは、メモリ割り当て、メモリパーミッション、および JIT コードの例外ハンドラの登録を管理します)。このために、SectionMemoryManager (この章に必要なすべての基本的なメモリ管理機能を提供する既製のユーティリティ) を返すラムダを使用します。次に、CompileLayer を初期化します。CompileLayer には、(1) ExecutionSession への参照、(2) オブジェクトレイヤーへの参照、および (3) IR からオブジェクトファイルへの実際のコンパイルを実行するために使用するコンパイラインスタンスの 3 つが必要です。コンパイラとして、既製の ConcurrentIRCompiler ユーティリティを使用します。これは、このコンストラクタの JITTargetMachineBuilder 引数を使用して構築します。ConcurrentIRCompiler ユーティリティは、コンパイルに必要な llvm TargetMachine (これはスレッドセーフではありません) を構築するために JITTargetMachineBuilder を使用します。その後、DLMangler、および Ctx の各サポートメンバを、入力 DataLayout、ExecutionSession、DL メンバ、および新しくデフォルトで構築された LLVMContext でそれぞれ初期化します。メンバが初期化されたので、残りの作業は、コードを格納する *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 に追加することにより、モジュールを「実行可能」にします。これにより、モジュールはメインの JITDylib に格納されます。このプロセスにより、モジュール内の各定義に対して JITDylib に新しいシンボルテーブルエントリが作成され、定義のいずれかが参照されるまでモジュールのコンパイルが遅延されます。これは遅延コンパイルではないことに注意してください。たとえ使用されなくても、定義を参照するだけでコンパイルがトリガーされます。後の章では、実際に呼び出されるまで関数のコンパイルを遅延させるように JIT を教えます。モジュールを追加するには、まず、スレッドフレンドリーな方法でモジュールの LLVMContext (Ctx メンバ) の有効期間を管理する ThreadSafeModule インスタンスでモジュールをラップする必要があります。この例では、すべてのモジュールが Ctx メンバを共有し、これは JIT の実行期間中存在します。後の章で同時コンパイルに切り替えたら、モジュールごとに新しいコンテキストを使用します。

最後のメソッドは lookup です。これにより、JITに追加された関数および変数の定義のアドレスを、それらのシンボル名に基づいて検索できます。上記で述べたように、lookup は、まだコンパイルされていないシンボルに対して暗黙的にコンパイルをトリガーします。この lookup メソッドは ExecutionSession::lookup を呼び出し、検索する dylib のリスト(この場合はメインの dylib のみ)と、検索するシンボル名を渡します。ただし、少し工夫が必要です。まず、検索するシンボルの名前をマングルする必要があります。ORC JIT コンポーネントは、静的コンパイラやリンカーと同様に、内部でマングルされたシンボルを使用します。プレーンな IR シンボル名を使用するのではなくです。これにより、JIT コンパイルされたコードが、アプリケーションや共有ライブラリ内のプリコンパイルされたコードと簡単に連携できるようになります。マングリングの種類は、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.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/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