PyPyでBrainf*ck

#jit #python

JIT関連を検索していて見つけた PyPy Status Blog: Tutorial: Writing an Interpreter with PyPy, Part 1 を読んで試してみようと思った。しかしリンク切れでソースが見れなかったりPyPyで変換できなかったりでそのままでは動かなかったので修正してみた。

使用したバージョンは 7.3.4：

$ pypy --version
Python 2.7.18 (63df5ef41012b07fa6f9eaba93f05de0eb540f88, Apr 08 2021, 15:52:53)
[PyPy 7.3.4 with GCC Apple LLVM 12.0.0 (clang-1200.0.32.29)]

RPythonは3.7.5

ソースのリポジトリ：https://github.com/tyfkda/pypy-bf

実装言語

件のブログで実装されている言語はBrainf*ckで、VM方式に命令をフェッチして実行の繰り返し。唯一の事前処理は命令以外の文字の除去と、対応する[]の飛び先を事前に調べることだけ。

# bf.py
def mainloop(program, bracket_map, putc, getc):
  tape = [0] * 30000
  pc = 0
  head = 0
  while pc < len(program):
    code = program[pc]

    if code == '>':
      head += 1
    elif code == '<':
      head -= 1
    elif code == '+':
      tape[head] += 1
    elif code == '-':
      tape[head] -= 1
    elif code == '.':
      putc(tape[head])
    elif code == ',':
      tape[head] = getc()
    elif code == '[':
      if tape[head] == 0:
        pc = bracket_map[pc]
    elif code == ']':
      if tape[head] != 0:
        pc = bracket_map[pc]

    pc += 1


def parse(program):
  parsed = []
  bracket_map = {}
  leftstack = []

  pc = 0
  for char in program:
    if char in ('[', ']', '<', '>', '+', '-', ',', '.'):
      parsed.append(char)

      if char == '[':
        leftstack.append(pc)
      elif char == ']':
        left = leftstack.pop()
        right = pc
        bracket_map[left] = right
        bracket_map[right] = left
      pc += 1

  return ''.join(parsed), bracket_map

ブログ記事から変更して、Tapeクラスを作らずに固定長のリストとヘッド位置をローカル変数として持つようにした。標準入力からの１文字読み込みと標準出力への１文字出力は、RPython用とで切り替える必要があるため、引数として受け取る形にした。

CPythonで動かすための起動部分。文字の入出力は sys.stdin.read、sys.stdout.writeを使用：

# for_cpython.py
import bf

def putc(c):
  sys.stdout.write(chr(c))  #; sys.stdout.flush()

def getc():
  return ord(sys.stdin.read(1))

def run(input):
  program, map = bf.parse(input.read())
  bf.mainloop(program, map, putc, getc)

if __name__ == '__main__':
  import sys
  run(open(sys.argv[1], 'r'))

動かすにはシェルからBrainf*ckのコードを渡す：

$ time python for_cpython.py examples/mandelbrot.b

試しにマンデルブロ集合を計算するプログラムを動かしてみたが、遅過ぎて実行完了を待ちきれずかかる時間は不明。

PyPyによる実行

PyPyはPythonの処理系の1つとして使用できる。なので上記のプログラムも動かすことができる。

インストール

Mac OSXだと brew install pypy でインストールできる。

またはpyenv経由でインストールすることもできるらしい。

実行

$ time pypy for_cpython.py examples/mandelbrot.b

手元で動かしたところCPythonよりずっと高速で完了まで待つことができて、70.70秒だった。

RPython Toolchain

RPython ToolchainというJITコンパイラ生成のツールチェインがあって、制限されたPythonの文法でインタプリタを書くとそれを変換して高速に実行できるようになる。 PyPyがRPythonで実装されているらしい。

PyPyのインストールではソースが含まれないので、ダウンロードのページのSourceから、またはリポジトリから落としてきてくる。

呼び出すコード

RPythonは制限されたPythonコードを受け付けるため、CPython版から少し変更する必要がある：

# for_rpython.py
import bf
import os

def putc(c):
  os.write(1, chr(c))

def getc():
  return ord(os.read(0, 1)[0])

def run(fp):
  program_contents = ''
  while True:
    read = os.read(fp, 4096)
    if len(read) == 0:
      break
    program_contents += read
  os.close(fp)
  program, bm = bf.parse(program_contents)
  bf.mainloop(program, bm, putc, getc)

def entry_point(argv):
  try:
    filename = argv[1]
  except IndexError:
    print('You must supply a filename')
    return 1

  run(os.open(filename, os.O_RDONLY, 0o777))
  return 0

def target(*args):
  return entry_point, None

if __name__ == '__main__':
  import sys
  entry_point(sys.argv)

RPython では sys が使えないとのことなので、 os.read や os.write を使う必要がある
- ただなにか仕様が違うらしく、これをCPythonで動かそうとしてもエラーが出る
- pypyに食わせて走らせることはできる
- os.writeだとバッファリングされないのではないかと推測するので不利じゃないかと思うが、行ごとに出力するように変更しても時間的には大差なかった
targetは変換用の関数で変換後のバイナリの実行開始情報となる

変換

記事にはインタプリタを変換するには pypy/translator/goal/translate.py を使うということが書かれているけど、今では rpython/bin/rpython 呼び出しに変わっているぽい。

PyPyのソースが置かれているパスを PYPY_SRC_PATH とすると、

$ $PYPY_SRC_PATH/rpython/bin/rpython for_rpython.py

しばらく待つとファイル名に-cが付けられた実行ファイルが生成される。実行：

$ time ./for_rpython-c examples/mandelbrot.b

手元では22.44秒だったので、PyPyでの実行時間と比べて31.7%(3.2倍速)に短縮された。

C++で同じ処理内容のVMを書いて動かしたところ43.71秒となり、RPythonで変換したもののほうが速い（すごい）。単なるコンパイルじゃないんだろうか…？

JIT化

チュートリアルのパート2 Adding a JIT に、JIT化について書いてある。

rpython による変換時に --opt=jit と指定するとJIT化される、らしい

JitDriver にgreensとredsを指定する
ループ箇所に jit_merge_point を入れる
JitPolicyを指定する

という必要がある：

# bf.py

+try:
+  from rpython.rlib.jit import JitDriver
+except ImportError:
+  class JitDriver(object):
+    def __init__(self,**kw): pass
+    def jit_merge_point(self,**kw): pass
+    def can_enter_jit(self,**kw): pass
+
+jitdriver = JitDriver(
+  greens=['pc', 'program', 'bracket_map', 'tape'],
+  reds=['head'])
+
 def mainloop(program, bracket_map, putc, getc):
   tape = [0] * 30000
   pc = 0
   head = 0
   while pc < len(program):
+    jitdriver.jit_merge_point(pc=pc, head=head, tape=tape, program=program,
+                              bracket_map=bracket_map)

    code = program[pc]

# for_rpython.py

+def jitpolicy(driver):
+  from rpython.jit.codewriter.policy import JitPolicy
+  return JitPolicy()

変換には --opt=jit をつけて、

$ $PYPY_SRC_PATH/rpython/bin/rpython --opt=jit for_rpython.py

（なぜかマンデルブロ集合のようなカラーのAAが出力され）7分半ほどかかり、実行ファイル for_rpython-c が生成される。

実行時間は11.85秒、普通の変換に比べて52.8%(1.9倍速)に短縮される。

最適化

Brainf*ckの [] のジャンプ命令でディクショナリ参照が行われるが、ジャンプ先は変わらないので最適化したい。 @purefunction デコレータでRPythonの変換にヒントを与えて、mainloop のブラケットの処理で使用するよう変更すると：

# bf.py
+from rpython.rlib.jit import purefunction
+
+@purefunction
+def get_matching_bracket(bracket_map, pc):
+    return bracket_map[pc]

 def mainloop(program, bracket_map, putc, getc, puts):
   ...
   while pc < len(program):
     ...
     elif code == '[':
       if tape.get() == 0:
-        pc = bracket_map[pc]
+        pc = get_matching_bracket(bracket_map, pc)
     elif code == ']':
       if tape.get() != 0:
-        pc = bracket_map[pc]
+        pc = get_matching_bracket(bracket_map, pc)

実行時間は3.48秒、単なるJIT化からさらに29.4%(3.4倍速)に短縮される。

実行時間まとめ：

処理系	秒数
CPython	計測断念
PyPy	70.70
C++	43.71
RPython変換後	22.44
–opt=jit	11.85
@purefunction指定	3.48

むちゃくちゃ速い、すごい！与えているのは単純なBrainf*ckのインタプリタだけなのに、どういう変換が行われればここまで速くなるのよ？

参考として、手動JIT版だと

[VM][JIT]Brainf*ckで学ぶスクリプト言語処理系高速化。インタプリタ→VMインタプリタ→JITコンパイラ。 - hogelogの日記を参考に、64ビット版Xbyakの速度を測ったところ、3.22秒だった。

RPython-jit版が手動でJIT化したものと遜色ないのがすごい。

RPython Toolchainについて

ドキュメントが古い状態のままだったり、検索してもあまり新しい記事がなく、あまり使われてないんだろうか…。

readthedocs
Comparing Partial Evaluation and Tracing, Part 1 | PyPy メタトレーシングがどうとか。
Tracing the meta-level: PyPy’s tracing JIT compiler 読んでない。
RPythonについて軽く | κeenのHappy Hacκing Blog 二村射影とか、部分評価とか、インタプリタからコンパイラが作られるとか。
RPythonを使ったインタプリタ実装にTopazとかいろいろあったみたいなんだけど最近は動きがない？どうなっちゃったんだろう？

実装言語

PyPyによる実行

インストール

実行

RPython Toolchain

呼び出すコード

変換

JIT化

最適化

参考として、手動JIT版だと

RPython Toolchainについて

関連記事

「JITコンパイルでの冒険」を読んだ

Xbyakで動的に生成するコードからCの関数を呼び出す

Direct Threaded版とJITの速度比較

JITについてお勉強

Xbyakで簡単な関数を実行時に生成する

新着記事