WebAssembly 初探 – AOT & JIT

October 19, 2020

一、AOT

AOT (Ahead-of-Time) 可理解为“提前编译”,将高级语言的代码编译为二进制机器码(汇编指令),运行时直接执行,这里主要区分于“解释执行”,如我们常用的各类脚本语言,都需要对应的 Runtime 边解析边执行,另外如 Java、WebAssembly,其编译后的字节码也是通过 VM 来解释执行,而像C\C++、Rust 等语言,则可以看做 AOT 执行,编译后即为可直接执行的机器码。这里我们通过一个小实验来演示 AOT 的过程:

1、编译

我们以 C 作为实验语言,准备代码:

#include <stdio.h>

int importFunc(int n);

int aotTest(int m) {
    int a = importFunc(m);
    putchar('a');
    return a * a;
}

其中 importFunc 是外部依赖的函数,putchar 是系统库 libc 的函数,然后编译生成 Object 文件 ( x86_64 架构):

clang -v --target=x86_64 -emit-llvm -c -S aot.c
llc -march=x86-64 -filetype=obj --relocation-model=pic aot.ll

注意重定向模式选择 PIC (Position Independent Code),通过 objdump 查看生成的 aot.o 文件:

> objdump -d aot.o

aot.o:     file format elf64-x86-64

Disassembly of section .text:

0000000000000000 <aotTest>:
   0: 55                    push   %rbp
   1: 48 89 e5              mov    %rsp,%rbp
   4: 48 83 ec 10           sub    $0x10,%rsp
   8: 89 7d f8              mov    %edi,-0x8(%rbp)
   b: 8b 7d f8              mov    -0x8(%rbp),%edi
   e: e8 00 00 00 00        callq  13 <aotTest+0x13>
  13: 89 45 fc              mov    %eax,-0x4(%rbp)
  16: bf 61 00 00 00        mov    $0x61,%edi
  1b: e8 00 00 00 00        callq  20 <aotTest+0x20>
  20: 8b 45 fc              mov    -0x4(%rbp),%eax
  23: 0f af 45 fc           imul   -0x4(%rbp),%eax
  27: 48 83 c4 10           add    $0x10,%rsp
  2b: 5d                    pop    %rbp
  2c: c3                    retq

可以看出 aot.o 文件格式是 "elf64-x86-64",注意 0x0e 和 0x1b 处的 callq(0xe8) 指令,后面的32位地址都为 0,需要后续的链接步骤进行改写,继续使用 readelf 工具查看:

Relocation section '.rela.text' at offset 0x150 contains 2 entries:
  Offset          Info           Type           Sym. Value    Sym. Name + Addend
00000000000f  000400000004 R_X86_64_PLT32    0000000000000000 importFunc - 4
00000000001c  000500000004 R_X86_64_PLT32    0000000000000000 putchar - 4

‘.rela.text’ 中的两个重定位 Offset 即是 callq 指令32位地址的位置,0x00000000000f 处对应 importFunc,0x00000000001c 对应 putchar,Type 都为 ‘R_X86_64_PLT32’,偏移计算方式为 L + A – P,A 即为 ‘Addend’,这里都为-4 (下一条指令偏移),假设 importFunc 的地址偏移为 addr,我们需要填充 0x0f 位置的值为:addr + (-4) – 0x0f

2、加载

首先将 aot.o 文件加载到内存,这里使用 mmap 来实现:

FILE* file = fopen('aot.o', "rb");
if (file) {
    int fd = fileno(file);
    void * ptr = mmap(0, 2048, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE, fd, 0);
}

mmap 设置权限为读+写+执行,当然也可以先读+写,运行的时候读+执行(使用 mprotect 动态设置),由步骤1可知,.o 文件中存在需要重定位的项(importFunc、putchar),接下来实现重定位 (Relocation)

3、重定位

由于 importFunc、putchar 是外部导入的函数,我们构建类似 PLT 的结构来实现重定位,PLT 中使用 jmp 进行跳转,rax 中存放函数的真实地址(64位),然后 jmp rax

48 b8 xx xx xx xx xx xx xx xx   mov rax, 0x?
ff e0                           jmp rax

Relocation 流程如下图所示,text 段中写入 PLT 的地址 (RIP相对寻址),PLT 中 jmp 到函数的真实地址,即实现函数的重定位

我们 mmap 的内存比 aot.o 文件多,可将 PLT 写入后面空余的部分

static int32 pltOffset = 1300;
int importFunc(int n) {
    return n + 3;
}

// create plt
uint8 *p = (uint8*)ptr + pltOffset;
*p++ = 0x48;
*p++ = 0xB8;
*(uint64*)p = (uint64)(uintptr_t)&importFunc;
p += sizeof(uint64);
*p++ = 0xFF;
*p++ = 0xE0;

其中 ptr 为前面 mmap 分配的内存,接下来需要对 text 段的 callq 地址进行填充 (注意这里 0x40 为 .text 在文件中的的起始偏移):

int32 d = pltOffset + (-4) - 0x0f - 0x40;  // L + A - P
*((int32*)((char*)ptr + 0x40 + 0x0f)) = d;

两个外部导入函数 importFunc 和 putchar 都可用同样的方法完成重定位

4、执行

执行相对简单,将 .text 起始地址转换为函数指针,调用即可

typedef int (*AotFunc) (int);

AotFunc func = (AotFunc)((char*)ptr + 0x40);
int ret = func(5);
printf(",%d", ret);

输出:

a,64

64 = (5 + 3) * (5 + 3),执行结果正确

到此,我们完成了 x86_64 架构上代码编译、内存加载、符号重定位、执行的完整 AOT 流程,不同架构的差别主要在符号重定位上,不过 iOS 系统上由于代码签名验证机制,该方案不可行。

二、JIT

AOT 是提前编译,JIT (Just-in-Time) 则是运行期进行编译,通常是先解释执行,然后对热点函数进行编译,再将热点函数从解释执行切换到机器码执行,这里使用 LLVM C 接口演示一下代码的在线编译:

热点函数 sum 定义如下:

extern int extInt;

int importFunc(int n);

int sum(int a, int b) {
    return importFunc(a) + b + extInt;
}

首先初始化 LLVMModule

LLVMModuleRef mod = LLVMModuleCreateWithName("my_module");

然后定义外部导入函数 importFunc 和全局变量 extInt

LLVMTypeRef param_types_fn[] = { LLVMInt32Type() };
LLVMTypeRef fn_type = LLVMFunctionType(LLVMInt32Type(), param_types_fn, 1, false);
LLVMValueRef fn = LLVMAddFunction(mod, "importFunc", fn_type);

LLVMValueRef extInt = LLVMAddGlobal(mod, LLVMInt32Type(), "extInt");

接下来生成 sum 函数的实现代码

LLVMTypeRef param_types[] = { LLVMInt32Type(), LLVMInt32Type() };
LLVMTypeRef ret_type = LLVMFunctionType(LLVMInt32Type(), param_types, 2, 0);
LLVMValueRef sum = LLVMAddFunction(mod, "sum", ret_type);
LLVMBasicBlockRef entry = LLVMAppendBasicBlock(sum, "entry");

LLVMBuilderRef builder = LLVMCreateBuilder();
LLVMPositionBuilderAtEnd(builder, entry);

LLVMValueRef args[1];
args[0] = LLVMGetParam(sum, 0);
LLVMValueRef ret = LLVMBuildCall(builder, fn, args, 1, "call");
LLVMValueRef tmp = LLVMBuildAdd(builder, ret, LLVMGetParam(sum, 1), "tmp");
LLVMValueRef result = LLVMBuildAdd(builder, tmp, LLVMBuildLoad(builder, extInt, "extInt"), "result");
LLVMBuildRet(builder, result);

最后设定目标机器,生成机器码,注意指定 LLVMRelocMode 为 LLVMRelocPIC

LLVMInitializeAllTargets();
LLVMInitializeAllTargetMCs();
LLVMInitializeAllTargetInfos();
LLVMInitializeAllAsmPrinters();

LLVMTargetRef targetRef;

char triple[] = "x86_64-linux-gnu";
char cpu[] = "";
char** errPtrTriple;
LLVMBool resTriple = LLVMGetTargetFromTriple(triple, &targetRef, errPtrTriple);
LLVMTargetMachineRef targetMachineRef = LLVMCreateTargetMachine(targetRef, 
                                                                triple, 
                                                                cpu, 
                                                                "", 
                                                                LLVMCodeGenLevelNone, 
                                                                LLVMRelocPIC, 
                                                                LLVMCodeModelDefault);
char** errPtrFileObj;
LLVMBool resFileObj = LLVMTargetMachineEmitToFile(targetMachineRef, 
                                                  mod, 
                                                  "sum_llvm.o", 
                                                  LLVMObjectFile, 
                                                  errPtrFileObj);

使用 objdump 和 readelf 查看生成的 sum_llvm.o 文件:

sum_llvm.o:     file format elf64-x86-64

Disassembly of section .text:

0000000000000000 <sum>:
   0:   50                      push   %rax
   1:   89 74 24 04             mov    %esi,0x4(%rsp)
   5:   e8 00 00 00 00          callq  a <sum+0xa>
   a:   48 8b 0d 00 00 00 00    mov    0x0(%rip),%rcx        # 11 <sum+0x11>
  11:   8b 54 24 04             mov    0x4(%rsp),%edx
  15:   01 d0                   add    %edx,%eax
  17:   03 01                   add    (%rcx),%eax
  19:   59                      pop    %rcx
  1a:   c3                      retq
Relocation section '.rela.text' at offset 0x120 contains 2 entries:
  Offset          Info           Type           Sym. Value    Sym. Name + Addend
000000000006  000400000004 R_X86_64_PLT32    0000000000000000 importFunc - 4
00000000000d  000300000009 R_X86_64_GOTPCREL 0000000000000000 extInt - 4

可知 LLVM 最终产生的其实就是 elf 文件,当然也可以使用 LLVMTargetMachineEmitToMemoryBuffer 直接写入内存,然后按前文 AOT 的方式完成重定位、执行。

综上,得益于 LLVM 这个大杀器,我们可以在运行期将代码编译为机器码,然后加载执行,实现 JIT 过程。这也是多个 WebAssembly 独立 Runtime 的 JIT 实现方式,如 WAMR\WAVM 等。

参考

https://www.gabriel.urdhr.fr/2015/09/28/elf-file-format/
https://github.com/maiquynhtruong/compilers/wiki/LLVM-Resources