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

とってもお久しぶりです。ずいぶん間が空いてしまいましたが、新たな機能のプリミティブのインライン化が実装できたので再開したいと思います。あなとみー おぶ mrubyのJIT の開始当時オリジナルmrubyの数割速いとか遅くなるとか言っていましたが、現在大体のベンチマークでmrubyの2〜4倍の速度が出ます。その秘密が今回紹介するプリミティブのインライン化です。
Rubyはかなり基本的な機能もすべてメソッド呼び出しで実装されています。これは、必要に応じて動作をカスタマイズしたりして柔軟性をもたらしますが、速度的には大きなハンディになります。とくに、mrubyはメソッド呼び出しのオーバーヘッドが大きいので問題になります。
そこで、mrubyのJITでは良く使うメソッドについてはインライン化して速度を稼ぐようにしました。

インライン化を支えるコード

具体的なコードを見てみましょう。
これは、array.cの[]メソッドの定義ですが、こんな感じでプリミティブとしてインライン化したいメソッドを登録します。

  mrb_define_method(mrb, a, "[]",              mrb_ary_aget,         MRB_ARGS_ANY());  /* 15.2.12.5.4  */
  mrbjit_define_primitive(mrb, a, "[]", mrbjit_prim_ary_aget);

こうすると、JITコンパイル時に、mrbjit_prim_ary_agetを呼び出してインライン化するようになります。

mrbjit_define_primitiveの定義です。class.cにあります。

void
mrbjit_define_primitive_id(mrb_state *mrb, struct RClass *c, mrb_sym mid, mrb_func_t func)
{
  struct RProc *p;
  int ai = mrb_gc_arena_save(mrb);

  p = mrb_proc_new_cfunc(mrb, func);
  p->target_class = c;
  mrb_obj_iv_set(mrb, (struct RObject*)c, mid, mrb_obj_value(p));
  mrb_gc_arena_restore(mrb, ai);
}

void
mrbjit_define_primitive(mrb_state *mrb, struct RClass *c, const char *name, mrb_func_t func)
{
  mrbjit_define_primitive_id(mrb, c, mrb_intern(mrb, name), func);
}  

コードを見てわかるとおり、インスタンス変数領域にインライン化の関数をProcオブジェクトとして格納します。インスタンス変数は@, @@で普通始まるので通常のインスタンス変数とは衝突しません。まれに特殊目的で@,@@で始まらない変数を登録する場合もありますので注意が必要です。

実際にインライン化する場所を見てみましょう。jitcode.hのemit_sendです。

    if (MRB_PROC_CFUNC_P(m)) {
      prim = mrb_obj_iv_get(mrb, (struct RObject *)c, mid);
      mrb->vmstatus = status;
      if (mrb_type(prim) == MRB_TT_PROC) {
	mrb_value res = mrb_proc_ptr(prim)->body.func(mrb, prim);
	if (!mrb_nil_p(res)) {
	  return code;
	}
      }

      //puts(mrb_sym2name(mrb, mid)); // for tuning
      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 {

呼び出そうとするメソッドがCで定義された場合の処理ですが、インライン化すると登録された場合は登録された関数を呼び出してインライン化するコードを生成します。生成に成功したら、nil以外を返すことになっているのでチェックして、その後のmrbjit_exec_send_cの呼び出しはキャンセルします。コード生成が失敗したら、nilを返して何事もなかったようにmrbjit_exec_send_cを生成して通常のメソッド呼び出し
シーケンスを実行します。

インライン化コードの例

では、実際にインライン化を行うコードを見てみましょう。primitive.ccでいろいろ定義していますが、その中で簡単なFixnum#succのインライン化を解説します。

mrb_value
MRBJitCode::mrbjit_prim_fix_succ_impl(mrb_state *mrb, mrb_value proc)
{
  mrbjit_vmstatus *status = mrb->vmstatus;
  mrb_code *pc = *status->pc;
  int i = *pc;
  int regno = GETARG_A(i);
  const Xbyak::uint32 off0 = regno * sizeof(mrb_value);

  add(dword [ecx + off0], 1);

  return mrb_true_value();
}

extern "C" mrb_value
mrbjit_prim_fix_succ(mrb_state *mrb, mrb_value proc)
{
  MRBJitCode *code = (MRBJitCode *)mrb->compile_info.code_base;

  return code->mrbjit_prim_fix_succ_impl(mrb, proc);
}

インライン化するコードはC++で書かれた実際にインライン化を行うMRBJitCodeのメソッド(MRBJitCode::mrbjit_prim_fix_succ_impl)と、そのメソッドを単に呼び出すCで書かれた関数(mrbjit_fix_succ)の組になります。C++で定義されたメソッドはProcオブジェクトとして扱えないからです。

インライン化するメソッドを見てみましょう。

mrb_value
MRBJitCode::mrbjit_prim_fix_succ_impl(mrb_state *mrb, mrb_value proc)
{
  mrbjit_vmstatus *status = mrb->vmstatus;
  mrb_code *pc = *status->pc;
  int i = *pc;
  int regno = GETARG_A(i);
  const Xbyak::uint32 off0 = regno * sizeof(mrb_value);

  add(dword [ecx + off0], 1);

  return mrb_true_value();
}

単にadd命令を生成するだけです。足すだけといってもどこを足すのか指定しなければならないので少し複雑です。オーバフロー時の処理を忘れていました。TODOということで・・・。

mrb->vmstatusにvmの内部状態が入っています。この中から現在実行中の命令(当然OP_SENDですね)を取り出します。これはpcメンバーにあります。
OP_SEND命令のA引数にはselfのレジスタ番号が入っているのでこれを取り出してmrb_valueのサイズを掛けてバイト単位のオフセットを求めています。

そして、add命令を生成してTrueを戻り値として返します。Non-Nilなのでこのインライン化は有効です。

今回は常にインライン化は有効になりますが、引数の条件によってインライン化は行わないようにすることも可能です。この場合はreturn mrb_nil_value();とします。

primitive.ccを見るともっと複雑な例をありますので、興味のある人は見てください。

そんなこんなで今回は終わり。そろそろねたも尽きかけたのでまた新しい機能ができたらお会いしましょう。ではでは