IT博客汇 | 实现JVM中的JIT

实现JVM中的JIT

Kevin Lynx发表于 2017-03-09 00:00:00

在JVM中，JIT (Just-in-Time) 即时编译指的是在Java程序运行过程中JVM优化部分指令为本地指令，从而大幅提升性能。在上一篇文章写一个玩具Java虚拟机中实现了一个基本可以运行Java字节码的JVM。本篇文章描述我是如何在这个玩具JVM中实现JIT的。

推荐文章“How to JIT - an introduction”，介绍了JIT的基本实现原理。作者把JIT分为两个阶段：

运行期生成机器代码(本地指令)
执行机器代码

生成机器代码很好理解，就是一个JVM指令到机器指令的翻译；而执行机器代码，原理上是利用了OS提供了API可以分配可以执行的内存，然后往这块内存中写入机器码，从而实现运行期可以执行动态生成的机器码功能。

我们可以利用这个原理来实现JIT，但是未免太底层了点，需要做很多工作来完成这件事情。我们可以利用libjit来简化实现。这个作者博客里还有些libjit的教程，其中part 1值得阅读。简单来说，libjit对机器指令做了抽象，利用它的API来描述一个函数包含了哪些指令，实现了什么功能。然后具体的指令生成以及指令执行则交给libjit完成。

例如以下使用libjit的代码：

// t = u
jit_insn_store(F, t, u); // 类似 mov 指令
// u = v
jit_insn_store(F, u, v);

// v = t % v
jit_value_t rem = jit_insn_rem(F, t, v); // 求余指令
jit_insn_store(F, v, rem);

所以，我们需要做的，就是将JVM的字节码，翻译成一堆libjit的API调用。但是我希望能够稍微做点抽象，我们写个翻译器，能够将JVM这种基于栈的指令，翻译成基于寄存器的指令，才方便后面无论是使用libjit还是直接翻译成机器码。

指令翻译

要将基于栈的指令翻译成基于寄存器的指令（类似），仔细想想主要解决两个问题：

去除操作数栈
跳转指令所需要的标签

去除操作数栈，我使用了一个简单办法，因为JVM中执行字节码时，我们是可以知道每条指令执行时栈的具体情况的，也就是每条指令执行时，它依赖的操作数在栈的哪个位置是清楚的。例如，假设某个函数开头有以下指令：

opcode [04] - 0000: iconst_1   # [1]
opcode [3C] - 0001: istore_1   # []
opcode [1B] - 0002: iload_1    # [1]
opcode [1A] - 0003: iload_0    # [1, N]
opcode [68] - 0004: imul       # [1 * N]

当执行imul指令时，就可以知道该指令使用栈s[0]、s[1]的值，做完计算后写回s[0]。所以，类似JVM中局部变量用数字来编号，我也为栈元素编号，这些编号的元素全部被视为局部变量，所以这些指令全部可以转换为基于局部变量的指令。为了和JVM中本身的局部变量统一，我们将栈元素编号从局部变量后面开始。假设以上函数有2个局部变量，那么栈元素从编号2开始，局部变量编号从0开始。以上指令可以翻译为：

mov 1, $2   # 常量1写入变量2
lod $2, $1  # 变量2写入变量1
lod $1, $2  # 变量1写回变量2
lod $0, $3  # 变量0写入变量3
mul $3, $2  # 变量3与变量2相乘，写回变量2

这里，我们定义了自己的中间指令集(IR)，这个中间指令集存在的意义在于，在将来翻译为某个平台的机器码时，它比JVM的指令集更容易理解。中间指令集是一种抽象，方便基于它们使用libjit或其他手段翻译成机器码。

不过，我们看到上面的指令非常冗余。要优化掉这种冗余相对比较复杂，所以本文暂时不讨论这个问题。

这个中间指令基于局部变量的方式，是利于JIT下游做各种具体实现的，例如是否直接转换为通用寄存器，即一定范围的局部变量数是可以直接使用寄存器实现，超出该范围的局部变量则放在栈上，用栈模拟；或者全部用栈模拟。注意在机器指令中栈元素是可以直接偏移访问的，不同于“基于栈的虚拟机”中的栈。

以上指令，我们可以简单地为每条指令设定如何翻译为libjit的调用，例如mov指令：

static void build_mov(BuildContext* context, const Instruction* inst) {
  jit_value_t c = jit_value_create_nint_constant(context->F, jit_type_int, inst->op1); 
  jit_insn_store(context->F, context->vars[inst->op2], c);
}

例如mul指令：

static void build_mul(BuildContext* context, const Instruction* inst) {
  // context->vars就是前面说的局部变量表，包含了JVM中的局部变量及操作数栈
  jit_value_t tmp = jit_insn_mul(context->F, context->vars[inst->op1], context->vars[inst->op2]);
  jit_insn_store(context->F, context->vars[inst->op1], tmp);
}

接下来说另一个问题：跳转指令的标签。在机器指令中，跳转指令跳转的目标位置是一个绝对地址，或者像JVM中一样，是一个相对地址。但是在我们的中间指令集中，是没有地址的概念的，在翻译为机器指令时，也无法获取地址。所以，我们一般是增加了一个特殊指令label，用于打上一个标签，设置一个标签编号，相当于是一个地址。在后面的跳转指令中，则跳转的是这个标签编号。

所以，我们需要在翻译JVM指令到我们的中间指令时，识别出哪些地方需要打标签；并且在翻译跳转类指令时，翻译为跳转到某个编号的标签。

例如以下指令：

opcode [04] - 0000: iconst_1
opcode [3C] - 0001: istore_1
opcode [1B] - 0002: iload_1     # 会被调整，需要在此打标签
opcode [1A] - 0003: iload_0
...
opcode [1A] - 0010: iload_0
opcode [9D] - 0011: ifgt -9     # pc-9，也就是跳转到0002位置

为了打上标签，我们的翻译需要遍历两遍指令，第一遍用来找出所有标签，第二遍才做真正的翻译。

  // 该函数遍历所有指令，找出所有需要打标签的指令位置
  private List<Integer> createLabels(List<InstParser.Instruction> jbytecode) {
    List<Integer> labels = new LinkedList<>();
    for (InstParser.Instruction i : jbytecode) {
      LabelParser labelParser = labelParsers.get(i.opcode);
      if (labelParser != null) { // 不为空时表示是跳转指令
        int pc = labelParser.parse(i); // 不同的跳转指令地址解析不同，解析得到跳转的目标地址
        labels.add(pc); // 保存起来返回
      }
    }
    return labels;
  }

然后在翻译指令的过程中，发现当前翻译的指令地址是跳转的目标位置时，则生成标签指令：

   List<Integer> labels = createLabels(jbytecode);
   ...
   Iterator<InstParser.Instruction> it = jbytecode.iterator();
    while (it.hasNext()) {
      InstParser.Instruction inst = it.next();
      int label = labels.indexOf(inst.pc);
      if (label >= 0) {
        state.addIR(new Inst(op_label, label)); // 生成标签指令，label就是标签编号
      }

在处理跳转指令时，则填入标签编号：

translators.put(Opcode.op_ifgt, (state, inst, iterator) -> {
  short offset = (short)((inst.op1 << 8) + inst.op2);
  int pc = inst.pc + offset;
  int label = state.findLabel(pc); // 找到标签编号
  int var = state.popStack();
  state.addIR(new Inst(op_jmp_gt, var, label));
});

我们的中间指令集中，跳转指令和标签指令就为：

label #N            // 打上标签N
jmp_gt $var, #N     // 如果$var>0，跳转到标签#N

看下使用libjit如何翻译以上两条指令：

static void build_label(BuildContext* context, const Instruction* inst) {
  // 打上标签，inst->op1为标签编号N，对应写到context->labels[N]中
  jit_insn_label(context->F, &context->labels[inst->op1]);
}

static void build_jmp_gt(BuildContext* context, const Instruction* inst) {
  jit_value_t const0 = jit_value_create_nint_constant(context->F, jit_type_int, 0);
  // 是否>0
  jit_value_t cmp_v_0 = jit_insn_gt(context->F, context->vars[inst->op1], const0);
  // 大于0则跳转到标签inst->op2
  jit_insn_branch_if(context->F, cmp_v_0, &context->labels[inst->op2]);
}

代码贴得有点多，大概懂原理就行了。

在JIT中还有个很重要的过程，就是判定哪些代码需要被JIT。这里只是简单地尝试对每一个函数进行JIT，发现所有指令都能够被JIT时就JIT。

指令执行

在上一篇文章中，执行每个JVM函数时，都会有一个Frame与之关联。所以，在这里只要函数被JIT了，对应的帧就会包含被编译的代码，也就是libjit中的jit_function_t。在该Frame被执行时，就调用libjit执行该函数：

  private void runNative() {
    int arg_cnt = getArgsCount();
    int[] args = new int[arg_cnt];
    for (int i = 0; i < arg_cnt; ++i) {
      if (mLocals[i].type != Slot.Type.NUM) throw new RuntimeException("only supported number arg in jit");
      args[i] = mLocals[i].i;
    }
    int ret = mJIT.invoke(args); // mJIT后面会看到，主要就是将参数以数组形式传递到libjit中，并做JIT函数调用
    mThread.popFrame();
    if (hasReturnType() && mThread.topFrame() != null) {
      mThread.topFrame().pushInt(ret); // 目前只支持int返回类型
    }
  }

实现

以上就是整个JIT的过程，主要工作集中于JVM指令到中间指令，中间指令到libjit API调用。整个实现包含以下模块：

+-----------+       +----------+
| ASM       |       | libjit   |
|           | <-----+ API call |
+-----------+       +----+-----+
                         ^
                         |
+-----------+       +----+-----+
|  JVM      |       |  IR code |
|  bytecode +-----> |          |
+-----------+       +----------+

JVM byte code及IR code的处理是在Java中完成的；处理完后将IR code输出为byte[]，通过JNI调用包装好的C API。这个C API则是基于libjit，将IR code翻译为libjit的API调用。指令翻译完后调用libjit的API得到最终的ASM机器指令。

同样，要执行指令时，也是通过JNI调用这个C API。JNI交互全部包装在以下类中：

public class ToyJIT {
  private long jitPtr = 0;

  public void initialize(byte[] bytes, int maxLocals, int maxLabels, int argCnt, int retType) {
    jitPtr = compile(bytes, maxLocals, maxLabels, argCnt, retType);
  }

  public int invoke(int... args) {
    return invoke(jitPtr, args);
  }

  static {
    System.loadLibrary("toyjit");
  }
  private static native long compile(byte[] bytes, int maxLocals, int maxLabels, int argCnt, int retType);
  private static native int invoke(long jitPtr, int[] args);

即，libtoyjit.so 主要提供翻译接口 compile 及执行接口 invoke。

性能对比

简单测试了下一个阶乘计算函数：

  public static int fac2(int n) {
    int r = 1;
    do {
      r = r * n;
      n = n - 1;
    } while (n > 0);
    return r;
  }

 ...
    int i = 0;
    for (; i < 10000; ++i) {
        fac2(100);
    }
 ...

fac2函数会被JIT，测试发现不开启JIT时需要16秒，开启后1秒，差距还是很明显的。

最后奉上代码，toy_jit，就是前面说的C API部分，翻译IR到libjit API call，包装接口用于JNI调用。redhat 7.2下编译，需要先编译出libjit，我是直接clone的libjit master编译的。Java部分还是在toy_jvm中。