2. 构建 JIT:添加优化 – ORC 层简介

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

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 的一种手段。感兴趣的读者可以阅读该章节了解详细信息,但简而言之:要优化一个模块,我们创建一个 llvm::FunctionPassManager 实例,使用一组优化配置它,然后在模块上运行 PassManager,将其变异为(希望)更优化但语义等价的形式。在原始教程系列中,FunctionPassManager 是在 KaleidoscopeJIT 外部创建的,模块在添加到它之前进行了优化。在本章中,我们将使优化成为 JIT 的一个阶段。目前,这将为我们提供学习更多关于 ORC 层的动力,但从长远来看,将优化作为 JIT 的一部分将产生一个重要的好处:当我们开始延迟编译代码(即,将每个函数的编译推迟到第一次运行时)时,由我们的 JIT 管理优化将允许我们也能延迟优化,而不是必须预先完成所有的优化。

为了为我们的 JIT 添加优化支持,我们将采用第 1 章中的 KaleidoscopeJIT,并在其之上组合一个 ORC IRTransformLayer。我们将在下面更详细地了解 IRTransformLayer 的工作原理,但接口很简单:此层的构造函数接受对执行会话和下层(所有层都这样做)的引用,以及一个 IR 优化函数,它将应用于通过 addModule 添加的每个模块

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),以及对新类 MaterializationResponsibility 的引用的引用。MaterializationResponsibility 参数可用于查询正在转换的模块的 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 层概念,稍后会详细介绍)。这个类的大部分内容都很简单:转换函数的 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,通过将其包装在 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.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/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