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 之前就被优化了。在本节中,我们将优化作为 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 之后,我们引入了一个新的成员 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 层概念,稍后将详细介绍)。该类的大部分内容都很简单:转换函数的类型定义、初始化成员的构造函数、转换函数值的设置器以及默认的无操作转换。最重要的是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 来链接我们的编译器生成的 obj 文件。

到目前为止,我们已经学习了如何优化和编译我们的 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