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

お久しぶりです。ずいぶんと更新をさぼっていたのですが、ときどきバグるけどmrubyのJITは元気です。さて、この半年、mrubyのJITもいろいろいじっていたのですが、多くの変更が　力技　| 黒魔術　って感じでエレガントな記事を望む皆様にはふさわしくないような気もしました。そこで、コミットログをあさってみたら、型推論もどきでガードを消しちゃうよというネタがちょっとアカデミックで良いのではないかと思って書いてみます。

型を格納する構造体

mrubyのJITでは各RITE VMの命令ごとにいろんなコンパイルやネイティブコードの実行に使う情報を格納するmrbjit_code_infoという構造体を用意しています。*1この構造体はこんな感じの定義です。

typedef struct mrbjit_code_info {
  mrbjit_code_area code_base;
  mrb_code *prev_pc;
  struct mrbjit_code_info *prev_coi;
  mrb_code *caller_pc;
  void *(*entry)();
  mrbjit_reginfo *reginfo;	/* For Local assignment */
  int used;
} mrbjit_code_info;

この中のreginfo構造体がRITE VMのレジスタの型情報を格納します。RITE VMのレジスタはどんな型のデータも自由に格納できるので、あくまでこの情報はその命令を実行した時点の情報です。レジスタは複数あるのでポインタ型になっていて配列として用意します。

reginfo構造体の中身はこんな感じです。

typedef struct mrbjit_reginfo {
  enum mrb_vtype type;     /* 型　*/
  struct RClass *klass;    /* クラス */
  int constp;              /* 定数か? */
} mrbjit_reginfo;

constpだけはなんだこりゃ？と思うかもしれないですね。文字通り値が決定している場合にnon-0になります。その値は？と思われるかもしれませんが、RITE VMのレジスタから取ってくればいいです。なにせ定数ですから。定数のたたみこみとか出来そうで夢が広がりそうですね。まだ定数のたたみこみサポートしてないですが。

こんな感じでレジスタの型情報を入れる入れ物を用意します。では、次に型情報を書き込み所を見てみましょう。

レジスタの型情報の書き込み

まず簡単なところで、シンボルを書き込みLOADSYM命令をコンパイルするコードを見てみましょう。

  const void *
    emit_loadsym(mrb_state *mrb, mrbjit_vmstatus *status, mrbjit_code_info *coi) 
  {
    const void *code = getCurr();
    mrb_code **ppc = status->pc;
    mrb_irep *irep = *status->irep;
    const Xbyak::uint32 dstoff = GETARG_A(**ppc) * sizeof(mrb_value);
    int srcoff = GETARG_Bx(**ppc);
    const Xbyak::uint32 src = (Xbyak::uint32)irep->syms[srcoff];
    mrbjit_reginfo *dinfo = &coi->reginfo[GETARG_A(**ppc)];  /* 書き込みレジスタの型情報 */
    dinfo->type = MRB_TT_SYMBOL;　　　　　　　　　　　　　　 /* もちろんシンボル型　*/
    dinfo->klass = mrb->symbol_class;                        /* Symbol class */
    dinfo->constp = 1;                                       /* シンボル番号命令にあるので定数 */

    mov(dword [ecx + dstoff], src);
    mov(dword [ecx + dstoff + 4], 0xfff00000 | MRB_TT_SYMBOL);

    return code;
  }

ここで、dinfoが書き込みレジスタの型情報です。シンボルを格納する命令だからシンボル型になる。単純な話です。

こんな感じで、書き込んだ型情報は次の命令を実行するときに基本的には引き継がれます。引き継ぐコードはこんな感じ。vm.cのmrbjit_dispatchにあります。

    if (ci->prev_coi && ci->prev_coi->reginfo) {
　　　/* 前回実行した命令にレジスタ情報がある場合 */
      mrbjit_reginfo *prev_rinfo;
      prev_rinfo = ci->prev_coi->reginfo;
      for (i = 0; i < irep->nregs; i++) {
	ci->reginfo[i] = prev_rinfo[i];
      }
    }
    else {
      /* ない場合　*/
      for (i = 0; i < irep->nregs; i++) {
	ci->reginfo[i].type = MRB_TT_FREE;
	ci->reginfo[i].klass = NULL;
	ci->reginfo[i].constp = 0;
      }
    }

こんな感じで作った型情報をどう使うのか? ちょっと疲れたので中断
ご飯食べてた、続き、

ガード

Tracing JITでは型チェックや分岐命令などでガードが大活躍します。ガードとは、想定した型や値になっているかチェックする仕組みでmrubyのJITでは想定と違った場合はVMに戻ります。ここで、もし型がコンパイル時に分かった場合はガードを省くことができるわけです。まず、レジスタが想定した型になっているかチェックするコードを生成する、gen_type_guardを見ています。

  void 
    gen_type_guard(mrb_state *mrb, int regpos, mrbjit_vmstatus *status, mrb_code *pc, mrbjit_code_info *coi)
  {
    enum mrb_vtype tt = (enum mrb_vtype) mrb_type((*status->regs)[regpos]);
    mrbjit_reginfo *rinfo = &coi->reginfo[regpos];

    if (rinfo->type == tt) {
      return;
    }

    mov(eax, dword [ecx + regpos * sizeof(mrb_value) + 4]); /* Get type tag */
    rinfo->type = tt;
    rinfo->klass = mrb_class(mrb, (*status->regs)[regpos]);
    /* Input eax for type tag */
    if (tt == MRB_TT_FLOAT) {
      cmp(eax, 0xfff00000);
      jb("@f");
    } 
    else {
      cmp(eax, 0xfff00000 | tt);
      jz("@f");
    }

    /* Guard fail exit code */
    gen_exit(pc, 1, 0, status);

    L("@@");
  }

ここで、

    mrbjit_reginfo *rinfo = &coi->reginfo[regpos];

    if (rinfo->type == tt) {
      return;
    }

が静的に分かっている型と想定している型が同じかどうかのチェックです。同じだった場合はガードを生成する必要が無いので戻ります。
また、ガードが通った場合はこの先想定した型であることが保障されるので、レジスタの型情報も行進しておきます。こうすることで、同じ値を何度もガードでチェックすることが避けられます。

    rinfo->type = tt;
    rinfo->klass = mrb_class(mrb, (*status->regs)[regpos]);

同様にレジスタが想定したクラスかどうかチェックするgen_class_guardも静的なレジスタの型情報でガードが減らせます。

  void 
    gen_class_guard(mrb_state *mrb, int regpos, mrbjit_vmstatus *status, mrb_code *pc, mrbjit_code_info *coi, struct RClass *c)
  {
    enum mrb_vtype tt;
    mrb_value v = (*status->regs)[regpos];
    mrbjit_reginfo *rinfo = &coi->reginfo[regpos];

    tt = (enum mrb_vtype)mrb_type(v);

    if (rinfo->type != tt) {

      rinfo->type = tt;

      mov(eax, ptr [ecx + regpos * sizeof(mrb_value) + 4]);

      if (tt == MRB_TT_FLOAT) {
	cmp(eax, 0xfff00000);
	jb("@f");
      }
      else {
	cmp(eax, 0xfff00000 | tt);
	jz("@f");
      }

      /* Guard fail exit code */
      gen_exit(pc, 1, 0, status);

      L("@@");
    }

    /* Import from class.h */
    switch (tt) {
    case MRB_TT_FALSE:
    case MRB_TT_TRUE:
    case MRB_TT_SYMBOL:
    case MRB_TT_FIXNUM:
    case MRB_TT_FLOAT:
      /* DO NOTHING */ /*　クラスオブジェクトは分かり切っているので比較はしない */
      break;

    default:  /* クラスオブジェクトのチェック */
      {
	if (c == NULL) {
	  c = mrb_object(v)->c;
	}
	if (rinfo->klass == c) {
	  return;
	}
	rinfo->klass = c;
	mov(eax, dword [ecx + regpos * sizeof(mrb_value)]);
	mov(eax, dword [eax + OffsetOf(struct RBasic, c)]);
	cmp(eax, (int)c);
	jz("@f");
	/* Guard fail exit code */
	gen_exit(pc, 1, 0, status);

	L("@@");
      }
      break;
    }
  }

また、定数かどうかを示すconstpを使うことで、条件分岐命令使うレジスタが想定しているbool値になっているかチェックするgen_bool_guardで条件判定の削除が実現できます。

  void
    gen_bool_guard(mrb_state *mrb, int b, mrb_code *pc, 
		   mrbjit_vmstatus *status, mrbjit_reginfo *rinfo)
  {
    if (rinfo->constp) {  /* 定数だった場合　*/
      if (b && rinfo->type != MRB_TT_FALSE) {  /* 真と予想して、そうだった場合 */
	return;
      }
      if (!b && rinfo->type == MRB_TT_FALSE) {  /* 偽と予想して、そうだった場合 */
	return;
      }
    }

    cmp(eax, 0xfff00001);
    if (b) {
      jnz("@f");
    } 
    else {
      jz("@f");
    }

    /* Guard fail exit code */
    gen_exit(pc, 1, 0, status);

    L("@@");
  }

こんな感じです。それではまた、ごきげんよう。

*1:正確には直前に実行した命令が異なれば別の構造体を割り当てています。これが、型推論を簡単にする秘訣になっています