mrubyのコードジェネレータをいじる

先日のうちで動かしたベンチだが、bench2と3の場所を入れ替えたら2のほうが1.5倍ぐらい遅くなった。2と3はアセンブラレベルで命令は同じ(case0のaに代入する値が変わるだけ)なので、どうにもさっぱりわからない。mruby本体に手を入れると速くなったり遅くなったりする。現象からすると命令が配置されるアドレスに問題がある可能性があるが、それが原因だとすると対策が思いつかない。どーしてくれんだAMD。とりあえずその手の最適化系は原因が判明するまで封印しようと思う。

さて、mrubyの高速化をあれこれと試しているわけなのだが、基本的に高速化の戦略は以下の3つにわかれる。
(a)VMGCといった基本部分。コア部分になるのでここを速くできれば単純に速くなる。これぞVMで動作する言語処理系の華だ。
(b)組み込みライブラリ。何気に重要なところで、Arrayとかはかなり使うものだし、そういった部分が速くなればおいしい。
(c)コードジェネレータの出力。あまり気にされていないが(見えないし)、俺的価値観ではここの改善は(a)に匹敵するほど大事。

そういうことで、今回はコードジェネレータを真面目にいじってみる。難しくてよくわかんないのでひたすら追っかけてみよう。とても長い記事になる。mrubyのコードジェネレータに興味がある人はどうぞ。


いまやmrubyのbenchmarkディレクトリに収まっているAOベンチ。手元にあるまともなテストコードはこれとbm_so_lists.rbしかないのでとりあえずAOベンチを参考にしよう。基本的には--verboseして無駄なところを探し、無駄なコードが出ないようにコードジェネレータを修正する方針。
テキトーに眺めていると、このようなコードが見つかった。

irep 152 nregs=7 nlocals=6 pools=0 syms=3
000 OP_ENTER	3:0:0:0:0:0:0
001 OP_SETIV	@x	R1
002 OP_SETIV	@y	R2
003 OP_MOVE	R6	R3
004 OP_SETIV	@z	R6
005 OP_RETURN	R6

irep 153 nregs=5 nlocals=4 pools=0 syms=1
000 OP_ENTER	1:0:0:0:0:0:0
001 OP_MOVE	R4	R1
002 OP_SETIV	@x	R4
003 OP_RETURN	R4

これはVecクラスのinitializeとx=だ。対応するRubyコードは

  def initialize(x, y, z)
    @x = x
    @y = y
    @z = z
  end

  def x=(v); @x = v; end

となる。どこが無駄かというと、最後のOP_SETIVの前にOP_MOVEが出力されている。これはなくてもいいだろう。initializeの例で言うならOP_MOVEを消してOP_SETIV @z R3とOP_RETURN R3になれば1命令減る。
そもそもなんでAOベンチではattr_accessorを使わずに自前で定義しているのかというと、mrblib(mruby組み込みライブラリの一部でRubyで記述されたもの)の中に以下のようなコードがあって、

class Module
  # 15.2.2.4.13
  def attr_reader(*names)
    names.each{|name|
      name2 = ('@'+name.to_s).intern
      define_method(name){self.instance_variable_get(name2)}
    }
  end
#(略)

つまり、self.instance_variable_getを呼ぶように作られていて、自前で定義したほうが速いと言っているのではないかと思っている。ここも何とかできればいいのだが。

さて、どう直すかというと、とりあえずさっぱりわからないのでまず仮説。initializeの例を見るに、OP_RETURN直前のOP_SETIVだけOP_MOVEがあるので、メソッドの戻り値になる操作の場合になにかしら特殊なことをしているのではないか。であればそこだけいじればよい。
codegen.cはNODE_xxx別にわかれているので、該当するNODEのタイプがわからないと追いかけようがない。--verboseの出力から該当する部分を抜き出してみよう。ここだ。

    NODE_CLASS:
      :Vec
      body:
        NODE_BEGIN:
          NODE_DEF:
            initialize
            local variables:
              x               y               z 
            mandatory args:
              NODE_ARG x
              NODE_ARG y
              NODE_ARG z
            NODE_BEGIN:
              NODE_ASGN:
                lhs:
                  NODE_IVAR @x
                rhs:
                  NODE_LVAR x
              NODE_ASGN:
                lhs:
                  NODE_IVAR @y
                rhs:
                  NODE_LVAR y
              NODE_ASGN:
                lhs:
                  NODE_IVAR @z
                rhs:
                  NODE_LVAR z
          NODE_DEF:
            x=
            local variables:
              v 
            mandatory args:
              NODE_ARG v
            NODE_BEGIN:
              NODE_ASGN:
                lhs:
                  NODE_IVAR @x
                rhs:
                  NODE_LVAR v

これは構文木で、Rubyのソースをパースして構造を解析した結果である。一番上のNODE_CLASSはクラス定義でその下がクラス名のシンボル、bodyはクラス定義本体、NODE_BEGINはその下の処理をまとめるノードで、たぶんbodyには1つのノードしか入れられないのだろう。NODE_DEFはメソッドの定義。メソッド名、ローカル変数、引数、メソッドの中身のNODE_BEGINと続く。
ではNODE_BEGINあたりから見ていこう。NODE_xxxの処理がされているのはcodegen関数である。

static void
codegen(codegen_scope *s, node *tree, int val)
{
  int nt;

  if (!tree) return;
  nt = (intptr_t)tree->car;
  tree = tree->cdr;
  switch (nt) {
  case NODE_BEGIN:
    while (tree) {
      codegen(s, tree->car, tree->cdr ? NOVAL : val);
      tree = tree->cdr;
    }
    break;

  case NODE_RESCUE:
    {

いきなり最初にあった。case NODE_BEGINのところだ。treeがノードの中身をあらわしていて、carは最初の1つ、cdrは残り。NODE_BEGINの中にはNODE_DEFがいっぱい並んでいたので、それを1つずつ取り出してcodegen再帰呼び出ししているわけだ。
気になるのはtree->cdr ? NOVAL : valというところ。検索するとNOVALは0、VALは1とマクロで定義されている。cdrがNULLかどうかを判定しているので、最後の一つかどうかという意味だ。最後の一つじゃなければNOVALを渡すが、最後の場合は引数のvalを渡す、となっている。いきなり怪しいものを見つけた。
では次、NODE_DEF。

  case NODE_DEF:
    {
      int sym = new_msym(s, (mrb_sym)tree->car);
      int idx = lambda_body(s, tree->cdr, 0);

      genop(s, MKOP_A(OP_TCLASS, cursp()));
      push();
      genop(s, MKOP_Abc(OP_LAMBDA, cursp(), idx, OP_L_METHOD));
      pop();
      genop(s, MKOP_AB(OP_METHOD, cursp(), sym));
      if (val) {
        genop(s, MKOP_A(OP_LOADNIL, cursp()));
      }
    }
    break;

NODE_DEF内の1個目はメソッド名でsymに入れている。んで残りをlambda_bodyに渡す。その後ろはクラスを取り出してOP_LAMBDAしてOP_METHODなのでメソッドの中身を作り終わってからそれを定義する処理である。ここにもvalがいる。
lambda_body。

static int
lambda_body(codegen_scope *s, node *tree, int blk)
{
  int idx, base = s->idx;
  mrb_code c;

  s = scope_new(s->mrb, s, tree->car);
  idx = s->idx;
  s->mscope = !blk;

  if (blk) {
    struct loopinfo *lp = loop_push(s, LOOP_BLOCK);
    lp->pc1 = new_label(s);
  }
  tree = tree->cdr;
  if (tree->car) {

//(略)

  }
  codegen(s, tree->cdr->car, VAL);
  pop();
  c = s->iseq[s->pc-1];
  if (GET_OPCODE(c) != OP_RETURN || GETARG_B(c) != OP_R_NORMAL || s->pc == s->lastlabel) {
    if (s->nregs == 0) {
      genop(s, MKOP_A(OP_LOADNIL, 0));
      genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
    }
    else {
      genop_peep(s, MKOP_AB(OP_RETURN, cursp(), OP_R_NORMAL), NOVAL);
    }
  }

略した部分は引数の処理だ。引数は省略できるのでそのあたりが面倒なことになっているが今回は関係ないので無視。
引数処理の直後のcodegenが、メソッド本体のNODE_BEGINを生成するところだ。VALを固定で渡している。が、NODE_BEGINはさっき見たように最後じゃなければNOVALに置き換えてしまう。
んで、その後ろにあるif文がどうやら今回の話に関係しそうなOP_RETURNの判定をしていて、これはつまり最後にユーザーが自前でreutnを書いたかどうかを見ている。ifの中は書かれていなかった場合となる。nregsが0という条件はメソッド内でスタックの使用が0だったと言う意味で、その場合返す値が無いのでOP_LOADNILを生成してOP_RETURNして終わる。スタックを使っていたらgenop_peepでピープホール最適化をしているが、OP_RETURNを出力するコードが無いのでこの中で出しているようだ。ここは後で見る。
では次、NODE_BEGINの中はNODE_ASGNなのでそれを見てみよう。NODE_ASGNは=で代入するときのノードのようだ。

  case NODE_ASGN:
    codegen(s, tree->cdr, VAL);
    pop();
    gen_assignment(s, tree->car, cursp(), val);
    break;

先にcdrで=の右側をVALで生成。右側はNODE_LVARだった。

  case NODE_LVAR:
    if (val) {
      int idx = lv_idx(s, (mrb_sym)tree);

      if (idx > 0) {
        genop(s, MKOP_AB(OP_MOVE, cursp(), idx));
      }
      else {
        int lv = 0;
        codegen_scope *up = s->prev;

        while (up) {
          idx = lv_idx(up, (mrb_sym)tree);
          if (idx > 0) {
            genop(s, MKOP_ABC(OP_GETUPVAR, cursp(), idx, lv));
            break;
          }
          lv++;
          up = up->prev;
        }
      }
      push();
    }
    break;

全体がif(val)で囲まれているのでNOVAL(最後じゃない)の場合は何も行われないことになる。これはたぶん、コードの中に数字だけ書いたりした場合に命令を省略しようとしているのではなかろうか。逆に、最後の場合は戻り値になるので省略されると困る。NODE_ASGNから飛んできた場合はVAL固定なので最後じゃなくても出力される。
lv_idxのlvはレベルじゃなくてローカル変数のことだ。ローカル変数の値がLVARだったらOP_MOVEでスタックトップのレジスタに格納する。どうやらこれが無駄に出ているOP_MOVEの正体のようだ。
このスコープのローカル変数じゃなかったらブロックの外側を辿って変数を探す処理が続く。変数が無い場合はここに来る前にソースをパースしている時点でエラーになるのでここでは想定していない。
NODE_LVARの処理が終わったらNODE_ASGNに戻ってgen_assignment()が呼ばれる。引数はcarなのでNODE_IVAR、格納先だ。あと上から渡されたvalも。
長いので削る。

static void
gen_assignment(codegen_scope *s, node *node, int sp, int val)
{
  int idx;
  int type = (intptr_t)node->car;

  node = node->cdr;
  switch ((intptr_t)type) {
  case NODE_GVAR:

//(略)

  case NODE_IVAR:
    idx = new_sym(s, (mrb_sym)node);
    genop_peep(s, MKOP_ABx(OP_SETIV, sp, idx), val);
    break;
  case NODE_CVAR:

//(略)

    break;
  }
  if (val) push();
}

NODE_IVARの場合、格納先のシンボルを取得して、genop_peepを呼ぶ。また出た。今度はここに入る。

static void
genop_peep(codegen_scope *s, mrb_code i, int val)
{
  /* peephole optimization */
  if (!val && s->lastlabel != s->pc && s->pc > 0) {
    mrb_code i0 = s->iseq[s->pc-1];
    int c1 = GET_OPCODE(i);
    int c0 = GET_OPCODE(i0);

    switch (c1) {
    case OP_MOVE:

//(略)

    case OP_SETIV:
    case OP_SETCV:
    case OP_SETCONST:
    case OP_SETMCNST:
      if (c0 == OP_MOVE) {
        if (GETARG_A(i) == GETARG_A(i0)) {
          s->iseq[s->pc-1] = MKOP_ABx(c1, GETARG_B(i0), GETARG_Bx(i));
          return;
        }
      }
      break;


//(略)
}

頭のところでvalを見て最後の1つの場合は最適化しない、となっている。lambda_bodyでVAL(最後)を渡していたが、NODE_BEGINで最後じゃない場合はNOVALに差し替えられているので、最後じゃないNODE_IVARはここで処理がされる。
ここの処理はいま出力しようとしている命令がc1に、1つ前の命令がc0に入る。出力しようとしている命令がOP_SETIVなどの場合、1つ前がOP_MOVEだったら、OP_MOVEの場所に命令を上書きしてOP_MOVEを消す。この処理でOP_SETIVの前のOP_MOVEが消えているわけだ。逆に言えば、最後の1つがVALでgenop_peepに入ってきた場合、OP_MOVEは消えない。これを消せばよい。
とはいえ消すのは簡単だが消してしまうとreturnする際に返す値がスタックトップではなくなるので、まずはそこを細工する必要がある。

コードはしばらく戻ってlambda_bodyの最後のところを再掲。

  if (GET_OPCODE(c) != OP_RETURN || GETARG_B(c) != OP_R_NORMAL || s->pc == s->lastlabel) {
    if (s->nregs == 0) {
      genop(s, MKOP_A(OP_LOADNIL, 0));
      genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
    }
    else {
      genop_peep(s, MKOP_AB(OP_RETURN, cursp(), OP_R_NORMAL), NOVAL);
    }
  }

最後の命令がreturnじゃない場合、スタックが使われていたらgenop_peepをNOVALで呼ぶ。return時の細工をしているらしいところだ。われわれはついに前人未到の奥地に到達したのだー。

//(略)

    case OP_RETURN:
      switch (c0) {
      case OP_MOVE:
        s->iseq[s->pc-1] = MKOP_AB(OP_MOVE, 0, GETARG_B(i0));
        genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
        return;
      case OP_LOADI:
        s->iseq[s->pc-1] = MKOP_AsBx(OP_LOADI, 0, GETARG_sBx(i0));
        genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
        return;
      case OP_ARRAY:
      case OP_HASH:
      case OP_RANGE:
      case OP_AREF:
      case OP_GETUPVAR:
        s->iseq[s->pc-1] = MKOP_ABC(c0, 0, GETARG_B(i0), GETARG_C(i0));
        genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
        return;
      case OP_LOADSYM:
      case OP_GETGLOBAL:
      case OP_GETIV:
      case OP_GETCV:
      case OP_GETCONST:
      case OP_GETSPECIAL:
      case OP_LOADL:
      case OP_STRING:
        s->iseq[s->pc-1] = MKOP_ABx(c0, 0, GETARG_Bx(i0));
        genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
        return;
      case OP_SCLASS:
        s->iseq[s->pc-1] = MKOP_AB(c0, GETARG_A(i), GETARG_B(i0));
        genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
        return;
      case OP_LOADNIL:
      case OP_LOADSELF:
      case OP_LOADT:
      case OP_LOADF:
      case OP_OCLASS:
        s->iseq[s->pc-1] = MKOP_A(c0, 0);
        genop(s, MKOP_AB(OP_RETURN, 0, OP_R_NORMAL));
        return;
      default:
        break;
      }
      break;
    default:
      break;
    }
  }
  genop(s, i);
}

これは今回の命令がOP_RETURNだった場合、一つ前の命令を見て、レジスタR0に入れるようにしてOP_RETURN R0とするロジックだ。例えば最後がOP_GETIVだった場合、スタックトップのレジスタに入れてからOP_RETURNしなくても、R0に入れてそれを返せばいいじゃーん、ということなのだろう。んー、スタックの節約のため?
なんしか、ここが最後のOP_RETURNを生成するところなのだから、ここでOP_SETIVを判定してインスタンス変数に入れるレジスタをRETURNに使うように変えれば、最後のOP_MOVEは消しておいても大丈夫となる。
こうだ。

      case OP_SETIV:
        genop(s, MKOP_AB(OP_RETURN, GETARG_A(s->iseq[s->pc-1]), OP_R_NORMAL));
        return;

んでもってOP_MOVEを消しにかかる。これを消すには最後のNODE_IVARがVALにならないようにすればよい。修正するのはgen_assignmentの中だ。

  case NODE_IVAR:
    idx = new_sym(s, (mrb_sym)node);
//    genop_peep(s, MKOP_ABx(OP_SETIV, sp, idx), val);
    genop_peep(s, MKOP_ABx(OP_SETIV, sp, idx), NOVAL);
    break;

どうだろう。
該当箇所はこのようになった。

irep 152 nregs=7 nlocals=6 pools=0 syms=3
000 OP_ENTER	3:0:0:0:0:0:0
001 OP_SETIV	@x	R1
002 OP_SETIV	@y	R2
003 OP_SETIV	@z	R3
004 OP_RETURN	R3

irep 153 nregs=5 nlocals=4 pools=0 syms=1
000 OP_ENTER	1:0:0:0:0:0:0
001 OP_SETIV	@x	R1
002 OP_RETURN	R1

ブラボー。