2. JITの構築：最適化の追加 – ORCレイヤー入門¶
このチュートリアルは現在開発中です。不完全であり、詳細は頻繁に変更される可能性があります。それにもかかわらず、現状のまま試していただき、フィードバックをお待ちしております。
2.1. 第2章 イントロダクション¶
警告：このチュートリアルは現在、ORC APIの変更に対応するために更新中です。最新の情報は第1章と第2章のみです。
第3章から第5章のサンプルコードはコンパイルと実行は可能ですが、更新されていません。
「LLVMにおけるORCベースのJITの構築」チュートリアルの第2章へようこそ。第1章では、LLVM IRモジュールを入力として受け取り、メモリ内に実行可能なコードを生成できる基本的なJITクラスであるKaleidoscopeJITを調べました。KaleidoscopeJITは、多くの作業を行うために、IRCompileLayerとObjectLinkingLayerという2つの既製のORCレイヤーを組み合わせることで、比較的少ないコードでこれを実現することができました。
このレイヤーでは、新しいレイヤーであるIRTransformLayerを使用してKaleidoscopeJITにIR最適化サポートを追加することにより、ORCレイヤーの概念についてさらに学習します。
2.2. IRTransformLayerを使用したモジュールの最適化¶
「LLVMを使用した言語の実装」チュートリアルシリーズの第4章では、LLVM IRを最適化するための手段として、llvmのFunctionPassManagerが紹介されています。興味のある読者はその章を参照して詳細を確認できますが、簡単に言うと、モジュールを最適化するには、llvm::FunctionPassManagerインスタンスを作成し、一連の最適化を設定してから、PassManagerをモジュールで実行して、意味的に同等の最適化された（おそらく）形式に変換します。元のチュートリアルシリーズでは、FunctionPassManagerはKaleidoscopeJITの外側に作成され、モジュールはKaleidoscopeJITに追加される前に最適化されていました。この章では、最適化をJITのフェーズにします。現時点では、これによりORCレイヤーについてさらに学ぶ動機が提供されますが、長期的に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の後に、CompileLayerの上に配置される新しいメンバーであるTransformLayerを導入します。実行セッションと出力レイヤー（レイヤーの標準的な慣例）への参照と、変換関数を使用してOptimizeLayerを初期化します。変換関数として、クラスのoptimizeModule静的メソッドを提供します。
// ...
return cantFail(OptimizeLayer.addModule(std::move(M),
std::move(Resolver)));
// ...
次に、CompileLayer::addへの呼び出しをOptimizeLayer::addへの呼び出しに置き換えるために、addModuleメソッドを更新する必要があります。
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を設定し、いくつかのパスを追加して、モジュールのすべての関数で実行し、変更されたモジュールを返します。具体的な最適化は、「LLVMを使用した言語の実装」チュートリアルシリーズの第4章で使用されているものと同じです。読者は、これらの詳細と、一般的にIR最適化について、その章を参照してください。
KaleidoscopeJITへの変更は以上です。addModuleを介してモジュールが追加されると、OptimizeLayerは変換されたモジュールを下のCompileLayerに渡す前に、optimizeModule関数を呼び出します。もちろん、addModule関数で直接optimizeModuleを呼び出し、IRTransformLayerを使用する手間をかける必要はありませんでしたが、そうすることで、レイヤーの構成方法をさらに確認する機会が得られます。また、IRTransformLayerは実装できる最も単純なレイヤーの1つであるため、レイヤーの概念自体への明確なエントリポイントも提供します。
// 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());
}
}
これは、llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.hとllvm/lib/ExecutionEngine/Orc/IRTransformLayer.cppからのIRTransformLayerの完全な定義です。このクラスは、2つの非常に単純なタスクに関心があります。（1）このレイヤーを介して出力されるすべてのIRモジュールを変換関数オブジェクトで実行し、（2）ORC IRLayerインターフェースを実装します（それ自体は一般的なORCレイヤーの概念に準拠し、後述します）。クラスの大部分は単純です。変換関数のtypedef、メンバーを初期化するコンストラクタ、変換関数値のセッター、デフォルトの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からのもので、MaterializationUnit（この場合はBasicIRLayerMaterializationUnit）にラップすることによって、指定されたJITDylibにThreadSafeModuleを追加します。IRLayerから派生したほとんどのレイヤーは、addメソッドのこのデフォルトの実装に依存できます。
これらの2つの操作であるaddとemitは、合わせてレイヤーの概念を構成します。レイヤーは、ORCにとってAPIが不透明なコンパイラパイプラインの一部（この場合はLLVMコンパイラの「opt」フェーズ）を、ORCが必要に応じて呼び出すことができるインターフェースでラップする方法です。addメソッドは、ある入力プログラム表現（この場合はLLVM IRモジュール）のモジュールを取り、ターゲットJITDylibに格納し、そのモジュールによって定義されたシンボルが要求されたときに、レイヤーのemitメソッドに渡されるように配置します。各レイヤーは、ベースレイヤーのemitメソッドを呼び出すことで独自の作業を完了できます。たとえば、このチュートリアルでは、IRTransformLayerは変換されたIRをコンパイルするためにIRCompileLayerを呼び出し、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.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/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