从零开始实现一门编程语言(四)增加支持JIT与优化器

120 阅读16分钟

简介

欢迎来到“用LLVM实现语言”教程的第4章。第1-3章描述了一个简单语言的实现,并增加了对生成LLVM IR的支持。本章描述了两种新技术:为语言添加优化器支持,以及添加JIT编译器支持。这些添加的内容将演示如何为Kaleidoscope语言获得漂亮、高效的代码。

常量折叠

我们第3章的演示是优雅的,易于扩展。不幸的是,它不会产生完美的代码。然而,IRBuilder在编译简单代码时确实为我们提供了明显的优化:

ready> def test(x) 1+2+x;
Read function definition:
define double @test(double %x) {
entry:
        %addtmp = fadd double 3.000000e+00, %x
        ret double %addtmp
}

该代码不是通过解析输入构建的AST的字符转换。那就是:

ready> def test(x) 1+2+x;
Read function definition:
define double @test(double %x) {
entry:
        %addtmp = fadd double 2.000000e+00, 1.000000e+00
        %addtmp1 = fadd double %addtmp, %x
        ret double %addtmp1
}

如上所述,常量折叠是一种非常常见且非常重要的优化:以至于许多语言实现者在其AST表示中实现了常量折叠支持。

使用LLVM,你不需要AST中的这种支持,因为所有构建LLVM IR的调用都要经过LLVM IR构建器,当你调用它时,构建器本身会检查是否有固定的折叠机会。如果是,它只对常量进行折叠并返回该常量,而不创建指令。

在实践中,我们建议在生成这样的代码时始终使用IRBuilder。它的使用没有“语法开销”(你不必让你的编译器到处都有常量检查),它可以显著减少在某些情况下生成的LLVM IR的数量(特别是对于带有宏预处理器或使用大量常量的语言)。

另一方面,IRBuilder的局限性在于它在构建代码时将所有的分析都内联在一起。举一个稍微复杂一点的例子

ready> def test(x) (1+2+x)*(x+(1+2));
ready> Read function definition:
define double @test(double %x) {
entry:
        %addtmp = fadd double 3.000000e+00, %x
        %addtmp1 = fadd double %x, 3.000000e+00
        %multmp = fmul double %addtmp, %addtmp1
        ret double %multmp
}

在这种情况下,乘法的LHS和RHS是相同的值。我们真的希望看到它生成" tmp = x+3;Result = tmp*tmp; “而不是两次计算” x+3 "。

不幸的是,再多的本地分析也无法检测和纠正这一点。这需要两个转换:表达式的重新关联(使add的词法相同)和公共子表达式消除(CSE),以删除冗余的add指令。幸运的是,LLVM以“passes”的形式提供了您可以使用的广泛的优化。

LLVM优化passes

LLVM提供了许多优化passes,它们做许多不同类型的事情并具有不同的权衡。与其他系统不同的是,LLVM并没有坚持一组优化对所有语言和所有情况都适用的错误观念。LLVM允许编译器实现者对使用何种优化、以何种顺序以及在何种情况下进行优化做出完整的决策。

作为一个具体的例子,LLVM支持两种“整个模块”passes,它们可以遍历尽可能大的代码体(通常是整个文件,但如果在链接时运行,则可以是整个程序的很大一部分)。它还支持并包括“每个函数”passes,每次只操作一个函数,而不查看其他函数。有关传递及其运行方式的更多信息,请参阅如何编写传递文档和LLVM传递列表。

对于Kaleidoscope,我们目前正在动态地生成函数,当用户键入它们时,每次生成一个函数。在这种情况下,我们并不追求最终的优化体验,但我们也希望尽可能地抓住简单和快速的东西。因此,当用户输入函数时,我们将选择对每个函数运行一些优化。如果我们想要制作一个“静态Kaleidoscope编译器”,我们将完全使用我们现在拥有的代码,除了我们将推迟运行优化器,直到整个文件被解析。

除了函数和模块passes之间的区别之外,passes还可以分为转换和分析passes。变换传递改变IR,分析passes其他passes可以使用的计算信息。为了添加转换passes,它所依赖的所有分析passes都必须提前注册。

为了实现每个函数的优化,我们需要设置一个FunctionPassManager来保存和组织我们想要运行的LLVM优化。一旦我们有了这些,我们就可以添加一组优化来运行。我们需要为每个我们想要优化的模块创建一个新的FunctionPassManager,所以我们将添加到上一章创建的函数(InitializeModule())中:

void InitializeModuleAndManagers(void) {
  // Open a new context and module.
  TheContext = std::make_unique<LLVMContext>();
  TheModule = std::make_unique<Module>("KaleidoscopeJIT", *TheContext);
  TheModule->setDataLayout(TheJIT->getDataLayout());

  // Create a new builder for the module.
  Builder = std::make_unique<IRBuilder<>>(*TheContext);

  // Create new pass and analysis managers.
  TheFPM = std::make_unique<FunctionPassManager>();
  TheLAM = std::make_unique<LoopAnalysisManager>();
  TheFAM = std::make_unique<FunctionAnalysisManager>();
  TheCGAM = std::make_unique<CGSCCAnalysisManager>();
  TheMAM = std::make_unique<ModuleAnalysisManager>();
  ThePIC = std::make_unique<PassInstrumentationCallbacks>();
  TheSI = std::make_unique<StandardInstrumentations>(*TheContext,
                                                    /*DebugLogging*/ true);
  TheSI->registerCallbacks(*ThePIC, TheMAM.get());
  ...

在初始化全局模块TheModule和FunctionPassManager之后,我们需要初始化框架的其他部分。四个analysismanager允许我们添加在IR层次结构的四个级别上运行的分析通道。PassInstrumentationCallbacks和StandardInstrumentations是pass instrumentation框架所必需的,它允许开发人员自定义在pass之间发生的事情。

一旦这些管理器被设置,我们使用一系列的“addPass”调用来添加一堆LLVM转换传递:

// Add transform passes.
// Do simple "peephole" optimizations and bit-twiddling optzns.
TheFPM->addPass(InstCombinePass());
// Reassociate expressions.
TheFPM->addPass(ReassociatePass());
// Eliminate Common SubExpressions.
TheFPM->addPass(GVNPass());
// Simplify the control flow graph (deleting unreachable blocks, etc).
TheFPM->addPass(SimplifyCFGPass());

在本例中,我们选择添加四个优化步骤。我们在这里选择的是一组非常标准的“清理”优化,对各种各样的代码都很有用。我不会深入研究他们做了什么,但是,相信我,他们是一个很好的起点:)。

接下来,我们注册转换passes使用的分析passes。

  // Register analysis passes used in these transform passes.
  PassBuilder PB;
  PB.registerModuleAnalyses(*TheMAM);
  PB.registerFunctionAnalyses(*TheFAM);
  PB.crossRegisterProxies(*TheLAM, *TheFAM, *TheCGAM, *TheMAM);
}

一旦设置了PassManager,我们就需要使用它。返回客户端之前,我们可以在新创建的函数构造完成后(在FunctionAST::codegen()中)运行它,

if (Value *RetVal = Body->codegen()) {
  // Finish off the function.
  Builder.CreateRet(RetVal);

  // Validate the generated code, checking for consistency.
  verifyFunction(*TheFunction);

  // Optimize the function.
  TheFPM->run(*TheFunction, *TheFAM);

  return TheFunction;
}

如您所见,这非常简单。FunctionPassManager优化和更新LLVM Function*,希望 改进它的主体。有了这些,我们可以再次尝试上面的测试:

ready> def test(x) (1+2+x)*(x+(1+2));
ready> Read function definition:
define double @test(double %x) {
entry:
        %addtmp = fadd double %x, 3.000000e+00
        %multmp = fmul double %addtmp, %addtmp
        ret double %multmp
}

正如预期的那样,我们现在得到了经过精心优化的代码,每次执行该函数都保存了一个浮点添加指令。

LLVM提供了可以在某些情况下使用的各种各样的优化。一些关于各种Passes的文档,但不是很完整。另一个好的想法来源可以来自Clang开始运行的路径。“opt”工具允许您从命令行试验passes,这样您就可以看到它们是否做了什么。

现在我们有了来自前端的合理代码,让我们讨论一下如何执行它!

4.4. 新增一个JIT 编译器

LLVM IR中可用的代码可以有各种各样的工具应用于它。例如,您可以对其运行优化(如上所述),可以将其以文本或二进制形式转储,可以将代码编译为针对某些目标的汇编文件,或者可以对其进行JIT编译。LLVM IR表示的好处是,它是编译器许多不同部分之间的“通用货币”。

在本节中,我们将向解释器添加JIT编译器支持。我们想要的Kaleidoscope的基本思想是让用户像现在一样输入函数体,但是立即计算他们输入的顶级表达式。例如,如果他们输入“1 + 2;”,我们应该计算并打印出3。如果他们定义了一个函数,他们应该能够从命令行调用它。

为此,我们首先准备环境,为当前本机目标创建代码,并声明和初始化JIT。这是通过调用一些InitializeNativeTarget等函数并添加一个全局变量TheJIT,并在main中初始化它来完成的:

static std::unique_ptr<KaleidoscopeJIT> TheJIT;
...
int main() {
  InitializeNativeTarget();
  InitializeNativeTargetAsmPrinter();
  InitializeNativeTargetAsmParser();

  // Install standard binary operators.
  // 1 is lowest precedence.
  BinopPrecedence['<'] = 10;
  BinopPrecedence['+'] = 20;
  BinopPrecedence['-'] = 20;
  BinopPrecedence['*'] = 40; // highest.

  // Prime the first token.
  fprintf(stderr, "ready> ");
  getNextToken();

  TheJIT = std::make_unique<KaleidoscopeJIT>();

  // Run the main "interpreter loop" now.
  MainLoop();

  return 0;
}

我们还需要为JIT设置数据布局:

void InitializeModuleAndPassManager(void) {
  // Open a new context and module.
  TheContext = std::make_unique<LLVMContext>();
  TheModule = std::make_unique<Module>("my cool jit", TheContext);
  TheModule->setDataLayout(TheJIT->getDataLayout());

  // Create a new builder for the module.
  Builder = std::make_unique<IRBuilder<>>(*TheContext);

  // Create a new pass manager attached to it.
  TheFPM = std::make_unique<legacy::FunctionPassManager>(TheModule.get());
  ...

KaleidoscopeJIT类是专门为这些教程构建的简单JIT,可以在LLVM源代码中找到:LLVM -src/examples/Kaleidoscope/include/Kaleidoscope ejit .h。在后面的章节中,我们将看到它是如何工作的,并使用新功能扩展它,但现在我们将把它作为给定的。它的API非常简单:addModule向JIT添加了一个LLVM IR模块,使其函数可以执行(其内存由ResourceTracker管理);查找允许我们查找指向编译代码的指针。

我们可以使用这个简单的API,并将解析顶级表达式的代码更改为如下所示

static ExitOnError ExitOnErr;
...
static void HandleTopLevelExpression() {
  // Evaluate a top-level expression into an anonymous function.
  if (auto FnAST = ParseTopLevelExpr()) {
    if (FnAST->codegen()) {
      // Create a ResourceTracker to track JIT'd memory allocated to our
      // anonymous expression -- that way we can free it after executing.
      auto RT = TheJIT->getMainJITDylib().createResourceTracker();

      auto TSM = ThreadSafeModule(std::move(TheModule), std::move(TheContext));
      ExitOnErr(TheJIT->addModule(std::move(TSM), RT));
      InitializeModuleAndPassManager();

      // Search the JIT for the __anon_expr symbol.
      auto ExprSymbol = ExitOnErr(TheJIT->lookup("__anon_expr"));
      assert(ExprSymbol && "Function not found");

      // Get the symbol's address and cast it to the right type (takes no
      // arguments, returns a double) so we can call it as a native function.
      double (*FP)() = ExprSymbol.getAddress().toPtr<double (*)()>();
      fprintf(stderr, "Evaluated to %f\n", FP());

      // Delete the anonymous expression module from the JIT.
      ExitOnErr(RT->remove());
    }

如果解析和编码成功,下一步就是将包含顶级表达式的模块添加到JIT中。我们通过调用addModule来实现这一点,它触发模块中所有函数的代码生成,并接受一个ResourceTracker,该ResourceTracker可用于稍后从JIT中删除模块。一旦模块被添加到JIT中,它就不能再被修改了,所以我们还通过调用InitializeModuleAndPassManager()打开一个新模块来保存后续代码。

将模块添加到JIT之后,我们需要获得指向最终生成代码的指针。为此,我们调用JIT的查找方法,并传递顶级表达式函数的名称:__anon_expr。因为我们刚刚添加了这个函数,所以我们断言查找返回了一个结果。

接下来,我们通过在符号上调用getAddress()来获得__anon_expr函数的内存地址。回想一下,我们将顶级表达式编译成一个自包含的LLVM函数,该函数不接受任何参数,并返回计算后的双精度对象。因为LLVM JIT编译器匹配本机平台ABI,这意味着您可以将结果指针强制转换为该类型的函数指针,然后直接调用它。这意味着,JIT编译的代码和静态链接到应用程序中的本机机器码之间没有区别。

最后,由于我们不支持对顶级表达式的重新求值,所以我们在完成后将模块从JIT中删除,以释放相关的内存。但是,回想一下,我们之前创建的模块(通过InitializeModuleAndPassManager)仍然是open,等待添加新的代码。

ready> 4+5;
Read top-level expression:
define double @0() {
entry:
  ret double 9.000000e+00
}

Evaluated to 9.000000

看起来基本上是可行的。该函数的转储显示了我们为输入的每个顶级表达式合成的“总是返回double的无参数函数”。这演示了非常基本的功能,但我们还能做更多吗?

ready> def testfunc(x y) x + y*2;
Read function definition:
define double @testfunc(double %x, double %y) {
entry:
  %multmp = fmul double %y, 2.000000e+00
  %addtmp = fadd double %multmp, %x
  ret double %addtmp
}

ready> testfunc(4, 10);
Read top-level expression:
define double @1() {
entry:
  %calltmp = call double @testfunc(double 4.000000e+00, double 1.000000e+01)
  ret double %calltmp
}

Evaluated to 24.000000

ready> testfunc(5, 10);
ready> LLVM ERROR: Program used external function 'testfunc' which could not be resolved!

函数定义和调用也可以工作,但是在最后一行出现了一些非常错误的地方。这个呼叫看起来是有效的,那么发生了什么?正如您可能从API中猜到的那样,模块是JIT的分配单元,testfunc是包含匿名表达式的同一模块的一部分。当我们从JIT中删除该模块以为匿名表达式释放内存时,我们同时删除了testfunction的定义。然后,当我们试图第二次调用testfunc时,JIT再也找不到它了。

解决这个问题的最简单方法是将匿名表达式与其他函数定义放在单独的模块中。只要每个被调用的函数都有一个原型,并且在调用之前将其添加到JIT中,JIT将很高兴地解决跨模块边界的函数调用。通过将匿名表达式放在不同的模块中,我们可以在不影响其他函数的情况下删除它。

实际上,我们要更进一步,把每个函数放在自己的模块中。这样做允许我们利用KaleidoscopeJIT的一个有用属性,它将使我们的环境更像REPL函数可以多次添加到JIT中(不像每个函数都必须具有唯一定义的模块)。当你在KaleidoscopeJIT中查找一个符号时,它将总是返回最新的定义:

ready> def foo(x) x + 1;
Read function definition:
define double @foo(double %x) {
entry:
  %addtmp = fadd double %x, 1.000000e+00
  ret double %addtmp
}

ready> foo(2);
Evaluated to 3.000000

ready> def foo(x) x + 2;
define double @foo(double %x) {
entry:
  %addtmp = fadd double %x, 2.000000e+00
  ret double %addtmp
}

ready> foo(2);
Evaluated to 4.000000

为了允许每个函数在自己的模块中存在,我们需要一种方法来重新生成以前的函数声明到我们打开的每个新模块中:

static std::unique_ptr<KaleidoscopeJIT> TheJIT;
...
Function *getFunction(std::string Name) {
  // First, see if the function has already been added to the current module.
  if (auto *F = TheModule->getFunction(Name))
    return F;

  // If not, check whether we can codegen the declaration from some existing
  // prototype.
  auto FI = FunctionProtos.find(Name);
  if (FI != FunctionProtos.end())
    return FI->second->codegen();

  // If no existing prototype exists, return null.
  return nullptr;
}
...

Value *CallExprAST::codegen() {
  // Look up the name in the global module table.
  Function *CalleeF = getFunction(Callee);
...

Function *FunctionAST::codegen() {
 // Transfer ownership of the prototype to the FunctionProtos map,      but keep a reference to it for use below.
  auto &P = *Proto;
  FunctionProtos[Proto->getName()] = std::move(Proto);
  Function *TheFunction = getFunction(P.getName());
  if (!TheFunction)
    return nullptr;

为了实现这一点,我们将首先添加一个新的全局函数FunctionProtos,它包含每个函数的最新原型。我们还将添加一个方便的方法getFunction(),以替代TheModule->getFunction()的调用。我们可以非常方便的在module中搜索一个现有的函数声明,如果没有找到,就从FunctionProtos中生成一个新的声明,FunctionAST::codegen()中,我们需要首先更新FunctionProtos映射,然后调用getFunction()。这样,我们总是可以在当前模块中为任何先前声明的函数获得函数声明。 我们还需要更新HandleDefinition和HandleExtern:

static void HandleDefinition() {
  if (auto FnAST = ParseDefinition()) {
    if (auto *FnIR = FnAST->codegen()) {
      fprintf(stderr, "Read function definition:");
      FnIR->print(errs());
      fprintf(stderr, "\n");
      ExitOnErr(TheJIT->addModule(
          ThreadSafeModule(std::move(TheModule), std::move(TheContext))));
      InitializeModuleAndPassManager();
    }
  } else {
    // Skip token for error recovery.
     getNextToken();
  }
}

static void HandleExtern() {
  if (auto ProtoAST = ParseExtern()) {
    if (auto *FnIR = ProtoAST->codegen()) {
      fprintf(stderr, "Read extern: ");
      FnIR->print(errs());
      fprintf(stderr, "\n");
      FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

在HandleDefinition中,我们添加了两行来将新定义的函数传输到JIT并打开一个新模块。在HandleExtern中,我们只需要添加一行代码来将原型添加到FunctionProtos中。

自LLVM-9以来,不允许在单独的模块中重复符号。这意味着你不能在你的Kaleidoscope中重新定义函数,如下图所示。跳过这部分。原因是新的OrcV2 JIT api试图非常接近静态和动态链接器规则,包括拒绝重复的符号。要求符号名是唯一的,允许我们使用(唯一的)符号名作为跟踪键来支持符号的并发编译。

做了这些更改后,让我们再次尝试我们的REPL(这次我删除了匿名函数的转储,现在你应该明白了:):

ready> def foo(x) x + 1;
ready> foo(2);
Evaluated to 3.000000

ready> def foo(x) x + 2;
ready> foo(2);
Evaluated to 4.000000

它的工作原理!

即使是这么简单的代码,我们也获得了一些令人惊讶的强大功能——看看这个:

ready> extern sin(x);
Read extern:
declare double @sin(double)

ready> extern cos(x);
Read extern:
declare double @cos(double)

ready> sin(1.0);
Read top-level expression:
define double @2() {
entry:
  ret double 0x3FEAED548F090CEE
}

Evaluated to 0.841471

ready> def foo(x) sin(x)*sin(x) + cos(x)*cos(x);
Read function definition:
define double @foo(double %x) {
entry:
  %calltmp = call double @sin(double %x)
  %multmp = fmul double %calltmp, %calltmp
  %calltmp2 = call double @cos(double %x)
  %multmp4 = fmul double %calltmp2, %calltmp2
  %addtmp = fadd double %multmp, %multmp4
  ret double %addtmp
}

ready> foo(4.0);
Read top-level expression:
define double @3() {
entry:
  %calltmp = call double @foo(double 4.000000e+00)
  ret double %calltmp
}

Evaluated to 1.000000

JIT怎么知道sin和cos?答案出奇地简单:KaleidoscopeJIT有一个直接的符号解析规则,它使用它来查找在任何给定模块中都不可用的符号:首先,它搜索已经添加到JIT中的所有模块,从最近的到最旧的,以查找最新的定义。如果在JIT中没有找到定义,则返回到在Kaleidoscope进程本身上调用“ dlsym(”sin“) ”。由于“sin”是在JIT的地址空间中定义的,因此它只需在模块中修补调用以调用libm版本

在将来,我们将看到如何通过调整符号解析规则来启用各种有用的特性,从安全性(限制JIT代码可用的符号集)到基于符号名的动态代码生成,甚至是惰性编译。

符号解析规则的一个直接好处是,我们现在可以通过编写任意的c++代码来实现操作来扩展语言。例如,如果我们加上如下代码

#ifdef _WIN32
#define DLLEXPORT __declspec(dllexport)
#else
#define DLLEXPORT
#endif

/// putchard - putchar that takes a double and returns 0.
extern "C" DLLEXPORT double putchard(double X) {
  fputc((char)X, stderr);
  return 0;
}

注意,对于Windows,我们需要实际导出函数,因为动态符号加载器将使用GetProcAddress来查找符号。

现在我们可以使用如下命令向控制台生成简单的输出:“extern putchard(x);putchard(120);”,它在控制台上打印小写的“x”(120是“x”的ASCII码)。类似的代码可以用于实现文件I/O、控制台输入和Kaleidoscope中的许多其他功能。

这就完成了Kaleidoscope教程的JIT和优化器章节。此时,我们可以编译一个非图灵完备的编程语言,并以用户驱动的方式对其进行优化和JIT编译。接下来,我们将研究用控制流结构扩展该语言,并在此过程中解决一些有趣的LLVM IR问题。

完整代码llvm.org/docs/tutori…