最初の最適化：型解析とメソッドのインライン化

前回の投稿から期間が空いてしまいました。実装にはいろいろ不十分なところも多いのですが、きりがないので記事にすることにします。

前回のベンチマークでは、JITを実装したものの、単純なJIT コンパイラではMRIのインタープリタによる実行速度に勝てないという結果でした。これに対する対処として、基本的な型解析とメソッドのインライン化の最適化を導入したというのが今回の話題です。

https://github.com/msumimz/ruby/tree/rbjit

まずは、ベンチマークの結果から。

$vc10/Release/miniruby examples/perf_while.rb
                      user     system      total        real
interprited      12.012000   0.000000  12.012000 ( 12.010687)
JIT compiled      1.186000   0.000000   1.186000 (  1.181068)

最適化の結果、インタープリタに比べて約10.1倍の実行速度になりました。

ベンチマークのソースコードは以下の通りです。timesの繰り返し回数を増やしたほかは前回のベンチマークと実質的に同じです。

# minirubyは$LOAD_PATHが空なので設定する
$LOAD_PATH << File.expand_path("../../lib", File.dirname(__FILE__))
require "benchmark"

# ベンチマークのメソッド
def m1
  i = 1
  sum = 0
  while i <= 10000
    sum += i
    i += 1
  end
  sum
end

# インタープリタで実行
puts "\t\t" + Benchmark::CAPTION
print "interprited\t"
puts Benchmark.measure { 30000.times { m1 } }

# 後で解説します
load File.expand_path("../lib/fixnum.rb", File.dirname(__FILE__))

# JITコンパイル
precompile Object, :m1

# JITコンパイルされたコードを実行
print "JIT compiled\t"
puts Benchmark.measure { 30000.times { m1 } }

最適化の内容解説

それでは、最適化の内容を説明しましょう。ベンチマークのコードは説明にはやや複雑なので、以下の単純なコードを例にします。

def m(a)
  1 + a + 2
end

その1：コントロールフローグラフの構築

最適化の前提として、Rubyのソースコードから、コントロールフローグラフ（Control Flow Graph、CFG）というものを作成します。これは、プログラムの制御フローを有向グラフとして表現したものです（そのまんますぎる）。今回の例の場合、制御フローというほどの構造はないので、CFGは以下のように単純なものになります。

f:id:msumimz:20140701221044p:plain

図の読み方ですが、ボックスがそれぞれ制御の単位を表します。今回は、entryというボックスと、exitというボックスの2つだけです。ボックス内の各行が実行すべき命令を表します。各行の最初の2つの十六進数が、それぞれ、命令アドレスと、命令の結果が代入される変数を表し、次に命令のオペコード、さらにそれに対する引数が続きます。変数欄が空の行は、命令の結果が保存されないことを表します。

xxxxxxxx xxxxxxxx CALL       xxxxxxxx (x) xxxxxxxx xxxxxxxx
アドレス 変数     オペコード 引数

現在の実装では、以下の命令があります。

命令	意味
IMM value	即値value左辺の変数に代入する。値はMRIの実装に従います。例えばFixnumの値の場合、それが表わす整数を2倍して1を足したものになります。その他、trueなら2、falseなら0、nilなら4などになります。
COPY w	変数wを左辺の変数に代入する。
LOOKUP w name	レシーバwから名前nameを持つメソッドを検索する。
CALL m (n) arg1 arg2 ... argn	メソッドmをn個の引数arg1...argnで呼び出す。arg1はレシーバです。メソッドmはLOOKUP命令で検索したものを用います。
EXIT	CFGの終了。
ENV	メソッドの実行環境を表す。今回は詳細省略。

ほかにもいくつかありますが、なんとなく雰囲気はわかっていただけると思います。

元のソースコードでは、メソッド内では変数を定義せずに1行で書いていますが、CFGに変換されると、各メソッド呼び出しの結果に対してそれぞれ変数が定義されます。例えば以下のような、より説明的なソースコードをコンパイルしても、生成されるCFGは同じになります。

def m(a)
  b = 1 + a
  c = b + 2
  return c
end

その2： SSA形式への変換

このJIT実装では、CFGの形式としてSSA形式というものを利用しています。そのため、まずこの形式に変換します。

とはいっても、今回の例では、すでにSSA形式になっているため何もする必要がないのでした。

その3：型解析

CFGの各命令の左辺の変数に対して、型を解析します。基本的な考え方は、即値（IMM）ならその値の型になり、代入（COPY）なら右辺の値になる、というように、「明らかに型がわかる変数から波及させて、順に型を推定していく」というものになります。

メソッド呼び出しの場合は、

レシーバの型を推定する。
レシーバの型とメソッド名から、呼び出され得るメソッドを検索する。
そのメソッドが返し得る型の和集合を左辺値の型とする。

という手順になります。

「メソッドが返し得る型」をどうやって取得するのかが問題になりますが、Rubyで書かれたメソッドであれば、そのメソッドを同様に型解析してやればわかります。C言語で書かれたメソッドの場合は、予め定義してやります。

例えば、Bignum#+については、setup.cppに、以下のように定義しています。

  CMethodInfo::construct(
    mri::Class::bignumClass(), "+", false,
    TypeSelection::create(TypeExactClass::create(mri::Class::fixnumClass()),
                          TypeExactClass::create(mri::Class::bignumClass())));

最後の引数であるTypeSelection::...が戻り値の型です。細かい説明は割愛しますが、FixnumまたはBignumを返すと定義されているな、となんとなく感じていただければよいと思います。

なお、解析の結果、唯一の型が決まるとは限らず、「複数の型のうちのどれか」「ある型またはその継承型のどれか」などの結果もあり得ます。そのため、以下では解析の結果のことを型制約と呼ぶことにします。「型に関する制約条件」くらいの意味です。

具体的なアルゴリズムとしては、定数伝播のアルゴリズムを変形して使っています。*1

さて、今回の例では、CFGの上から順に、以下のように変数の型制約が推定できます（結果に無関係な命令は省略）。

アドレス	左辺変数	命令	型制約
149aefe0	14aa7efe0	IMM 3	3はFixnumなので、左辺の型制約はFixnum。
14becfd8	14beefe0	LOOKUP 14a7efe0 '+'	レシーバーである変数14a7efe0はFixnumクラスのインスタンスであり、このオブジェクトから名前'+'で呼び出されるメソッドはFixnum#+のみなので、メソッドFixnum#+が型制約となる。
148e4fd8	14beafe0	CALL 14beefe0 (2) ...	メソッドFixnum#+の返し得る値はFixnumまたはBignumなので、これらのいずれかが左辺の型制約となる。
14ad4fe0	14ad6fe0	IMM 5	5はFixnumなので、左辺の型制約はFixnum。
14bdcfd8	14bf4fe0	LOOKUP 14beafe0 '+'	レシーバーである変数14beafe0はFixnumクラスまたはBignumであると推定されており、これから名前'+'で呼び出されるメソッドはFixnum#+またはBignum#+であるため、これらのいずれかが型制約となる。
13fb3fd8	14978fe0	CALL 14bf7fe0 (2) ...	メソッドFixnum#+の返し得る値はFixnumまたはBignumであり、Bignum#+の返し得る値も同じ。したがって、これらが左辺の型制約となる。
14156fe0	14a5efe0	COPY 14978fe0	代入命令の左辺の型制約は、右辺と同じになる。したがって、FixnumまたはBignum。

最後の変数14a5efe0がこのメソッドの戻り値となりますので、メソッドの戻り値の型制約はFixnumまたはBignumです。

その4：プリミティブを利用したメソッド定義

さて、このようにして、CFGの変数の型と、実行時に呼び出されるメソッドを推定することができました。しかし、MRIの実装では、Fixnum#+もBignum#+もC言語で定義されているため、これ以上の最適化はできません。そこで、少なくともFinxum#+についてRubyで定義し直し、メソッドのインライン化を実現しましょう。

事前準備として、プリミティブというものを導入します。プリミティブは、Ruby スクリプト内で、関数形式で（レシーバを指定せずに）呼び出され、インライン化されて低レベルな命令として実行される特殊な関数です。JIT コンパイルされないと使えず、インタープリタで実行するとメソッド未定義エラーになります。

実際に利用例を見てみましょう。以下は、プリミティブを使用してRuby スクリプトとして定義されたFixnum#+です。

class Fixnum
  alias :generic_add :+

  def +(other)
    if rbjit__is_fixnum(other) && rbjit__test_not(rbjit__bitwise_add_overflow(self, rbjit__bitwise_sub(other, 0)))
      sum = rbjit__bitwise_add(self, rbjit__bitwise_sub(other, 0))
      rbjit__typecast_fixnum(sum)
    else
      rbjit__typecast_fixnum_bignum(self.generic_add(other))
    end
  end
end

プリミティブは"rbjit__"で始まります。以下のプリミティブが使われています。

名前	機能
rbjit__is_fixnum(obj)	objがFixnumならtrue、そうでなければfalseを返す。
rbjit__test_not(obj)	objの真偽値を判定し、真ならfalse、偽ならtrueを返す。Rubyの'!'演算子と同じ意味ですが、'!'演算子は再定義可能なメソッドなので、プリミティブを用意しています。
rbjit__bitwise_add_overflow(a, b)	aとbを整数として加算し、オーバーフローまたはアンダーフローが発生すればtrueを、発生しなければfalseを返す。
rbjit__bitwise_add(a, b)	aとbを整数として加算し、結果を返す。
rbjit__bitwise_sub(a, b)	aとbを整数として減算し、結果を返す。
rbjit__typecast_fixnum(obj)	objがFixnum型であることを型解析器に伝える（最適化のヒント）。
rbjit__typecast_fixnum_bignum(obj)	objがFixnum型またはBignum型であることを型解析器に伝える（最適化のヒント）。

上記のFixnum#+の定義で行っていることは、「otherがFixnumであり、かつ、selfとotherを（otherのLSBをクリアして）加算してもオーバーフローが発生しないなら、otherのLSBをクリアした上で整数として加算して結果を戻り値とする。そうでなければ、通常のFixnum#+を呼び出して結果を戻り値とする」となります。

ここで、Fixnumオブジェクトは、それが表す整数値を1ビット左シフトして、LSBを立てたもので表現されることを思い出してください。例えば、rbjit__bitwise_sub(other, 0)というプリミティブ呼び出しは、FixnumであるotherオブジェクトのLSBをクリアする意味になります。Fixnumの値0をビット列で表すと0b0000_0000_0000_0001だからです。

各プリミティブは、LLVMのビットコードとして、primitives/primitives.llで定義されています。例えばrbjit__is_fixnumの定義は以下の通りです。

define %VALUE @rbjit__is_fixnum(%VALUE %v) nounwind readnone alwaysinline {
  %v1 = and %VALUE %v, 1
  %cmp = icmp eq %VALUE %v1, 1
  %ret = select i1 %cmp, %VALUE 2, %VALUE 0
  ret %VALUE %ret
}

rbjit__bitwise_addの定義は以下の通りです。

define %VALUE @rbjit__bitwise_add(%VALUE %v1, %VALUE %v2) nounwind readnone alwaysinline {
  %ret = add %VALUE %v1, %v2
  ret %VALUE %ret
}

これらは、CFGにおいてはプリミティブとして特別扱いされます。LLVMによるJIT コンパイル時にはインライン化され、他のIRと合わせて最適化されます。

上記で定義されたFixnum#+を、単独でCFGに変換すると、以下のようになります。命令PRIMがプリミティブを表します。なんとなく理解できるでしょうか。

f:id:msumimz:20140701220236p:plain

これをx86のマシン語にコンパイルすると、以下のコードになります。Rubyコードでは冗長に見えますが、それぞれのプリミティブは低レベルなので、LLVMの最適化が効果的に機能します。そのため、コンパイル結果はシンプルで、otherがFixnumである場合は高速に実行されます。

0A1B0010  push        ebx
0A1B0011  sub         esp,10h  
0A1B0014  int         3  
0A1B0015  mov         eax,dword ptr [esp+1Ch]  
0A1B0019  lea         ebx,[eax+eax]  
0A1B001C  xor         edx,edx  
0A1B001E  mov         ecx,dword ptr [esp+18h]  
0A1B0022  test        bl,2  
0A1B0025  je          0A1B0038  
0A1B002B  lea         edx,[eax-1]  
0A1B002E  add         edx,ecx  
0A1B0030  setno       dl  
0A1B0033  movzx       edx,dl  
0A1B0036  add         edx,edx  
0A1B0038  test        edx,edx  
0A1B003A  je          0A1B0049  
0A1B0040  lea         eax,[ecx+eax-1]  
0A1B0044  jmp         0A1B0065  
0A1B0049  mov         dword ptr [esp+0Ch],eax  
0A1B004D  mov         dword ptr [esp+8],ecx  
0A1B0051  mov         dword ptr [esp+4],1  
0A1B0059  mov         dword ptr [esp],13028FE0h  
0A1B0060  call        rbjit::rbjit_callMethod (0405DA8h)  
0A1B0065  add         esp,10h  
0A1B0068  pop         ebx  
0A1B0069  ret

とはいえ、かなり無駄なコードもありますね…。まあ、おいおい改良していきましょう。

なおベンチマークのソースコードにある以下の1行は、上記のような再定義されたFixnumのメソッドを読み込むためのものです。

load File.expand_path("../lib/fixnum.rb", File.dirname(__FILE__))

その5： Fixnum#+のインライン化

さて、これでFixnum#+がRuby スクリプトとして定義できました。最初の例に戻って、メソッド内の'+'メソッド呼び出しをインライン化してみましょう。

例のCFGのアドレス148e4fd8にある最初のメソッド呼び出しは、レシーバがFixnum型であり、呼び出されるメソッドはFixnum#+であることが推定されていますので、CALL命令を、単純に上記のFixnum#+のCFGに置き換えることができます。

アドレス13fb3fd8にある2番目のメソッド呼び出しは、レシーバがFixnum型またはBignum型のいずれかであり、したがってメソッドもFixum#+かBignum#+のいずれかであると推定されています。そのため、まずレシーバの型をテストして、Fixnumであると判定されたときのみインライン化するようにします。Fixnumでない場合はBignum#+を呼び出します。

インライン化された結果は以下のようになります。長いですが、Fixnum#+を2回インライン化しているだけです。変数の値がFixnumの範囲を超えない限り、メソッド呼び出しが発生しませんので高速です。

f:id:msumimz:20140701221050p:plain

レシーバの型が推定できない場合

例として示したメソッドの定義が、以下のような内容だったらどうでしょうか。

def m(a)
  a + 1 + 2
end

この場合、最初の'+'メソッドのレシーバは変数aですが、この変数の型は不明ですので、呼び出されるメソッドが特定できず、インライン化もできないことになります。

これに対応するには、2つの方法があります。

メソッド名が'+'や'-'などの場合は、レシーバはFixnumの可能性が高いと仮定して最適化してしまう。
JIT コンパイル前に、インタープリタで何度か実行してみることで実行時プロファイルを取得し、レシーバとして利用される型に関する情報を収集する。

1はすでに実装されています。そのため、実際には上記のメソッド定義でもFixnum#+はインライン化されます。レシーバがFixnumであるかどうかのテストが実行されますので、その分がオーバーヘッドになりますが、通常のメソッド検索・呼び出しのシーケンスに比べれば軽量です。

ただし、このような決めうちが通用するのは、Fixnumくらいしかありませんので、一般性はありません。とはいえ、四則演算のようなメソッドに対しては有効です。MRIでも同様の最適化が実装されています。

2は、実行時にレシーバが特定の型であるケースが多いと判明したなら、その型のメソッド呼び出しを特別扱いする方法です。

一般性のある方法は2番目になります。これを行うには、MRIの実行エンジンに手を入れて、プロファイリング用コードを埋め込む必要があります。現実のコードでは型解析がうまく行くとは限らないので、必要な機能ですが、実装はたいへんそうで、今から気が重いです。

よくある質問と答え（脳内妄想版）

Fixnumに特化した最適化に見えますが、現実のプログラムに対して効果があるでしょうか。

メソッド内部でFixnumが利用されるケースは多くあります。例えばArray#eachは、Rubyで実装すれば以下のようになるでしょう。

class Array
  def each
    i = 0
    len = self.length
    while i < len
      yield self[i]
      i += 1
    end
  end
end

Fixnumの計算が高速に実行されることで、Array#eachも効率化されます。

さらに、クロージャもCFGになってしまえばメソッドと大差ありませんので、同じようにインライン化できます（未実装ですが、将来的にはということです）。

このようにして、Array#eachメソッド自体、使われているFixnumのメソッド、クロージャがそれぞれインライン化されれば、Array#eachによるループを、メソッド呼び出しのない、C言語のfor/while文に近いコードで実現できることになります。これにより、かなりの性能向上が望めるでしょう。

効率的なJIT コンパイルのために、C言語で実装されたメソッドをRubyで実装しなおす必要があるのですか。

そうです。Rubiniusと同じです。Rubiniusに対して「Smalltalkの影響を受けて、Ruby自身でRubyを記述している」云々の紹介がされることがありますが、Rubiniusが特別のことをしているというより、JIT コンパイラを備えた言語では、普通のことだと思います。

とはいえ、すべてのメソッド定義を書き直すのは現実的ではなく、FixnumやArrayなどの、一部の主要なクラスが対象となるでしょう。

ちなみに、この点、インタープリタであるMRIとは最適化の方向性が反対になります。MRIでは、メソッド定義はCで書いた方が高速だからです（この辺り、なかなか悩ましいところです）。

JIT コンパイルにかかる時間はどのくらいですか。

ベンチマークの最後の行を以下のように変更して実行してみました。

puts Benchmark.measure { precompile Object, :m1; 30000.times { m1 } }

結果は次の通りです。さすがに1秒かかるようなタスクであれば、コンパイルにかかる時間は大きな影響がないようです（メソッドの大きさの割に呼び出し回数が多く、あまり参考になりません）。

                      user     system      total        real
interprited      11.965000   0.000000  11.965000 ( 11.980685)
compiled          1.201000   0.000000   1.201000 (  1.203068)

また、MRIはCPUのコアを一つしか使いません。これを利用して、JIT コンパイルを別スレッドで実行することで、実行時の性能コストなしでコンパイルを行うように実装することも可能と思います。

インライン化されたメソッドに対する例外のコールスタックはどうなりますか？また、Rubyの実行トレース機能は使えますか？

コンパイルされたコードのアドレスとソースコードの各行を対応付けることでき、さらに実行時にコールスタックの履歴を遡ることができれば、例外が発生しない限りコストがかからない方法で、例外のコールスタックを実行時に構築できます。が、このあたりはLLVM次第のところが大きく、あまりきちんと調べていません。MCJITならばDWARFオブジェクトを出力できるそうなので、おそらく方法はあるかと思います。

Rubyのトレース機能も、上記の情報が入手できれば、部分的には実現できなくはないです（デバッガが実現している方法）。といっても現実的には難しいので、JIT コンパイルを行う場合はトレースは無効になるという制限を設けざるを得ないと思います。

インライン化されたメソッドが再定義される場合は？

当然考慮しなくてはなりません。これに関する議論も複雑です。いずれ説明したいと思います。

まとめ

proof of conceptとしては、こんなものかと思います。ベンチマークとして数字が出ると、いろいろと最適化して速度を向上させたくなりますが、早すぎる最適化はなんとやらなので、最適化はいったん区切りをつけ、言語機能の実装を再開したいと思います。

*1:田中育男「コンパイラの構成と最適化第2版」（浅倉書店、2009）pp.367--369。ただし、型推論にこのアルゴリズムを適用すると、停止性（進行性）が保証されない気がします。そのうち検討します。

msumimz's diary

RubyにJITコンパイラを実装する個人プロジェクトの情報発信ブログです。

最初の最適化：型解析とメソッドのインライン化

最適化の内容解説

その1：コントロールフローグラフの構築

その2： SSA形式への変換

その3：型解析

その4：プリミティブを利用したメソッド定義

その5： Fixnum#+のインライン化

レシーバの型が推定できない場合

よくある質問と答え（脳内妄想版）

Fixnumに特化した最適化に見えますが、現実のプログラムに対して効果があるでしょうか。

効率的なJIT コンパイルのために、C言語で実装されたメソッドをRubyで実装しなおす必要があるのですか。

JIT コンパイルにかかる時間はどのくらいですか。

インライン化されたメソッドに対する例外のコールスタックはどうなりますか？また、Rubyの実行トレース機能は使えますか？

インライン化されたメソッドが再定義される場合は？

まとめ

最適化の内容解説

その1： コントロールフローグラフの構築

その2： SSA形式への変換

その3： 型解析

その4： プリミティブを利用したメソッド定義

その5： Fixnum#+のインライン化

レシーバの型が推定できない場合

よくある質問と答え（脳内妄想版）

Fixnumに特化した最適化に見えますが、現実のプログラムに対して効果があるでしょうか。

効率的なJITコンパイルのために、C言語で実装されたメソッドをRubyで実装しなおす必要があるのですか。

JITコンパイルにかかる時間はどのくらいですか。

インライン化されたメソッドに対する例外のコールスタックはどうなりますか？ また、Rubyの実行トレース機能は使えますか？

インライン化されたメソッドが再定義される場合は？

まとめ

その1：コントロールフローグラフの構築

その3：型解析

その4：プリミティブを利用したメソッド定義

効率的なJIT コンパイルのために、C言語で実装されたメソッドをRubyで実装しなおす必要があるのですか。

JIT コンパイルにかかる時間はどのくらいですか。

インライン化されたメソッドに対する例外のコールスタックはどうなりますか？また、Rubyの実行トレース機能は使えますか？