对Kaleidoscope解释器实现JIT(及时)编译的方法

243 阅读5分钟

JIT(Just in Time)编译涉及将字节码转化为机器可执行指令的字节码。在这篇文章中,我们将对我们的Kaleidoscope解释器实现一个JIT。

目录:

  1. 介绍
  2. JIT编译器
  3. 总结

先决条件

  1. LLVM编译器的优化

介绍

JIT(Just-In-Time)编译器是一个将字节码转换为可由目标机器执行的指令的编译器。JIT编译器主要用于我们想在运行时改善或优化二进制代码的性能的情况。例如,Java JIT编译器在运行时提高了Java程序的性能。

一个具有JIT的系统在执行过程中不断分析代码,并识别出从编译或重新编译中获得的速度提升超过编译代码的开销的区块。

JIT编译器

我们从先决条件的文章中获得的LLVM IR是许多不同的编译器部分之间的共同货币,例如,我们可以转换它的文本格式、二进制格式、运行优化、JIT编译等等。

JIT编译器的基本思想是,用户输入一个函数体,并立即评估输入的顶层表达式。例如,如果用户输入了1+2,编译器会评估并返回3。如果一个函数已经被定义,它应该可以从REPL中调用。

为此,我们创建一个环境,为当前的本地目标创建代码并初始化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;
}

KaleidoscopeJIT是为本教程建立的一个JIT,它LLVM的源代码中。它的API非常简单,addModule将一个LLVM IR模块添加到JIT中,使其功能可供执行,removeModule删除一个模块,这也释放了与模块中代码相关的任何内存,findSymbol使我们能够循环指向编译后的代码。

然后我们为JIT设置数据布局,为此我们在InitializeModuleAndPassManager函数中添加以下一行:

TheModule->setDataLayout(TheJIT->getTargetMachine().createDataLayout()); // JIT

我们使用API并改变解析顶层表达式的代码。我们的handleTopLevelExpression现在看起来像下面这样:

static void handleTopLevelExpression()
{
    if (auto FnAST = ParseTopLevelExpr()) // evaluate top-level expression into anonymous function
    {
        if (FnAST->codegen())
        {
            auto RT = TheJIT->getMainJITDylib().createResourceTracker(); // create resource tracker, tracks JIT allocated memory

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

            auto ExprSymbol = ExitOnErr(TheJIT->lookup("__anon_expr")); // search JIT for __anon_ expression

            double (*FP)() = (double (*)())(intptr_t)ExprSymbol.getAddress(); // get symbol address
            fprintf(stderr, "Evaluated to %f\n", FP());

            ExitOnErr(RT->remove()); // remove anonymous expression module from JIT
        }
    }
    else
    {
        getNextToken(); // skip token, error recovery
    }
}

如果解析和代码生成成功,我们将带有顶层表达式的模块添加到JIT中。为此,我们调用addModule,它触发了模块中所有函数的代码生成,并返回一个句柄,用于以后从JIT中删除模块。

当模块被添加到JIT中时,它不能被修改,因此我们打开一个新的模块来保存后续代码。为此我们调用InitializeModuleAndPassManager()

当模块被添加到JIT时,我们通过调用JITfindSymbol方法并传递顶层表达式的名称来获得最终生成代码的指针。由于我们刚刚添加了这个函数,findSymol被作为结果返回。

然后我们通过调用符号上的getAddress来获得顶层表达式名称的内存地址--__anon_expr
我们将顶层表达式编译成一个独立的LLVM函数,它不需要任何参数,并返回一个双数。

由于我们不支持对顶层表达式的重新评估,所以我们删除该模块,以便释放内存。

我们也可以让函数住在自己的模块中,为此我们需要将之前的函数声明重新生成到我们打开的每个新模块中:

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()的调用。它搜索现有的函数声明,如果没有找到,就从FunctionProtos生成一个新的声明。

CallExprAST::codegen()中,我们替换对TheModule->getFunction()的调用,在FunctionAST::codegen()中,我们更新FunctionProtos映射,然后调用getFunction(),我们就完成了。

现在要更新辅助函数handleDefinitionhandleExtern

static void HandleDefinition() {
  if (auto FnAST = ParseDefinition()) {
    if (auto *FnIR = FnAST->codegen()) {
      fprintf(stderr, "Read function definition:");
      FnIR->print(errs());
      fprintf(stderr, "\n");
      TheJIT->addModule(std::move(TheModule));
      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

总结

JIT编译涉及到执行代码,据此我们在执行过程中--在运行时--对代码进行编译。