Python の型アノテーションと ast モジュールで, python コードから C/C++ コード生成に思いを馳せる #Python

背景

グラフィックスや数値計算などでは, 速度やランタイム環境の関係で C/C++ が主流ですが, C/C++ でゴリゴリ書くというのはなかなか骨の折れる作業です(e.g. C/C++ ではベクトルクラスなどはないので, 毎回プロジェクトごとにベクトルクラスを定義したりとか)

また, グラフィックスや数値計算アプリではアルゴリズムや実装が試行錯誤になることが多く, C/C++ で直接書くとなかなかテストや評価が大変です(画像でテストする場合は画像ローダも書かないといけないし, 結果をプロットしたいときは一度データを吐いて gnuplot で確認とか.

このとき, IPython でインタラクティブにアルゴリズム実装を確認したり, matplotlib でプロットしたり, Jupytor notebook でロードした画像をインラインで表示できたりと, python だといろいろエコシステムが整っていて便利です.

10 年くらい前! に python コードを LLVM IR に変換して高速実行というのを試したことがあるのですが, このとき型推論が面倒(型推論の実装が面倒)で, 途中でやめてしまった経緯があります.

py2llvm: Python to LLVM translator
https://llvm.org/devmtg/2008-08-23/py2llvm.pdf

(その後の aherlihy さんによる改良版)
https://github.com/aherlihy/PythonLLVM

python のコード自体を C/C++ に変換して早くするというプロジェクトはいくつかありますが(Cython など), python を DSL のように扱い, 特定のアルゴリズム/実装を C/C++ コード化して早くする, という取り組みはあまりなかった気がします.

最近になり pytorch が python で記述したモデルを C++ にエクスポートして機械学習をモバイルや C++ の環境だけで動かせるようにする取り組みが始まるなども出てきています.

型アノテーションと型ヒント

Python 3.5 から, 型のアノテーションと, 型ヒント, 型の処理(typing モジュール)ができるようになっています.

Pythonではじまる、型のある世界
https://qiita.com/icoxfog417/items/c17eb042f4735b7924a3

3.5 から導入された型アノテーションを使えば, ユーザが明示的に型をつけ, mypy http://mypy-lang.org/ でチェックすることにより, (たぶん)型推論を実装しなくて済む(or 実装が楽になる)ので良さそうですね!

型アノテーションは 3.5 以前(e.g. python2)ではエラーになってしまいますが, python もそろそろ(2018 年 5 月時点) バージョン 2 はサポートが終わり(少なくとも 2020 までには終了: https://pythonclock.org/ ), バージョン 3 がこれからの主流になりますので, 今から python を使うのであれば 3.5 以降を仮定してよいでしょう.

実装はじめ

typing モジュールを使い, グラフィックスでよく使う三次元ベクトル型(vec3)を tuple として定義してみます.

# proc.py

import typing

vec3 = typing.Tuple[float, float, float]

def fun_add(a: vec3, b: vec3) -> float:
    """
    >>> a = [1, 2, 3] # type: vec3
    >>> b = [1, 3, 7] # type: vec3
    >>> fun_add(a, b)
    (2, 5, 10)
    """
    return (a[0] + b[0], a[1] + b[1], a[2] + b[2])


if __name__ == "__main__":
    import doctest
    doctest.testmod()

普通に python スクリプトとして動かしてみます.

$ python proc.py -v

Trying:
    a = [1, 2, 3] # type: vec3
Expecting nothing
ok
Trying:
    b = [1, 3, 7] # type: vec3
Expecting nothing
ok
Trying:
    fun_add(a, b)
Expecting:
    (2, 5, 10)
ok
1 items had no tests:
    __main__
1 items passed all tests:
   3 tests in __main__.fun_add
3 tests in 2 items.
3 passed and 0 failed.
Test passed.

動作しました.

mypy で型チェックを行う.

mypy を型チェッカとしてつかいます.

$ mypy proc.py

proc.py:12: error: Incompatible return value type (got "Tuple[float, float, float]", expected "float")

Cool! 無事戻り値の型が違うと報告してくれました. -> vec3 に修正しておきましょう.

AST でパースする.

Python では ast モジュールで Python スクリプトの内部 AST 表現が取得できます.

import ast

t = ast.parse(open("proc.py").read())
print(ast.dump(t))

Module(body=[Import(names=[alias(name='typing', asname=None)]),
Assign(targets=[Name(id='vec3', ctx=Store())],
value=Subscript(value=Attribute(value=Name(id='typing', ctx=Load()),
attr='Tuple', ctx=Load()), slice=Index(value=Tuple(elts=[Name(id='float',
ctx=Load()), Name(id='float', ctx=Load()), Name(id='float', ctx=Load())],
ctx=Load())), ctx=Load())), FunctionDef(name='fun_add',
args=arguments(args=[arg(arg='a', annotation=Name(id='vec3', ctx=Load())),
arg(arg='b', annotation=Name(id='vec3', ctx=Load()))], vararg=None,
kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
body=[Expr(value=Str(s='\n    >>> a = [1, 2, 3] # type: vec3\n    >>> b = [1,
3, 7] # type: vec3\n    >>> fun_add(a, b)\n    (2, 5, 10)\n    ')),
Return(value=Tuple(elts=[BinOp(left=Subscript(value=Name(id='a', ctx=Load()),
slice=Index(value=Num(n=0)), ctx=Load()), op=Add(),
right=Subscript(value=Name(id='b', ctx=Load()), slice=Index(value=Num(n=0)),
ctx=Load())), BinOp(left=Subscript(value=Name(id='a', ctx=Load()),
slice=Index(value=Num(n=1)), ctx=Load()), op=Add(),
right=Subscript(value=Name(id='b', ctx=Load()), slice=Index(value=Num(n=1)),
ctx=Load())), BinOp(left=Subscript(value=Name(id='a', ctx=Load()),
slice=Index(value=Num(n=2)), ctx=Load()), op=Add(),
right=Subscript(value=Name(id='b', ctx=Load()), slice=Index(value=Num(n=2)),
ctx=Load()))], ctx=Load()))], decorator_list=[], returns=Name(id='float',
ctx=Load())), If(test=Compare(left=Name(id='__name__', ctx=Load()), ops=[Eq()],
comparators=[Str(s='__main__')]), body=[Import(names=[alias(name='doctest',
asname=None)]), Expr(value=Call(func=Attribute(value=Name(id='doctest',
ctx=Load()), attr='testmod', ctx=Load()), args=[], keywords=[]))], orelse=[])])

ast からアノテーションを取得できそうですね.

あとは Visitor パターンで AST をトラバースし, アノテーションを参考に C/C++ コードに変換していけばいろいろできそうですね!

TODO

自動微分と組みわせる. e.g. https://github.com/google/tangent
Rust 言語で似たようなことをやりたい.
優秀な若人が, 人類史上最速で優秀な型つき python 若人へと昇華なされるスキームを確立する旅に出たい.