3. 构建 JIT:按函数延迟编译

本教程正在积极开发中。它尚未完成,细节可能会频繁更改。尽管如此,我们邀请您尝试目前的状态,并欢迎任何反馈。

3.1. 第 3 章 介绍

警告:由于 ORC API 更新,本文档目前已过时。

示例代码已更新,可以使用。文本将在 API 变动平息后更新。

欢迎来到“在 LLVM 中构建基于 ORC 的 JIT”教程的第 3 章。本章讨论延迟 JIT 编译,并向您展示如何通过从第 2 章的 JIT 中添加 ORC CompileOnDemand 层来启用它。

3.2. 延迟编译

当我们将模块添加到第 2 章的 KaleidoscopeJIT 类时,IRTransformLayer、IRCompileLayer 和 RTDyldObjectLinkingLayer 分别会立即为我们优化、编译和链接它。这种方案,其中使模块可执行的所有工作都在预先完成,简单易懂,其性能特征也很容易推断。但是,如果要编译的代码量很大,这将导致非常高的启动时间,并且如果运行时只调用少数编译函数,也可能会进行大量不必要的编译。一个真正的“即时”编译器应该允许我们将任何给定函数的编译推迟到首次调用该函数时,从而缩短启动时间并消除冗余工作。实际上,ORC API 为我们提供了一个延迟编译 LLVM IR 的层:CompileOnDemandLayer。

CompileOnDemandLayer 类符合第 2 章中描述的层接口,但其 addModule 方法的行为与我们到目前为止看到的层截然不同:它不会预先做任何工作,而只是扫描正在添加的模块,并安排在首次调用模块中的每个函数时编译它们。为此,CompileOnDemandLayer 为其扫描的每个函数创建两个小实用程序:一个桩(stub)和一个编译回调(compile callback)。桩是一对函数指针(一旦函数被编译,它将指向函数的实现)和一个通过该指针的间接跳转。通过在程序的生命周期内固定间接跳转的地址,我们可以给函数一个永久的“有效地址”,即使函数的实现从未被编译,或者即使它被编译多次(例如,由于在更高的优化级别重新编译函数)并且地址发生更改,该地址也可以安全地用于间接寻址和函数指针比较。第二个实用程序,编译回调,表示从程序到编译器的重新入口点,它将触发函数的编译然后执行。通过初始化函数的桩以指向函数的编译回调,我们启用了延迟编译:首次尝试调用该函数将跟随函数指针并触发编译回调。编译回调将编译函数,更新桩的函数指针,然后执行该函数。在所有后续对该函数的调用中,函数指针将指向已编译的函数,因此编译器不会产生进一步的开销。我们将在本教程的下一章中更详细地研究这个过程,但现在我们将信任 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 头文件,然后添加两个新成员:std::unique_ptr<JITCompileCallbackManager> 和 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 将调用此函数来询问我们想要编译什么。至少我们需要编译正在调用的函数(由分区函数的参数给出),但我们也可以请求 CompileOnDemandLayer 编译从正在调用的函数中无条件调用(或极有可能被调用)的其他函数。对于 KaleidoscopeJIT,我们将保持简单,只请求编译被调用的函数。接下来,我们传递对我们的 CompileCallbackManager 的引用。最后,我们需要提供一个“间接桩管理器构建器”:一个实用程序函数,用于构建 IndirectStubManagers,而 IndirectStubManagers 又用于为每个模块中的函数构建桩。CompileOnDemandLayer 将为每次调用 addModule 调用一次间接桩管理器构建器,并使用生成的间接桩管理器为集合中所有模块中的所有函数创建桩。如果/当模块集合从 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.net.cn/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/IRPartitionLayer.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;
  IRPartitionLayer IPLayer;
  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),
        IPLayer(*this->ES, OptimizeLayer),
        CODLayer(*this->ES, IPLayer, 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 编译