LLVM 新一代 JIT API：ORC

2023-10-22 13:16:40

最近看到 ORC(On Request Compilation) 在增加 MachO 平台的 OC 和 Swift 语言支持，这是 MachO JIT(Just In Time) 相关的进展。本文将探索这个 LLVM 新一代的 JIT APIs，即 ORC，其 ORC JIT Weekly 现在还一直处于更新状态。

1、JIT 解释

以防语境不一致，解释下 JIT(Just In Time) 这个术语。

Whenever a program, while running, creates and runs some new executable code which was not part of the program when it was stored on disk, it’s a JIT.

JIT 即一个程序在运行时，创建并运行了一些新的可执行代码，而这些代码并非是程序的原有部分。

JIT 也被称之为懒编译(late/lazy compilation)。

其实包含两个概念，一个是动态生成代码，再一个是动态运行代码。

1.1、AOT

这个概念是相对于 AOT(Ahead Of Time) 而言的，AOT 会在执行之前把代码(文本、字节码等)编译为本地代码(机器指令)，运行时就执行这些编译好的本地代码。

1.2、JIT 的好处

AOT 可以预先生成执行效果较好(优化过)的本地代码，但由于是静态编译，所以无法准确预测代码在运行时的行为，不能达到性能极致。

相比而言，JIT 结合了解释器和 AOT 的优点，在运行时收集各种信息和指标，既可以快速启动，又可以生成更加优化的代码，所以能在总体上达到更佳的运行效果，这也是解释器 + JIT 称为目前主流执行方式的原因之一。

生成更高效的代码是需要花费很多时间的，为了在快速启动和高度优化之间取得平衡，很多 JIT/AOT 实现（比如 Java 虚拟机和 Firefox 浏览器）使用了分层（Tiered）编译技术。分层编译一般分为两层，第一层（tier1，叫法因实现而异）编译器可以较快地生成本地代码(简陋的优化编译，如 -O0)并执行；而第二层（tier2）编译器在后台线程中生成执行效果更好的代码，并替换 tier 生成的代码。

还有另一个好处就是，可以动态下发代码，实现程序运行时的热重载和热修复。热重载比如在调试当中，可以边修改代码边实时调试，而不需要停止调试再重新编译。热修复比如在线上发现了某个函数存在 bug，正常是要修改好后，重新发一个编译版本，但 JIT 就可以下发代码直接替换掉函数。

1.3、简陋 JIT 的简单实现

使用 mmap 函数创建可读可写可执行(一般操作系统可能会有限制)的内存，在这段内存写入要执行的机器指令码(可以内置编译器生成机器码)，把这段内存的首地址作为函数指针，之后进行调用。

//分配内存
void* createSpace(size_t size) {
    void* ptr = mmap(0, size,
            PROT_READ | PROT_WRITE | PROT_EXEC,
            MAP_PRIVATE | MAP_ANON,
            -1, 0);   
    return ptr;
}

long add(long num) { return num + 2; }

//内存中创建函数
void copyCodeToMem(unsigned char* addr) {
    // 上面 add 函数的机器码
    unsigned char macCode[] = {
        0x55,
        0x48,0x89,0xe5,
        0x48,0x89,0x7d,0xf8,
        0x48,0x8b,0x45,0xf8,
        0x48, 0x83, 0xc0, 0x02,
        0x5d,
        0xc3 
    };
    memcpy(addr, macCode, sizeof(macCode));
}

int main(int argc, char** argv) {                                                                                              
    const size_t SIZE = 1024;
    typedef long (*demo)(long);
    void* addr = createSpace(SIZE);
    copyCodeToMem(addr);
    demo d1 = addr;
    long result = d1(1);
    printf("result = %ld\n", result); 
    return 0;
}

2、LLVM JIT 设计与实现

2.1、设计需要

LLVM 的 JIT 的设计要满足一些用户及其需要：

  • Kaleidoscope(设计自己的编译前端)：简单、安全
  • LLDB(调试器)：交叉编译(Cross-target compilation)
  • High Performance JITs(高性能 JIT)：可以自定义配置优化和生成代码的操作
  • Interpreters and REPLs(解释器)：懒编译(Lazy compilation)、与静态编译效果一致(static compile)

LLVM 有三个 JIT 的实现：

LLVM 新一代 JIT API：ORC

2.2、LLVM JIT engine

LLVM JIT 编译器是基于函数粒度的，因为它可以一次只编译一个函数。(理论上可以进一步提高粒度，为 trace，即函数的某条特定执行路径，但还待研究)

JIT engine 会在运行时编译执行 LLVM IR 函数，在进行编译，它会使用 LLVM code generator(代码生成器) 去生成指定平台的二进制指令，然后会返回编译好的函数指针，这样就可以通过函数指针的方式来调用函数了。

LLVM JIT 系统使用 ExecutionEngine 类来提供支持，用来组织整个程序的执行、分析下一个需要被执行的程序段、选择对应的操作来执行，它会把代码生成到内存当中，但是否要执行取决于用户(开发者是否要调用)。

LLVM JIT engine 有以下特点：

  • 懒编译(lazy compilation)：只在调用时才进行函数的编译。如果关闭该特性，则获得函数指针时就会马上进行编译。
  • 外部全局变量的编译：包括对当前 LLVM 模块(module)之外的实例，进行符号解析和内存分配
  • 通过 dlsym 查找和解析外部符号：这是在运行时进行动态共享对象(dynamic shared object，DSO)加载一样的过程

LLVM 里面有两套 JIT execution engine(执行引擎)的实现：llvm::JIT 类和 llvm::MCJIT 类。一个 ExecutionEngine 对象是用 ExecutionEngine::EngineBuilder() 函数和 IR Module 参数来进行实例化的。然后 ExecutionEngine::create() 会创建一个 JIT 或者 MCJIT 引擎实例。(所以基于 LLVM 的实现思路的热修复一般是采用 IR(bitcode) 下发的)

把二进制指令写进内存里是由 ExecutionManager 类完成的，它会把函数指针给用户。内存管理的任务包括内存分配、释放、提供内存空间给库加载、内存权限处理。JIT 和 MCJIT 都实现了自己的内存管理类，继承于 RTDyldMemoryManager 基类。

2.3、llvm::JIT

JIT 类以及其框架都是旧版的引擎，用了不同部分的 LLVM 代码生成器，在 LLVM 3.5 之后被移除，它是平台无关(target-independent)的，每种平台都需要实现自己的二进制指令生成操作。

JIT 类使用 JITCodeEmitter 类来生成二进制指令，是 MachineCodeEmitter 的子类，但和新的 MC(Machine Code)框架毫无关系，只能支持少量的平台，且很多平台特性不可用。

JIT 类使用 JITMemoryManager 进行内存管理，使用 JITResolver 实例来跟踪和解析所有还没被编译好的函数的调用点(call sites)，这对于实现懒编译非常重要。它提供了很多函数，例如，用 allocateGlobal() 函数来为一个全局变量分配内存，使用 startFunctionBody() 函数来创建 JIT 的调用(为生成指令分配内存)。

每种平台需要实现 machine function(机器码级别的函数) 的 pass，名为 <Target>CodeEmitter，用来把指令编码成 blobs(binary large object)，使用 JITCodeEmitter 写入内存之中。比如，MipsCodeEmitter，会遍历函数所有的代码块(basic blocks)，然后调用 emitInstruction() 来处理每条机器指令(machine instruction)。

LLVM 新一代 JIT API：ORC

// (...) 
MCE.startFunction(MF); 
for (MachineFunction::iterator MBB = MF.begin(), E = MF.end(); MBB != E; ++MBB){ 
    MCE.StartMachineBasicBlock(MBB); 
    for (MachineBasicBlock::instr_iterator I = MBB->instr_begin(), E = MBB->instr_end(); I != E;)        
         emitInstruction(*I++, *MBB); 
} 
// (...)
void MipsCodeEmitter::emitInstruction(MachineBasicBlock::instr_ iterator MI, MachineBasicBlock &MBB) {
    // ...
    MCE.processDebugLoc(MI->getDebugLoc(), true);   
    emitWord(getBinaryCodeForInstr(*MI));  // 可以阅读 llvm TableGen 了解生成指令
    ++NumEmitted;  // Keep track of the # of mi's emitted   
    // ...
}

LLVM 新一代 JIT API：ORC

如果使用懒编译的方式，会先生成一个 stub (桩)函数指针来返回，在进行编译时，再把指针修正(patch)为(jump/call)真实函数地址。在下次调用时，会直接调到真实函数地址。

2.4、llvm::MCJIT

MCJIT(Machine Code JIT) 是 LLVM 中比较新的 JIT 实现。

MC 为指令提供了统一的表达(uniform representation)，而且它是一个框架，可以共享于汇编器、反汇编器、汇编输出和 MCJIT。所以好处之一就是只需要指明一次指令的编码方式就够了，而且当你编写了 LLVM 后端生成指令，那么你也具有了 JIT 的功能。

MCJIT 同样使用ExecutionEngine 来进行实例创建，构造器(constructor) 同样用 llvm::Module 对象作为入参。

LLVM 新一代 JIT API：ORC

MCJIT 设计了一些关于 LLVM module 的实例的编译状态：

  • Added - 还未被编译，但已加入到执行引擎当中。允许模块暴露函数定义给其他模块，然后进行延迟编译，直到需要时。
  • Loaded - 该模块处于 JIT 编译，但未准备好被执行。重定位未完成，且需要分配合适权限的内存页。用户可以在内存中 remap JIT 编译的函数而不需要重新编译。
  • Finalized - 模块包含有准备被执行的函数。不能被 remap，因为重定位已经完成了。

这个状态的设计，使得 MCJIT 要获取符号地址时，必须整个模块已经处于 Finalized 状态。MCJIT::finalizeObject() 会调用 generateCodeForModule() 来生成已加载的模块，然后所有的模块会通过 finalizeLoadedModules() 函数被 finalized。

generateCodeForModule() 做了以下几件事情：

  • 创建 ObjectBuffer 实例来持有 Module 对象，如果 Module 已经被加载(编译过了)，那么会使用 ObjectCache 接口来查找，避免重编。
  • 如果没有 cache，那么执行 MC 代码生成 MCJIT::emitObject()，会返回 ObjectBufferStream 对象。
  • RuntimeDyld 动态链接器会加载结果 ObjectBuffer 对象(根据文件格式调用对应平台的)，然后通过 RuntimeDyld::loadObject() 来创建符号表，返回 ObjectImage 对象。
  • 标记模块为已加载 Loaded。

LLVM 新一代 JIT API：ORC

RuntimeDyld 动态链接器是在 Module 进行 finalization (符号解析、注册异常处理)时被使用。MCJIT 也会使用 RuntimeDyld 通过 RuntimeDyld::getSymbolLoadAddress() 函数来查找符号地址。

2.5、LLVM JIT 编译工具

2.5.1、lli

lli (解释工具，interpreter tool)，实现了 LLVM bitcode(IR) 的解释器，以及使用 LLVM 执行引擎实现的 JIT 编译器。

$ clang -emit-llvm -c sum-main.c -o sum-main.bc
$ lli sum-main.bc
$ lli -use-mcjit sum-main.bc
$ lli -force-interpreter sum-main.bc

2.5.2、llvm-rtdyld

llvm-rtdyld 工具可以用来测试 MCJIT 对象加载和链接框架。可以从磁盘读取二进制目标文件，然后执行指定的函数。它不会进行 JIT 编译和执行，但可以让你测试和运行目标文件。

$ clang -g -c add.c -o add.o 
$ llvm-rtdyld -printline add.o 
Function: _add, Size = 20   
Line info @ 0: add.c, line:2   
Line info @ 10: add.c, line:3   
Line info @ 20: add.c, line:3

这个工具实际上就是把二进制目标文件读取进了 ObjectBuffer 对象，然后生成 ObjectImage 实例，通过 RuntimeDyld::resolveRelocations() 进行重定位，最后函数入口是通过 getSymbolAddress() 拿到以及调用。

3、ORC(On Request Compilation)

ORC 为构造 JIT 编译器提供了模块化的 API 接口。提供有以下特性：

  • JIT-linking：提供了 APIs 在运行时链接重定位文件到目标进程。
  • LLVM IR compilation：提供现成的组件，来把 LLVM IR 添加到 JIT 进程当中。
  • Eager and lazy compilation：默认 ORC 会在查找 JIT session 对象(ExecutionSession)时编译符号。当然 ORC 也提供了懒编译的选项。
  • Support for Custom Compilers and Program Representations：ORC 可以运行用户通过 JIT session 定义符号提供的自定义编译器。
  • Concurrent JIT‘d code and Concurrent compilation：JIT 代码可能被多线程执行，可能创建一个新的线程，可能并发地重新进入 ORC。内置的依赖跟踪可以保证，ORC 不会释放掉 JIT 代码或数据的指针，直到所有的依赖都 JIT 结束，以及是可以安全地调用为止。
  • Removable Code：为 JIT 程序表达(program representation)提供资源。
  • Orthogonality and Composability：以上的特性都可以被独立或组合地使用。

从 LLVM 7.0 开始，ORC 就专注于支持并发的 JIT 编译，这种并发的能力被整合到 ORCv2 版本当中。而传统(legacy)版本整合为 ORCv1，在 LLVM 12.0 会被移除。

LLVM 新一代 JIT API：ORC

通过模块化的方式，编译层和链接层可以被单独地进行测试，而且可以不通过回调就能观察到事件的变化(通过增加通知层等)。

3.1、LLJIT & LLLazyJIT

ORC 提供了两个基本的 JIT 类，用来集成 ORC 组件来创建 JIT，以及如何替换掉早期的 LLVM JIT(比如 MCJIT)。

LLJIT 类使用 IRCompileLayer 和 RTDyldObjectLinkingLayer 来支持 LLVM IR 的编译，以及重定向文件的链接。所有操作都是在符号查找时立即进行的。在大部分情况下，LLJIT 都被用来替换掉 MCJIT。

LLLazyJIT 是 LLJIT 的扩展，添加了 CompileOnDemandLayer 来支持 LLVM IR 的懒编译。当 LLVM IR 模块被 addLazyIRModule 添加，函数的实现体在第一次调用时，才会被编译。

// Try to detect the host arch and construct an LLJIT instance.
auto JIT = LLJITBuilder().create();

// If we could not construct an instance, return an error.
if (!JIT)
  return JIT.takeError();

// Add the module.
if (auto Err = JIT->addIRModule(TheadSafeModule(std::move(M), Ctx)))
  return Err;

// Look up the JIT'd code entry point.
auto EntrySym = JIT->lookup("entry");
if (!EntrySym)
  return EntrySym.takeError();

// Cast the entry point address to a function pointer.
auto *Entry = (void(*)())EntrySym.getAddress();

// Call into JIT'd code.
Entry();

// 懒编译版本
// Build an LLLazyJIT instance that uses four worker threads for compilation,
// and jumps to a specific error handler (rather than null) on lazy compile
// failures.

void handleLazyCompileFailure() {
  // JIT'd code will jump here if lazy compilation fails, giving us an
  // opportunity to exit or throw an exception into JIT'd code.
  throw JITFailed();
}

auto JIT = LLLazyJITBuilder()
             .setNumCompileThreads(4)
             .setLazyCompileFailureAddr(
                 toJITTargetAddress(&handleLazyCompileFailure))
             .create();

LLVM 新一代 JIT API：ORC

上图为懒编译流程图。

3.2、设计

ORC 的 JIT 模型目标是——模拟静态和动态链接器所使用的链接和符号解析的规则。这可以让 ORC 对任意的 LLVM IR 进行 JIT 操作。

看下 ORC 是如何运作的，在命令行下的构建程序是这样的：

$ clang++ -shared -o libA.dylib a1.cpp a2.cpp
$ clang++ -shared -o libB.dylib b1.cpp b2.cpp
$ clang++ -o myapp myapp.cpp -L. -lA -lB
$ ./myapp

而在 ORC 当中，会转换成对应的 API 调用：

// JIT 程序，提供 JIT 内容：JITDylib、错误报告机制、符号定义查询操作
ExecutionSession ES;
// 这两个 layer 是编译器的一层封装，可以将未编译的中间码加到 JITDylib 中
RTDyldObjectLinkingLayer ObjLinkingLayer(
    ES, []() { return std::make_unique<SectionMemoryManager>(); });
CXXCompileLayer CXXLayer(ES, ObjLinkingLayer);

// Create JITDylib "A" and add code to it using the CXX layer. 符号表
auto &LibA = ES.createJITDylib("A");
CXXLayer.add(LibA, MemoryBuffer::getFile("a1.cpp"));
CXXLayer.add(LibA, MemoryBuffer::getFile("a2.cpp"));

// Create JITDylib "B" and add code to it using the CXX layer.
auto &LibB = ES.createJITDylib("B");
CXXLayer.add(LibB, MemoryBuffer::getFile("b1.cpp"));
CXXLayer.add(LibB, MemoryBuffer::getFile("b2.cpp"));

// Create and specify the search order for the main JITDylib. This is
// equivalent to a "links against" relationship in a command-line link.
auto &MainJD = ES.createJITDylib("main");
MainJD.addToLinkOrder(&LibA);
MainJD.addToLinkOrder(&LibB);
CXXLayer.add(MainJD, MemoryBuffer::getFile("main.cpp"));

// Look up the JIT'd main, cast it to a function pointer, then call it.
auto MainSym = ExitOnErr(ES.lookup({&MainJD}, "main"));
auto *Main = (int(*)(int, char*[]))MainSym.getAddress();

int Result = Main(...);

这里的操作都依赖于 CXXCompilingLayer 的实现。ORC 还可以生成报错信息。而 llvm::orc::JITDylib 提供了符号表，支持异步的符号查询。

参考

How to JIT - an introduction：https://eli.thegreenplace.net/2013/11/05/how-to-jit-an-introduction

JIT原理简单介绍：https://segmentfault.com/a/1190000040256281

《WebAssembly 原理与核心技术》笔记：https://calssion.netlify.app/2021/08/07/assembly/wasm/

Add initial Objective-C and Swift support to MachOPlatform：https://reviews.llvm.org/rGcdcc35476833

Kaleidoscope-My First Language Frontend with LLVM Tutorial：https://llvm.org/docs/tutorial/MyFirstLanguageFrontend/index.html

《Getting Started with LLVM Core Libraries》第 7 章 The Just-in-Time Compiler

MCJIT Design and Implementation：https://llvm.org/docs/MCJITDesignAndImplementation.html

ORC Design and Implementation：https://llvm.org/docs/ORCv2.html

Building a JIT: Starting out with KaleidoscopeJIT：https://llvm.org/docs/tutorial/BuildingAJIT1.html

LLVM Dev Meeting 2016 - ORC：https://llvm.org/devmtg/2016-11

上一篇：【论文笔记】：Mutual Supervision for Dense Object Detection


下一篇：【翻译】Focal Loss for Dense Object Detection(RetinaNet)