miura1729の日記

あなとみー　おぶ　mrubyのJIT (その12)

09:26 |

今回はオブジェクト作成のメソッドnewのインライン化のお話です。newのインライン化はそれほど規模が大きくないのですが、mrubyのJITの大きな転換点になっています。

前回取り組んだytljitがRubyの型推論(ブロックや要素ごとに違う型が格納できる配列などを含む)やオブジェクトのスタック割り付けなどを行って複雑すぎて手に負えなくなっちゃったことがありました。その反省を踏まえて、できるだけオーソドックスに分かりやすく気を使っていました(結果どうなったかは知らない)。

ところが、aoベンチでオブジェクトの生成がボトルネックっぽいと分かったことでnewをインライン化を初めて、検討の結果またも悪魔に魂を売り渡してしまいました。その、技法とは「バイトコードの書き換え」です。

newの難しいところはオブジェクトを生成した後、initilaizeメソッドを呼ばないといけないところです。簡単じゃないかと思われるかもしれませんが、mrubyではnewでのinitilizeメソッドはmrb_run直接ではなく、vm.cで定義されているmrb_funcall_with_block経由で呼ばれます(結局mrb_runで実行されるのですが)。そうすると、詳しい理由は避けますがJITコンパイルはうまくいきません。そういうことで、newをコンパイルしてその機械語を実行するときinitizlzeメソッドの機械語コードが生成されていないのです。mrb_funcall_with_blockを呼び出してインタープリタでinitilizeメソッドを実行してもいいのですが、それでは面白くないです。

いろいろ考えて、newをmrb_instance_allocで得た初期化していないオブジェクトをレシーバーにしてinitializeメソッドを呼び出す動作にコンパイルしたらどうだろうと考えました。このようにすることで、initalizeはmrb_funcall_with_block経由ではなく、mrb_run直接実行になります。前述のとおり、initializeメソッドは機械語が生成されていないので、メソッド呼び出しのためのcallinfoを構築したら、即座にVMに戻るコードを生成するようにしました。

これで、基本的にうまくいったのですが１つ問題があります。newメソッドは生成したオブジェクトを返すのですが、initializeメソッドは生成したオブジェクト(initializeメソッドのself)を返すとは限らないことです。initializeメソッドから戻ったところで改めて生成したオブジェクトを返すようにすればいいのですが、そうするとinitializeメソッドを呼びっぱなしにできないのでcallinfo構築にコストがかかります。要は必ずinitializeメソッドで必ずselfを返してくれればいいのです。

戻り値を返すVMの命令はOP_RETURNなわけですが、

　　OP_RETURN    R0

とすると現在のmrubyでは必ずselfを返します。つまり、OP_RETURNのAオペランドを0クリアすればよいのです。initializeメソッドの戻り値は使われることがないので副作用もありません。簡単に実装できて効果のあり、被害も少なそう、な命令書き換えをやらない手があるだろうか？　ということで命令書き換えで実装しました。

コードを見てみましょう

mrb_value
MRBJitCode::mrbjit_prim_instance_new_impl(mrb_state *mrb, mrb_value proc,
					  mrbjit_vmstatus *status, mrbjit_code_info *coi)
{
  mrb_value *regs = *status->regs;
  mrb_irep *irep = *status->irep;
  mrb_code *pc = *status->pc;
  int i = *pc;
  int a = GETARG_A(i);
  int nargs = GETARG_C(i);

  struct RProc *m;
  mrb_value klass = regs[a];
  struct RClass *c = mrb_class_ptr(klass);

  m = mrb_method_search_vm(mrb, &c, mrb_intern(mrb, "initialize"));

  // TODO add guard of class
  
  // obj = mrbjit_instance_alloc(mrb, klass);
  push(ecx);
  push(ebx);
  mov(eax, *((Xbyak::uint32 *)(&klass) + 1));
  push(eax);
  mov(eax, *((Xbyak::uint32 *)(&klass)));
  push(eax);
  push(esi);
  call((void *)mrbjit_instance_alloc);
  add(esp, 3 * sizeof(void *));
  pop(ebx);
  pop(ecx);

  // regs[a] = obj;
  mov(ptr [ecx + a * sizeof(mrb_value)], eax);
  mov(ptr [ecx + a * sizeof(mrb_value) + 4], edx);

  if (MRB_PROC_CFUNC_P(m)) {
    CALL_CFUNC_BEGIN;
    mov(eax, (Xbyak::uint32)c);
    push(eax);
    mov(eax, (Xbyak::uint32)m);
    push(eax);
    CALL_CFUNC_STATUS(mrbjit_exec_send_c, 2);
  }
  else {
    /* patch initialize method */
    mrb_irep *pirep = m->body.irep;
    mrb_code *piseq = pirep->iseq;
    for (int i = 0; i < pirep->ilen; i++) {
      if (GET_OPCODE(piseq[i]) == OP_RETURN) {
	/* clear A argument (return self always) */
	piseq[i] &= ((1 << 23) - 1);
      }
    }
    
    /* call info setup */
    CALL_CFUNC_BEGIN;
    mov(eax, (Xbyak::uint32)c);
    push(eax);
    mov(eax, (Xbyak::uint32)m);
    push(eax);
    CALL_CFUNC_STATUS(mrbjit_exec_send_mruby, 2);

    mov(eax, dword [ebx + OffsetOf(mrbjit_vmstatus, regs)]);
    mov(ecx, dword [eax]);

    gen_set_jit_entry(mrb, pc, coi, irep);

    gen_exit(m->body.irep->iseq, 1, 0);
  }

  return mrb_true_value();
}

細かく見てみます。

  // obj = mrbjit_instance_alloc(mrb, klass);
  push(ecx);
  push(ebx);
  mov(eax, *((Xbyak::uint32 *)(&klass) + 1));
  push(eax);
  mov(eax, *((Xbyak::uint32 *)(&klass)));
  push(eax);
  push(esi);
  call((void *)mrbjit_instance_alloc);
  add(esp, 3 * sizeof(void *));
  pop(ebx);
  pop(ecx);

  // regs[a] = obj;
  mov(ptr [ecx + a * sizeof(mrb_value)], eax);
  mov(ptr [ecx + a * sizeof(mrb_value) + 4], edx);

初期化していないオブジェクトの生成です。コメントのCのコードのほぼ直訳です。

  if (MRB_PROC_CFUNC_P(m)) {
    CALL_CFUNC_BEGIN;
    mov(eax, (Xbyak::uint32)c);
    push(eax);
    mov(eax, (Xbyak::uint32)m);
    push(eax);
    CALL_CFUNC_STATUS(mrbjit_exec_send_c, 2);
  }
  else {

mにはinitializeのProcオブジェクトが入っています。mがC関数だった場合の処理です。ここを通ることは意外と少ないです。配列とかもmrubyでinitializeが定義されていますし。

    /* patch initialize method */
    mrb_irep *pirep = m->body.irep;
    mrb_code *piseq = pirep->iseq;
    for (int i = 0; i < pirep->ilen; i++) {
      if (GET_OPCODE(piseq[i]) == OP_RETURN) {
	/* clear A argument (return self always) */
	piseq[i] &= ((1 << 23) - 1);
      }
    }
    

initializeがmrubyで定義されている場合の処理です。ここで、命令書き換えをやっています。

    /* call info setup */
    CALL_CFUNC_BEGIN;
    mov(eax, (Xbyak::uint32)c);
    push(eax);
    mov(eax, (Xbyak::uint32)m);
    push(eax);
    CALL_CFUNC_STATUS(mrbjit_exec_send_mruby, 2);

    mov(eax, dword [ebx + OffsetOf(mrbjit_vmstatus, regs)]);
    mov(ecx, dword [eax]);

    gen_set_jit_entry(mrb, pc, coi, irep);

次にmrubyのメソッドを呼び出すcallinfoの構築です。説明はとても面倒なので省略します。

    gen_exit(m->body.irep->iseq, 1, 0);

最後にVMに戻ります。

泥臭い話が続きますが、次回も多分泥臭いです。コンパイラは開発が進むと泥臭くなるもののようです。