译作:窥斑见豹--Peeking inside luaJIT

原文:Peeking inside LuaJIT

luajit是lua的一款高性能虚拟机。我最近开始捣鼓它,我把目前为止所发现的记录到这篇博文里。
luajit是一个tracing jit,而不是method JIT。它的工作模式并不是发现并优化整个热点方法(hot method),而是发现并优化执行过程中的热点轨迹(hot traces)或者路径。luaJIT也有一个相当有意思的翻译器,但我不会主动提到它,除非涉及到JIT。
翻译器的实现代码主要分布到vm_<arch>.dasc文件中,其中包含翻译器(和其它汇编)桩函数的规格说明(以类汇编语言描述)。例如,vm_x86.dasc包含如下函数:

// Generate assembly code for a specific Lua bytecode.
static void build_ins(BuildCtx *ctx, BCOp op, int defop) {
  switch (op) {
  case BC_HHGTTG:   // hypothetical bytecode
    | mov eax, 42  // hypothetical implementation
    break;
  }
}

vm_86.dasc首先由dynasm.lua解析,并转换成(lowered by)C程序。以|开头的行被映射成对dasm_put的调用,通过这种方式,生成的C函数自动生成(emit)汇编操作码(assembly opcodes)。这个新生成的C程序跟buildvm.c一起链接,然后在编译时生成所需的汇编桩(assembly stubs)。例如,以BC_HHGTTG(转换后)调用build_ins,会生成把42写到eax的汇编操作码(不管这个指令的语义是啥)。这个汇编指令写入到lj_vm.s,并且与luaJIT其余部分一起链接。

探测热点轨迹

现在,有了一点背景知识后,让我们看看lua实际上是如何找到要编译的轨迹(traces)的。LuaJIT用一张hash表,维护了相关指令(跳转和调用)的热度,但是奇怪的是,它没有解决冲突。由于热度是启发式的,而不密集(即是,不太可能程序中所有跳转都是热点),这个方法实际上工作得很好。下面的宏用于访问热点计数(这个宏是C语言中一个特定指令)。看看这个宏的实现,理解应该会更清楚些:

// gg is the global state and pc points to the bytecode
// instruction whose hotcount we want.
#define hotcount_get(gg, pc) (gg)->hotcount[(u32ptr(pc)>>2) & (HOTCOUNT_SIZE-1)]

翻译器执行跳转或者调用时,就会更新并检查这些热点计数器。这是通过vm_<arch>.dasc中的hotloop和hotcall宏实现的(我们将只关注X86架构,所以<arch>是x86):

|.macro hotloop, reg
|  mov reg, PC
|  shr reg, 1
|  and reg, HOTCOUNT_PCMASK
|  sub word [DISPATCH+reg+GG_DISP2HOT], HOTCOUNT_LOOP
|  jb ->vm_hotloop
|.endmacro

然后,在各个地方会调用hotloop

case BC_FORL:
  |.if JIT
  |  hotloop RB
  |.endif

dynasm将这个宏调用以简单明了的方式转换,BC_FORL除去HOTCOUT_LOOP的计数,一旦计数小于0,则跳转到vm_hotloop

|->vm_hotloop:            // Hot loop counter underflow.
|.if JIT
|  mov LFUNC:RB, [BASE-8]     // Same as curr_topL(L).
|  mov RB, LFUNC:RB->pc
|  movzx RD, byte [RB+PC2PROTO(framesize)]
|  lea RD, [BASE+RD*8]
|  mov L:RB, SAVE_L
|  mov L:RB->base, BASE
|  mov L:RB->top, RD
|  mov FCARG2, PC
|  lea FCARG1, [DISPATCH+GG_DISP2J]
|  mov aword [DISPATCH+DISPATCH_J(L)], L:RBa
|  mov SAVE_PC, PC
|  call extern lj_trace_hot@8     // (jit_State *J, const BCIns *pc)
|  jmp <3
|.endif

dynasm通过.type指示符来识别C语言的结构体和其内部字段的偏移位置。在vm_hotloop的实现体内部,LFUNC:RB告诉dynasm把RB的值当作一个LFUNC来处理。RB和LFUNC在前面定义了:
译注: LFUNC其实是GCfuncL,这是luajit的lua函数类型。与之对应lua还有函数原型。两者的差别类似C++的类与对象的关系,不太像js的原型跟对象的关系。详细分析可看这篇

|.define RB,      ebp
// ...
|.type LFUNC,     GCfuncL

当然,在LFUC:RB中把RB当作一个LFUNC并不会做什么事,它只是标注。但是,在下一条指令中,这个标注让我们可以在代码中写出LFUNC:RB->pc,而dynasm会自动找到pc正确的偏移位置。让我们逐行看一遍vm_hotloop,注意BASE指向lua求值栈的栈顶,RB和RD则是寄存器,FCARG1和FCARG2也是寄存器,它俩是lua调用C函数时调用规范中的前两个函数参数,SAVE_L和SAVE_PC是栈上的槽位。
首先,我们从栈顶弹出GCFuncL,并取出pc,其指向这个闭包字节码的开头:

|  mov LFUNC:RB, [BASE-8]
|  mov RB, LFUNC:RB->pc

luaJIT遵守一个通用惯例,函数字面量或者原型跟函数闭包分开保存。这样,VM可以在相同函数原型的两个闭包之间共享字节码,从而节省空间。例如,在V8中,你用SharedFunctionInfo来保存一个函数字面量的特有信息,用 Function来表示实际可执行的闭包。
在luaJIT中,函数字面量用GCproto来表示。在内存中,一个GCproto对象后面跟着函数字面量(所有闭包共享,用GCfuncL来表示)的字节码。因此,给定一个GCfuncL,我们能够将指向字节码数组开头的指针减去sizeof(GCproto)来得到对应的GCproto。PC2PROTO用这个技术来访问GCproto结构体的framesize字段,用它来计算栈上第一个可用槽。

|  movzx RD, byte [RB+PC2PROTO(framesize)]
|  lea RD, [BASE+RD*8]

然后,我们填好lua_state的各个字段(在lj_obj.h中定义)

|  mov L:RB, SAVE_L
|  mov L:RB->base, BASE
|  mov L:RB->top, RD

设置参数寄存器,保存一点东西到栈槽中,然后调用lj_trace_hot

|  mov FCARG2, PC
|  lea FCARG1, [DISPATCH+GG_DISP2J]
|  mov aword [DISPATCH+DISPATCH_J(L)], L:RBa
|  mov SAVE_PC, PC
|  call extern lj_trace_hot@8
|  jmp <3

luaJIT进入记录模式。

记录轨迹

轨迹是执行一段特定代码路径的字节码线性序列。轨迹由翻译器协助记录(记录始于lj_trace_hot调用之时)。翻译器使用一个指针数组,其元素指向汇编桩,并以汇编桩实现的字节码指令为数组索引--原则上,翻译一个lua程序则是把字节码一个接一个解码并跳转到对应的汇编桩上。把这个调度数组中各个元素都指向lj_vm_record(dynasm加上了lj_前缀,在vm_x86.dasc中符号其实是vm_record)从而记录轨迹。一个简化的vm_record如下所示:

|->vm_record:
|  mov L:RB, SAVE_L
|  mov L:RB->base, BASE
|  mov FCARG2, PC
|  mov FCARG1, L:RB
|  call extern lj_dispatch_ins@8
|  mov BASE, L:RB->base
|  movzx RA, PC_RA
|  movzx OP, PC_OP
|  movzx RD, PC_RD
|  jmp aword [DISPATCH+OP*8+GG_DISP2STATIC]
|.endif

可以看到这段代码基本上是先调用lj_dispatch_ins,然后跳转到 [DISPATCH+OP*8+GG_DISP2STATIC] 指向的地址。luajit会在相距调度数组GG_DISP2STATIC的地方保存调度数组的副本。调度数组当下只包含指向的lj_vm_record的指针。
Lj_vm_record用这个副本数组来跳转到真正的汇编桩。
我们不要老是纸上谈兵,现在我们看看实实在在的代码,理解代码的精妙之处。
当一次循环或者调用被度量为热点后,翻译器调用lj_trace_hot,其为记录,编译和安装轨迹的入口。lj_trace_hot设置重要对象jit_State的状态为LJ_TRACE_START,并且把CPU控制权移交给lj_trace_ins。
luajit的跟踪子系统(tracing subsystem)有五个状态:START,RECORD, END, ASM和ERR。lj_trace_ins是一个有限自动机(finite automaton),其基于从执行轨迹上读出的字节码指令来进行状态之间的迁移。总体方案很简单:

START   -------------------> RECORD ---> ERR
          set up dispatch      |             
              vector           |         [IDLE]
                               v           /
                              END -----> ASM

当然,这些状态迁移不会同时发生--随着调用lj_dispatch_ins解析当前执行轨迹的越来越多字节码,状态被记录下来,状态迁移(或者不迁移)。跟踪器看到的字节码流被译作SSA形式的中间表示(intermediate SSA representation),然后再优化和编译成汇编。

跟踪只会给我们一个代码的线性路径,而不管其它已得到的相同的有效路径。跟踪JIT通常通过安装守卫(guards)来处理这种情况。守卫会检查各种假设条件,当违反假设时跳出轨迹(就luajit而言,会返回翻译器)。例如,考虑下面这个简单的lua程序:

total = 0
function test()
  if total < 500 then
      total = total + 1
   end
end
for i=1,501 do
   test()
end

这段代码会编译为如下的字节码(去掉一些对我们的理解不重要的部分)

;; function test()
0001    GGET     0   0
0002    KSHORT   1 500
0003    ISGE     0   1
0004    JMP      0 => 0008
0005    GGET     0   0
0006    ADDVN    0   0   0
0007    GSET     0   0
0008 => RET0     0   1
;; body (instructions 0001 to 0007 removed)
0008    FORI     0 => 0012
0009 => GGET     4   2      ; "test"
0010    CALL     4   1   1
0011    FORL     0 => 0009
0012 => RET0     0   1

(用-b -l可以让luajit以上面的形式输出字节码,用-jdump输出全部的轨迹信息)

没啥意外的--有一个for循环(以FORI和FORL来定界),调用了501次test。luajit提取下面这段SSA格式的轨迹(也删减了代码)

;; Instructions 0001 to 0013 removed
0014 >  p32 HREFK  0013  "total" @11
0015 >  num HLOAD  0014
0016 >  num LT     0015  +500
0017  + num ADD    0015  +1
0018    num HSTORE 0014  0017
0019  + int ADD    0001  +1
0020 >  int LE     0019  +501
0021 ------ LOOP ------------
0022 >  num LT     0017  +500
0023  + num ADD    0017  +1
0024    num HSTORE 0014  0023
0025  + int ADD    0019  +1
0026 >  int LE     0025  +501
0027    int PHI    0019  0025
0028    num PHI    0017  0023

LOOP之后的指令不足为奇:一个典型的循环加上一个寻常的SSA phi节点。注意(在跟踪JIT(tracing JITs)中很寻常),test的一个版本被完全内联到一个紧凑循环中。指令0016相当有意思:>表示这是一条特殊的指令,它是一个守卫(guard),即如果这条指令对应的条件不满足则这个轨迹会退出,并返回到翻译器。在这里,如果0015(为总的值,看看0014和0015指令你就明白了)大于等于500,则会退出。有意思的原因是跟踪编译器没有自作聪明,而是当总值大于等于500时啥也没做,这也是正确的行为。它只知道当总值小于500时,正确的作为是总值加1,因为这正是它在轨迹中所遵从的(what it has observed in the trace)。特别的是,如果总值大于等于500,路径已经走了足够多次,这条路径会被标记为热点,然后将编译成一个副轨迹(side-trace,相对root trace而言),你应该自己试试分析这个过程。

Lj_record.c中的lj_record_ins在字节码执行前,以SSA格式记录每条字节码指令。一个构架上的精妙之处在于知道要生成哪个守卫(guard)-- 假设一个条件:
IF A < B THEN {0} ELSE {1} END
是要生成A < B的守卫还是生成NOT (A < B)的守卫呢?这只有在条件求值之后才知道。如果在轨迹记录阶段,A < B为真,那么我们需要一个守卫来保证A < B,反之亦然。luajit没有在跟踪器中实现一个偏翻译器(partial interpreter),而是把守卫指令的插入推迟到调用下一个字节码的lj_record_ins之前。这时它就知道条件计算的结果了。

另外一个精妙之处涉及代码路径包含已经编译的轨迹。luajit通过替换正常的循环或者调用字节码为带有特殊标记的字节码来跟踪已编译的轨迹。这些特殊的字节码表示轨迹的存在,并提供开始于这点的已编译轨迹的位置。然而,当跟踪时,我们想要能看到函数字节码。解决方案是:先临时把特殊的JIT标记字节码(例如BC_JFUNCF或者BC_JLOOP)打补丁成原始的字节码,再让跟踪器走完原始字节码,在跟踪结束后再把补丁去掉,JIT标记字节又重新生效。这段处理在rec_func_jit(在lj_record.c中)和trace_pendpatch(在lj_trace.c中)。

快照

涉及在两个(或者多个!)执行模式转换的VM都会面临一个在未优化和已优化的堆栈布局之间的切换的问题。对luajit来说,也是这样。luajit的栈的含义稍有不同。我们希望同一个轨迹的翻译版本和编译版本之间是观测性等价(observational equivalence)的。而且编译版本有相同的副作用,操作数栈的映射方式也要跟翻译版本一致。然而,这种等价性不需要在已编译的轨迹内部保持,只需要在轨迹出口(用守卫方式或者不用)中保持即可。
luajit解决方案是:维护一个从栈位置到SSA IR指令的映射。这些映射在luajit代码库中被称为快照(snapshots)。使用一个快照,它能重新构建操作数栈就好像轨迹被翻译的那样(对应的代码在lj_snap.c中的snap_restoreval和ln_snap_restore)。
这很慢,但套用Mike Pall的话来说:『使用数据驱动方式的状态恢复当然会很慢。但是反复调用副出口(side exits)会很快触发副轨迹(side traces)的生成。快照用于初始化副轨迹的IR。初始化需要必要的状态,而状态要用似载荷(pseudo-loads)。这些能一起用副轨迹的余下部分来优化。通过后端让似载荷与父轨迹的机器状态相统一,从来能够零成本的把似载荷链接到副轨迹中。』

推荐阅读更多精彩内容