mrubyのContextThreading化 その3

不具合も解消したので未実装のものをどうにかする。とかいいつつ、実装しなければならないモノがどれだけあるのかわからないので目標をAOベンチ完走に置いて、動かないところを調べていった。
昨日の時点で未実装の機能はメソッド途中のOP_RETURNと、引数省略時のOP_ENTERである。
とりあえずOP_RETURNはサクっと。コードは色々と手直ししたのでMAKE_CT_FUNCを全部。

#define MAKE_CT_FUNC \
  if (!irep->ct_func) {\
    int index;\
    unsigned char *data = mrb_malloc(mrb, irep->ilen * 11 + 3);\
    unsigned char **adr_buf = mrb_malloc(mrb, irep->ilen * sizeof(int));\
    unsigned char *return_adr;\
    irep->ct_func = (void*)data;\
    for (index = 0; index < irep->ilen; index++) {\
      adr_buf[index] = data;\
      *data = 0xe8; /* call rel32 */ \
      *(int**)(data + 1) = (int*)((unsigned char*)optable[GET_OPCODE(irep->iseq[index])] - (data + 5));\
      data += 5;\
      switch GET_OPCODE(irep->iseq[index]) {\
      case OP_JMP:\
      case OP_RETURN:\
        *data = 0xe9; /* jmp rel32 */ \
        data += 5;\
        break;\
      case OP_JMPIF:\
        *data = 0x0f; /* jne rel32 */ \
        *(data + 1) = 0x85;\
        data += 6;\
        break;\
      case OP_JMPNOT:\
        *data = 0x0f; /* je rel32 */ \
        *(data + 1) = 0x86;\
        data += 6;\
        break;\
      }\
    }\
    return_adr = data;\
    *data++ = 0xc2; /* ret $04 */\
    *data++ = 0x04;\
    *data++ = 0x00;\
    for (index = 0; index < irep->ilen; index++) {\
      switch GET_OPCODE(irep->iseq[index]) {\
      case OP_JMP:\
        *(int**)(adr_buf[index] + 6) = (int*)(adr_buf[index + GETARG_sBx(irep->iseq[index])] - (adr_buf[index] + 10));\
        break;\
      case OP_JMPIF:\
      case OP_JMPNOT:\
        *(int**)(adr_buf[index] + 7) = (int*)(adr_buf[index + GETARG_sBx(irep->iseq[index])] - (adr_buf[index] + 11));\
        break;\
      case OP_RETURN:\
        *(int**)(adr_buf[index] + 6) = (int*)(return_adr - (adr_buf[index] + 10));\
        break;\
      }\
    }\
    mrb_free(mrb, adr_buf);\
  }

ちょっとコメントを入れておいたのでアセンブラの命令がわかる人にはわかりやすくなったかもしれない。
今まではOP_RETURNはメソッドの最後と決め撃ちしていたが、メソッドの途中の場合にもきちんとct_funcを終了させてやる必要があるので、OP_RETURNを識別してret命令までジャンプするようにした。以上だ。

ではOP_ENTER。これが曲者である。


OP_ENTERはメソッドに引数がある場合にirepの先頭に出力される命令である。かなり難しい。難しい度合いをわかって頂くためにOP_ENTERのVMのコードをここに載せておこう。

    CASE(OP_ENTER) {
      /* Ax             arg setup according to flags (24=5:5:1:5:5:1:1) */
      /* number of optional arguments times OP_JMP should follow */
      int ax = GETARG_Ax(i);
      int m1 = (ax>>18)&0x1f;
      int o  = (ax>>13)&0x1f;
      int r  = (ax>>12)&0x1;
      int m2 = (ax>>7)&0x1f;
      /* unused
      int k  = (ax>>2)&0x1f;
      int kd = (ax>>1)&0x1;
      int b  = (ax>>0)& 0x1;
      */
      int argc = mrb->ci->argc;
      mrb_value *argv = regs+1;
      int len = m1 + o + r + m2;
      mrb_value *blk = &argv[argc < 0 ? 1 : argc];

      if (argc < 0) {
        struct RArray *ary = mrb_ary_ptr(regs[1]);
        argv = ary->ptr;
        argc = ary->len;
        mrb_gc_protect(mrb, regs[1]);
      }
      if (mrb->ci->proc && MRB_PROC_STRICT_P(mrb->ci->proc)) {
        if (argc >= 0) {
          if (argc < m1 + m2 || (r == 0 && argc > len)) {
            argnum_error(mrb, m1+m2);
            goto L_RAISE;
          }
        }
      }
      else if (len > 1 && argc == 1 && argv[0].tt == MRB_TT_ARRAY) {
        argc = mrb_ary_ptr(argv[0])->len;
        argv = mrb_ary_ptr(argv[0])->ptr;
      }
      mrb->ci->argc = len;
      if (argc < len) {
        regs[len+1] = *blk; /* move block */
        memmove(&regs[1], argv, sizeof(mrb_value)*(argc-m2)); /* m1 + o */
        memmove(&regs[len-m2+1], &argv[argc-m2], sizeof(mrb_value)*m2); /* m2 */
        if (r) {                  /* r */
          regs[m1+o+1] = mrb_ary_new_capa(mrb, 0);
        }
        if (o == 0) {
          pc++;
        }
        else {
          pc += argc - m1 - m2 + 1;
        }
      }
      else {
        memmove(&regs[1], argv, sizeof(mrb_value)*(m1+o)); /* m1 + o */
        if (r) {                  /* r */
          regs[m1+o+1] = mrb_ary_new_elts(mrb, argc-m1-o-m2, argv+m1+o);
        }
        memmove(&regs[m1+o+r+1], &argv[argc-m2], sizeof(mrb_value)*m2);
        regs[len+1] = *blk; /* move block */
        pc += o + 1;
      }
      JUMP;
    }

ぶっちゃけ、コードを見ただけではなにがなんだかさっぱりわからない。そもそもOP_ENTERって何をするためのものなん?て感じ。先頭のコメントにはオプション引数があるだけOP_JMPが続くんだよ(エスパー訳)と書いてあるように見える。
コードを見てもわからないときはVM命令列を吐き出してみるとよい。

def op_enter_test(a, b = 1, c = 2, &block)
end

絵に描いたようなひどいテストコードだが、これを--verboseで出力するとこうなる。

NODE_SCOPE:
  NODE_BEGIN:
    NODE_DEF:
      op_enter_test
      local variables:
        a         b         c         block
      mandatory args:
        NODE_ARG a
      optional args:
        b=NODE_INT 1 base 10
        c=NODE_INT 2 base 10
      blk=&block
      NODE_BEGIN:
irep 148 nregs=3 nlocals=2 pools=0 syms=1
000 OP_TCLASS   R2
001 OP_LAMBDA   R3      I(149)  1
002 OP_METHOD   R2      :op_enter_test
003 OP_LOADNIL  R2
004 OP_STOP

irep 149 nregs=7 nlocals=6 pools=0 syms=0
000 OP_ENTER    1:2:0:0:0:0:1
001 OP_JMP              004
002 OP_JMP              005
003 OP_JMP              006
004 OP_LOADI    R2      1
005 OP_LOADI    R3      2
006 OP_RETURN   R5

OP_ENTER命令のパラメータは最初が省略不可の引数の数で、次が省略可能引数の数、最後がブロックの数だろう。この例では他はわからないが、まあ問題は無い。この情報を元に↑のCのコードを眺めると少し意味がわかってくる。変数oに省略可能引数の数が入っていて、省略可能引数があるだけpcを余分に進めようとしているのだ。それ以外は残余引数として配列で入ってきたものを展開する処理だろう。
そしてVM命令列を眺めると、これがどういう動きをするのかが見えてくる。
例えばメソッドの呼び出しがop_enter_test(1,2)とかだった場合、第3引数が省略されて、OP_ENTER内でpcが+2される。そうすると002のOP_JMP 005に移動して、そっから005のOP_LOADI R3 2に飛んで、ローカル変数cに2が入る。004行目はbに1を入れる処理だから、この命令は引数で指定済みなので動かなくていい。なるほど。なるほど。

ContextThreadingではVM命令列は元と同じものとするから、このVM命令列が動作することが前提になる。つまり、OP_ENTERを呼び出した後で、pcを進めた値に応じてCPUのIPCのほうも動かしてやらなければならない、ということだ。IPCを動かす手段として一番手っ取り早いのはjmp命令を出力することだから、VM命令のOP_ENTERをcallした後で動的にjmpさせることでOP_ENTERと同じ動きを生成コード上で実現できる。

まずVM命令のOP_ENTERを改造して、eaxレジスタにpcをどれだけ進めたかを設定しておく。

    CASE(OP_ENTER) {
      int result; // ←追加
//(略)
        if (o == 0) {
          pc++;
          result = 0;                     // ←追加
        }
        else {
          pc += argc - m1 - m2 + 1;
          result = (argc - m1 - m2) * 10; // ←追加
        }
      }
      else {
        memmove(&regs[1], argv, sizeof(mrb_value)*(m1+o)); /* m1 + o */
        if (r) {                  /* r */
          regs[m1+o+1] = mrb_ary_new_elts(mrb, argc-m1-o-m2, argv+m1+o);
        }
        memmove(&regs[m1+o+r+1], &argv[argc-m2], sizeof(mrb_value)*m2);
        regs[len+1] = *blk; /* move block */
        pc += o + 1;
        result = o * 10;                                // ←追加
      }
      i = *pc;                                          // ←追加
      asm("movl %0, %%eax;" : : "r" (result) : "%eax"); // ←追加
      asm("addl $40, %esp;");                           // ←追加
      asm("ret $0");                                    // ←追加
      JUMP;
    }

result変数を新設して、(pcを進めた数-1)*10を格納している。なんでこんな計算をしているのかというと、とりあえずCPUのIPCは自動で増えるので-1と、あとcall命令とjmp命令のサイズがあわせて10byteだから進めるバイト数はそんだけ、ということである。
後ろのほうのインラインアセンブラはresultの値をeaxレジスタにセットする部分が増えた以外はJUMPマクロの動きと同じだ。
JUMPマクロを残してあるのはGCCがCのコードは何も無いと認識してしまうのを防ぐため。念のため。

次にコード生成部分。長いのでswitch-caseのところだけ。

      switch GET_OPCODE(irep->iseq[index]) {\
      case OP_JMP:\
      case OP_RETURN:\
        *data = 0xe9; /* jmp rel32 */ \
        data += 5;\
        break;\
      case OP_JMPIF:\
        *data = 0x0f; /* jne rel32 */ \
        *(data + 1) = 0x85;\
        data += 6;\
        break;\
      case OP_JMPNOT:\
        *data = 0x0f; /* je rel32 */ \
        *(data + 1) = 0x86;\
        data += 6;\
        break;\
      case OP_ENTER:                                   /* ←これ追加 */\
        *data = 0x05; /* addl imm32, eax */ \
        *(int**)(data + 1) = (int*)(data + 7);\
        data += 5;\
        *data = 0xff; /* jmp *eax */ \
        *(data + 1) = 0xe0;\
        data += 2;\
        break;\
      }\

OP_ENTERの処理が増えた。ここでやっているのは、OP_ENTERだったらcallした後に、eaxレジスタにセットされて戻ってきた値と次のjmp命令のアドレスをaddl命令で足してeaxに入れて、eaxが示すアドレスに間接絶対ジャンプをする、という処理だ。これによりVM命令列のOP_ENTERと同様の動きをネイティブコード側でもできるようになる。

さて、OP_RETURNとOP_ENTERが動くということは、AOベンチも動くということで、O3で最適化するとまだ動かないからO0でコンパイルして元のmrubyと比較してみた。ファイルのタイムスタンプで時間を測ってるから手書きになるけども。

元のmruby:10分8秒(608秒)
ContextThreading:9分46秒(586秒)

3%程度のしょぼい改善となった。最適化されていないコンパイルだから命令の実行時間は長くなっていて、ContextThreadingの効果は相対的に減っているはずなので、ほんとはもうちょっと速いのかもしれない。
パフォーマンスネックとして気になるのは、ContextThreadingの仕掛けではメソッド1ネストに対し、CPUレベルではcall命令が2ネストとなる部分で、call-retスタックはそんなに無いはずなので溢れてるんじゃないかと。AOベンチはネストがかなり深い。今後の性能対策としてcallしないでどうにかする、というのはありかもしれないが、実際そこまで頑張る必要があるかというとちょっと謎だ。

当面の目標であったAOベンチは完走した。続いて更にいじるとすれば、次は例外処理になるか。例外発生時にはirep間のジャンプが発生するから、これをどう実装するかは考え物だ。O3オプションで動くようにもしてみたいところ。動かない理由が全くわからないが。
でもAOベンチ動いちゃったしなあ〜そろそろ新しいことしようかな〜とか思いつつ、とりあえず一区切りということでgithubで公開しておく。続きは気が向いたら。

https://github.com/mirichi/mruby/tree/ct

コミットログを改めて見てみると、苦労の割に修正量はずいぶんと少ないのだなーて感じだった。