JIT Compile - Just In Time コンパイル

はじめに

「見た目は JavaScript、頭脳(中身)は Ruby、（安定感は AC/DC）」 でお届けしているスクリプト言語 Kinx。作ったものの紹介だけではなく実現のために使った技術を紹介していくのも貢献。その道の人には当たり前でも、そうでない人にも興味をもって貰えるかもしれない。

前回のテーマは VM (Virtual Machine)。今回のテーマは JIT。

参考
- 最初の動機 ... スクリプト言語 KINX（ご紹介）
  - 個別記事へのリンクは全てここに集約してあります。
- リポジトリ ... https://github.com/Kray-G/kinx
  - Pull Request 等お待ちしております。

正直、JIT は道半ばでコード晒すのもイマイチだが、参考にはなるかと思う。コントリビューター募集中というのもあるので興味があればよろしくお願いします。

追記：
何気にだいぶアップデートして、出力コードの質はここの記事よりもだいぶ良くなってる（はず）。あと、Switch-Case もサポートしたし、オブジェクトへのアクセスも一部だができるようになっているので、この記事はちょっと古い情報です。どこかで最新にアップデートするか、別記事でフォローします。

JIT コンパイラとは

Just In Time、つまり まさにその場で コンパイルして実行する方式。ネイティブコード（機械語）にコンパイルして実行するため実行速度を追い求める人たちの間では欲しくてたまらない機能の一つ。

ただし、プロセッサーごとに個別に実装しないといけない。x64 だけ、とかならまだいいんだけど、どっかで ARM くらいはサポートすっかなー、とか、PowerPC は、SPARC?、とかになると、そもそもテストするのが大変だ。

Ruby 方式

Ruby も最近 JIT サポートをアピールしはじめたが、その辺の大変さからうまく？逃げている。つまり、C ソースを作成して、他のコンパイラでコンパイル したものを自分のプロセスに取り込む、といったことをしている。コンパイラがホストコンピュータ用のコードを出力してくれるし、もともと他人のモノなので Ruby 自体はプロセッサーを意識しなくて良い。

まあ逃げ方としてはそうなるとは思うが、しかしそれでは実行環境にコンパイラが必要になってしまう。しかもなにより Windows 環境で Visual Studio をインストールしてください、とか、Cygwin インストールしてください、とは 言えない というのが私の感覚です。開発者以外にコンパイラをインストールさせるのは ... かなり気が引ける。

sljit

そこで sljit を採用した。sljit はアセンブラを抽象化しており、その抽象化アセンブラで書くことでコード出力時に x64 や ARM などに合わせて出力ができるようになっている。その代わり、レジスタ数とかどのプロセッサーでも共通に使えるようにミニマムに合わせておかなければならない。

ただし、x64 以外ではテストしていないので悪しからず。

尚、ネイティブ・コード用 IR 出力ルーチンは src/ast_native.c にあり、sljit を使ったネイティブ・コード用コンパイラは src/nir_compile.c にある。そして、まだ実装されていないものが沢山あることがわかる。。。

個別の実装（特徴）

いくつか個別の実装の内容を記載してみる。まだまだやらなきゃいけないことが沢山あるなあ。

IR (Intermediate Representation)

ネイティブ・コードに変換する上で扱いやすいように、VM 用の IR と異なるコードを出力する。いやいや本当に違うものを出力する。なぜなら、例外の扱いとかコンパイル時点における型チェックの厳密さとかが全く異なるからで、それに合わせて違う体系になっている。

例えば；

native nfib(n) {
    if (n < 3) return n;
    return nfib(n-2) + nfib(n-1);
}

function fib(n) {
    if (n < 3) return n;
    return fib(n-2) + fib(n-1);
}

これは以下のようになる。

nfib(1, 10):
  .L0
       1:   jmp                     .L1
  .L1
       2:   load                    r1, $(0,0)
       3:   lt                      r2 = r1, 3
       4:   jz                      .L4
  .L2
       5:   load                    r3, $(0,0)
       6:   ret                     r3
  .L3
       8:   jmp                     .L4
  .L4
       9:   load                    r4, $(0,0)
       a:   sub                     r5 = r4, 2
       b:   arg                     r5
       c:   call                    r6 = [rec-call]()
       d:   excpt                   r7 = check
       e:   jz                      .L6
  .L5
       f:   excpt                   off
      10:   excpt                   on
      11:   ret                     0
  .L6
      13:   load                    r7, $(0,0)
      14:   sub                     r8 = r7, 1
      15:   arg                     r8
      16:   call                    r9 = [rec-call]()
      17:   excpt                   r10 = check
      18:   jz                      .L8
  .L7
      19:   excpt                   off
      1a:   excpt                   on
      1b:   ret                     0
  .L8
      1d:   add                     r10 = r6, r9
      1e:   ret                     r10
  .L9
      20:   ret                     0

fib:
  .L460
     d06:   enter                   23, vars(1), args(1)
  .L461
     d07:   lt_v0i                  $0(0), 3
     d08:   jz                      .L463(d0a)
  .L462
     d09:   retvl0                  $0(0)
  .L463
     d0a:   pushvl0                 $0(0)
     d0b:   subi                    2
     d0c:   callvl1                 $1(34), 1
     d0d:   pushvl0                 $0(0)
     d0e:   subi                    1
     d0f:   callvl1                 $1(34), 1
     d10:   add
     d11:   ret
     d12:   halt

長くなるので各命令の意味は省略します。ちょっと無駄なコードも見え隠れ（隠れてない）するが、こんな感じの出力を現時点ではする。しかし、こうやって書くと改善点も客観的に見えるというメリットもあるね。

例外処理

ネイティブコードでも try-catch-finally をサポートするにあたり、多少工夫が必要。setjmp/longjmp が使えれば簡単だが、JIT コード内で、しかも複数プロセッサー対応を目指した抽象化アセンブラの範囲内では使えないので実現するのはちょっと骨が折れた。

具体的には以下の通りに実装されている。

例外が発生したかをチェックできる特別なフラグを用意する。
例外発生時、フラグを立て、キャッチ・ルーチンがあればそこにジャンプする。
finally があればそれを実行する。
同じ関数内に次の（外側の）キャッチ・ルーチンがあれば、そこにジャンプ。
無ければ return する。
リターンは関数呼び出しの次のコードなので、そこには例外が発生したかのフラグを確認するコードを出力しておく。
例外発生を検知したら、上記と同じ形で伝播させていく。

スタックトレースも作れるはずなのだが、他のことを優先しているので今はできていない。どっかできちんとやる。

型チェック

ネイティブコードでは型チェックは重要。関数の入口で引数の型をチェック。レキシカル変数をアクセスする場合もチェック。関数から帰ってきたときにもチェック。型が期待するものと違ってたらすぐさま例外を上げるようになっている。

Switch-Case

sljit のライブラリを色々見たのだが、sljit の世界の中でどうしてもジャンプテーブルの実現方法が分からなかったので、今はサポートできていない（具体的にはレジスタ値にジャンプする方法がインターフェースとして用意されていない模様）。

最悪は if-else で実現するくらいだが、ジャンプテーブルの無い Switch-Case は軽自動車並みのエンジンで見た目だけ F1 カーといった感じなので、やる意味があるかどうか。sljit を改造するという手もあるが、x64 以外はすぐに対応できそうもない、といったところ。

VM 関数呼び出し

これも現在は未サポート。子 VM を起動してジャンプすれば良いとは思うが、入れ子になった VM 内で例外が発生した場合の処理が大変そうなので、まだ手がつけられていない。おそらく以下のような感じでできるのでは、と思っているが、どうだろう。

フレームに native からのコールであるフラグを立てる。
return 時、native からであれば復帰値をセットして　VM から return する。
例外発生時、スタックを巻き戻している最中に native コールされたフラグを見つけた場合、そこでスタックトレース情報の作成を中断し、native 向けの例外情報を設定して native にリターン。native 側に帰ってきたら、native での例外伝搬の仕組みで伝搬させていく。

他に気にするところはあるかな。例外のスタックトレースを native 関数内でも保存するようにしないといけないな。やってみる時間があれば挑戦してみる、という予定。

オブジェクトへのアクセス

レキシカル変数へのアクセスはできるので、できるはずだがこれもまだ未サポート。演算で認められていない型同士の演算をきちんと拒絶、もしくはきちんとキャストするようにした上で、オブジェクトへのインデックス・アクセスをできるようにすれば良いだけだが、先にテストコードとドキュメントとライブラリを充実させたいもので。

そのうちやる。

レジスタ割り付け

複数のプロセッサーに対応させようとすると、レジスタ数の縛りが結構キツイ。なので割り切ってレジスタ割り付けは省略。全部ストア・ロードで対応。出力コードは汚いが、マイクロベンチマークである程度パフォーマンスが出ているので、良しとした。

尚、最初の例ではあえて載せていなかったが、x64 向けコードを出力させてみると以下のようになる。うむ、汚い。大半は型チェックと例外チェックなのだが、演算内容に対しても全くレジスタ割り付けしていないのも分かるだろう。毎回ストア・ロードしているし、無駄にストアしたのを再度ロードしたりしているのでその辺は簡単に改善できそうではある。今のところ重要度は高くないのでペンディング中。return くらいは一か所に集めるか。

nfib: (native-base:0x7f1b2aa50010)
       0:   53                                push    rbx
       1:   41 57                             push    r15
       3:   41 56                             push    r14
       5:   41 55                             push    r13
       7:   55                                push    rbp
       8:   41 54                             push    r12
       a:   48 8b df                          mov     rbx, rdi
       d:   4c 8b fe                          mov     r15, rsi
      10:   4c 8b f2                          mov     r14, rdx
      13:   48 81 ec 68 02 00 00              sub     rsp, 0x268
      1a:   49 8b 47 08                       mov     rax, qword [r15+0x00000008]
      1e:   48 83 c0 01                       add     rax, 0x01
      22:   49 89 47 08                       mov     qword [r15+0x00000008], rax
      26:   48 3d 00 04 00 00                 cmp     rax, 0x400
      2c:   72 29                             jb      0x00000057
      2e:   48 c7 43 20 01 00 00 00           mov     qword [rbx+0x00000020], 0x01
      36:   48 c7 43 28 06 00 00 00           mov     qword [rbx+0x00000028], 0x06
      3e:   48 c7 c0 00 00 00 00              mov     rax, 0x00
      45:   48 81 c4 68 02 00 00              add     rsp, 0x268
      4c:   41 5c                             pop     r12
      4e:   5d                                pop     rbp
      4f:   41 5d                             pop     r13
      51:   41 5e                             pop     r14
      53:   41 5f                             pop     r15
      55:   5b                                pop     rbx
      56:   c3                                ret
      57:   49 83 bf 10 01 00 00 01           cmp     qword [r15+0x00000110], 0x01
      5f:   0f 85 52 02 00 00                 jnz     0x000002b7
      65:   49 8b 57 10                       mov     rdx, qword [r15+0x00000010]
      69:   48 89 14 24                       mov     qword [rsp], rdx
      6d:   eb 00                             jmp     0x0000006f
      6f:   48 8b 14 24                       mov     rdx, qword [rsp]
      73:   48 89 94 24 18 02 00 00           mov     qword [rsp+0x00000218], rdx
      7b:   48 c7 84 24 20 02 00 00 01 00     mov     qword [rsp+0x00000220], 0x01
      87:   00 00 48 83 bc 24 18 02 00        cmp     qword [rsp+0x00000218], 0x03
      90:   00 03                             jb      0x0000009e
      92:   72 0c 48 c7 84 24 20 02 00 00     mov     qword [rsp+0x00000220], 0x00
      9e:   00 00 00 00 48 83 bc 24 20        cmp     qword [rsp+0x00000220], 0x00
      a7:   02 00                             jz      0x000000d3
      a9:   00 00 74 2a                       mov     rdx, qword [rsp]
      ad:   48 8b 14 24 48 89 94 24           mov     qword [rsp+0x00000228], rdx
      b5:   28 02 00 00 48 8b 84 24           mov     rax, qword [rsp+0x00000228]
      bd:   28 02 00 00 48 81 c4              add     rsp, 0x268
      c4:   68 02                             pop     r12
      c6:   00                                pop     rbp
      c7:   00 41                             pop     r13
      c9:   5c 5d                             pop     r14
      cb:   41 5d                             pop     r15
      cd:   41                                pop     rbx
      ce:   5e                                ret
      cf:   41 5f                             jmp     0x0000006d
      d1:   5b c3                             jmp     0x000000d3
      d3:   eb 9c eb 00                       mov     rdx, qword [rsp]
      d7:   48 8b 14 24 48 89 94 24           mov     qword [rsp+0x00000230], rdx
      df:   30 02 00 00 48 8b 94 24           mov     rdx, qword [rsp+0x00000230]
      e7:   30 02 00 00                       sub     rdx, 0x02
      eb:   48 83 ea 02 48 89 94 24           mov     qword [rsp+0x00000238], rdx
      f3:   38 02 00 00 48 8b 94 24           mov     rdx, qword [rsp+0x00000238]
      fb:   38 02 00 00 48                    mov     qword [rsp+0x00000018], rdx
     100:   89 54 24 18 48 c7 84 24 18 01     mov     qword [rsp+0x00000118], 0x01
     10c:   00 00 01 00                       mov     rcx, qword [rbx+0x00000018]
     110:   00 00 48                          mov     rax, rbx
     113:   8b 4b 18 48                       mov     rdx, qword [r15+0x00000008]
     117:   89 d8 49 8b 57                    mov     qword [rsp+0x00000010], rdx
     11c:   08 48 89 54 24                    lea     rsi, [rsp+0x00000008]
     121:   10 48 8d                          mov     rdi, rax
     124:   74 24                             call    rcx
     126:   08 48 89 c7 ff d1 48 89           mov     qword [rsp+0x00000240], rax
     12e:   84 24 40 02 00 00 48 c7 84 24     mov     qword [rsp+0x00000248], 0x01
     13a:   48 02 00 00                       mov     rax, qword [rbx+0x00000020]
     13e:   01 00 00 00                       cmp     rax, 0x00
     142:   48 8b                             jnz     0x00000150
     144:   43 20 48 83 f8 00 75 0c 48 c7     mov     qword [rsp+0x00000248], 0x00
     150:   84 24 48 02 00 00 00 00 00        cmp     qword [rsp+0x00000248], 0x00
     159:   00 48 83 bc 24 48                 jz      0x0000018e
     15f:   02 00 00 00 0f 84 2f 00           mov     qword [rbx+0x00000020], 0x00
     167:   00 00 48 c7 43 20 00 00           mov     qword [rbx+0x00000020], 0x01
     16f:   00 00 48 c7 43 20 01 00           mov     rax, qword [rsp+0x00000210]
     177:   00 00 48 8b 84 24 10              add     rsp, 0x268
     17e:   02 00                             pop     r12
     180:   00                                pop     rbp
     181:   48 81                             pop     r13
     183:   c4 68                             pop     r14
     185:   02 00                             pop     r15
     187:   00                                pop     rbx
     188:   41                                ret
     189:   5c 5d 41 5d 41                    jmp     0x0000006d
     18e:   5e 41 5f 5b                       mov     rdx, qword [rsp]
     192:   c3 e9 df fe ff ff 48 8b           mov     qword [rsp+0x00000248], rdx
     19a:   14 24 48 89 94 24 48 02           mov     rdx, qword [rsp+0x00000248]
     1a2:   00 00 48 8b                       sub     rdx, 0x01
     1a6:   94 24 48 02 00 00 48 83           mov     qword [rsp+0x00000250], rdx
     1ae:   ea 01 48 89 94 24 50 02           mov     rdx, qword [rsp+0x00000250]
     1b6:   00 00 48 8b 94                    mov     qword [rsp+0x00000018], rdx
     1bb:   24 50 02 00 00 48 89 54 24 18     mov     qword [rsp+0x00000118], 0x01
     1c7:   48 c7 84 24                       mov     rcx, qword [rbx+0x00000018]
     1cb:   18 01 00                          mov     rax, rbx
     1ce:   00 01 00 00                       mov     rdx, qword [r15+0x00000008]
     1d2:   00 48 8b 4b 18                    mov     qword [rsp+0x00000010], rdx
     1d7:   48 89 d8 49 8b                    lea     rsi, [rsp+0x00000008]
     1dc:   57 08 48                          mov     rdi, rax
     1df:   89 54                             call    rcx
     1e1:   24 10 48 8d 74 24 08 48           mov     qword [rsp+0x00000258], rax
     1e9:   89 c7 ff d1 48 89 84 24 58 02     mov     qword [rsp+0x00000260], 0x01
     1f5:   00 00 48 c7                       mov     rax, qword [rbx+0x00000020]
     1f9:   84 24 60 02                       cmp     rax, 0x00
     1fd:   00 00 01 00 00 00                 jnz     0x0000020f
     203:   48 8b 43 20 48 83 f8 00 0f 85     mov     qword [rsp+0x00000260], 0x00
     20f:   0c 00 00 00 48 c7 84 24 60        cmp     qword [rsp+0x00000260], 0x00
     218:   02 00 00 00 00 00                 jz      0x0000024d
     21e:   00 48 83 bc 24 60 02 00           mov     qword [rbx+0x00000020], 0x00
     226:   00 00 0f 84 2f 00 00 00           mov     qword [rbx+0x00000020], 0x01
     22e:   48 c7 43 20 00 00 00 00           mov     rax, qword [rsp+0x00000210]
     236:   48 c7 43 20 01 00 00              add     rsp, 0x268
     23d:   00 48                             pop     r12
     23f:   8b                                pop     rbp
     240:   84 24                             pop     r13
     242:   10 02                             pop     r14
     244:   00 00                             pop     r15
     246:   48                                pop     rbx
     247:   81                                ret
     248:   c4 68 02 00 00                    jmp     0x0000006d
     24d:   41 5c 5d 41 5d 41 5e 41           mov     rdx, qword [rsp+0x00000240]
     255:   5f 5b c3 e9 20 fe ff ff           add     rdx, qword [rsp+0x00000258]
     25d:   48 8b 94 24 40 02 00 00           mov     qword [rsp+0x00000260], rdx
     265:   48 03 94 24 58 02 00 00           mov     rax, qword [rsp+0x00000260]
     26d:   48 89 94 24 60 02 00              add     rsp, 0x268
     274:   00 48                             pop     r12
     276:   8b                                pop     rbp
     277:   84 24                             pop     r13
     279:   60 02                             pop     r14
     27b:   00 00                             pop     r15
     27d:   48                                pop     rbx
     27e:   81                                ret
     27f:   c4 68 02 00 00                    jmp     0x0000006d
     284:   41 5c 5d 41 5d 41 5e 41           mov     rax, qword [rsp+0x00000210]
     28c:   5f 5b c3 e9 e9 fd ff              add     rsp, 0x268
     293:   ff 48                             pop     r12
     295:   8b                                pop     rbp
     296:   84 24                             pop     r13
     298:   10 02                             pop     r14
     29a:   00 00                             pop     r15
     29c:   48                                pop     rbx
     29d:   81                                ret
     29e:   c4 68 02 00 00 41 5c              mov     rax, 0x00
     2a5:   5d 41 5d 41 5e 41 5f              add     rsp, 0x268
     2ac:   5b c3                             pop     r12
     2ae:   48                                pop     rbp
     2af:   c7 c0                             pop     r13
     2b1:   00 00                             pop     r14
     2b3:   00 00                             pop     r15
     2b5:   48                                pop     rbx
     2b6:   81                                ret
     2b7:   c4 68 02 00 00 41 5c 5d           mov     qword [rbx+0x00000020], 0x01
     2bf:   41 5d 41 5e 41 5f 5b c3           mov     qword [rbx+0x00000028], 0x07
     2c7:   48 c7 43 20 01 00 00              mov     rax, 0x00
     2ce:   00 48 c7 43 28 07 00              add     rsp, 0x268
     2d5:   00 00                             pop     r12
     2d7:   48                                pop     rbp
     2d8:   c7 c0                             pop     r13
     2da:   00 00                             pop     r14
     2dc:   00 00                             pop     r15
     2de:   48                                pop     rbx
     2df:   81                                ret

最適化

最適化は全くやってない。JITに関してはそんなに頑張らなくてもいいかなーと。

というか、実は VM コードに対しても現時点では全くやってない。定数の畳み込みすらやってないので、まずはそっちから。

今後

優先順位によってすぐにはできないかもしれないが、上記問題の解決はどこかでやるつもり。

それ以外、せっかく型指定できるようにしているので、native 以外での部分 JIT はできるようにした方がいいよね。逆にそれができれば native いらない。というか、単なる int 省略記法になる。JIT とは言っても、一般的なトレーシングとかではなく、コンパイル時にコード出力してしまうのだけどね。この辺の言葉の使い方も難しい。Just In Time の範囲がどこまでかによって変わるが、一般的に 実行時に ということのようなので、これも JIT のうちでしょう。

おわりに

こうして記事にしてみると、何気に具体的な課題が見つかっていい感じですね。予定立てます。

ではでは、今回もここまで読んでいただいてありがとうございます。最後はいつもの以下の定型フォーマットです。興味がありましたら ★ とか 「~~いいね~~ LGTM」 ボタンとか押してもらえるとモチベーションにつながります。どうぞよろしくお願いします。

最初の動機はスクリプト言語 KINX（ご紹介）を参照してください（もし宜しければ**「~~いいね~~ LGTM」**ボタンをポチっと）。
リポジトリはここ（https://github.com/Kray-G/kinx）です。こちらももし宜しければ★をポチっと。