pybind11を用いてPythonプロジェクトの一部をc++コードで置き換えて高速化を試みます。
本稿では以前作成したmaster mindのプロジェクトを題材にします。
本稿で出てくる用語(ほとんど出さないように気をつけましたが)はmaster mind by c++; part 1 基盤プログラム - サブロウ丸この記事に定義されてあります。
何はともあれ高速化するならプロファイル。 プロファイルを行なってボトルネックになっている箇所を探して、そこをc++に置き換えてみましょう。
プロファイル
プロファイル結果の図を下記に添付しています。これを見るとcommon.py::calc_dist::count_hitblow 関数あたりがボトルネックになっていますね。ここでcount_hitblowは2つのコードを与えたときの(hit, blow)を計算する関数で、calc_distは1つのコードxとコード集合Yを与えたときにxとy in Yのペアの全てのcount_hitblowを実行する関数です。
ということで、この部分をc++に置き換えます。簡単なやり方としては
- 既存のpythonで実装した関数と同様の処理を行うc++コードを作成
- そのさい、引数や返り値の型は元のpython関数のものに合わせて適当に選択
- c++コードをコンパイルしpythonモジュール化
- python側からそのモジュールをimportして関数などを使用
関数だけでなくクラスなどもモジュール内に含めることができます。
pybind11によるc++の置き換え
utils/common.pyにあるcount_hitblowとcalc_dist関数をc++で書き直します。コードはutils/common.cppとして実装していきましょう。
count_hitblow関数
calc_dist関数
pythonのオリジナルの関数にlist, dict, tupleが関数の引数や返り値に含まれる場合、c++の方でstd::vectorやstd::map, std::tupleに置き換えて書いておくといい感じにpythonコードと結合してくれます。c++だとmap(下のコードはunordered_map)のkeyにintのtupleを用いてますがkeyはhashableである必要があるので struct tuple_hashでハッシュ値を定義しています。
python
def calc_dist(guess, feasible_codes, config): distribution = defaultdict(list) for code in feasible_codes: hit, blow = count_hitblow(guess, code, config) distribution[hit, blow].append(code) return distribution
cpp
using Code = std::vector<int>; using Hint = std::tuple<int, int>; struct tuple_hash { template <class T1, class T2> std::size_t operator() (const std::tuple<T1, T2> &t) const { return std::hash<T1>()(std::get<0>(t)) ^ std::hash<T2>()(std::get<1>(t)); } }; using Distribution = std::unordered_map<Hint, std::deque<Code>, tuple_hash>; Distribution calc_dist( Code &guess, std::vector<Code> &feasible_codes, int &num_colors, int &num_pins ) { int hit, blow; Distribution dist; for( auto &code : feasible_codes ) { count_hitblow(guess, code, hit, blow, num_colors, num_pins); dist[std::make_tuple(hit, blow)].push_back(code); } return dist; }
コード全体
cppコードで実装した関数は下のリンクでPYBIND11_MODULEで定義してあるようにcpp_utilsという名前でモジュール化されます。
python側は
from master_mind.cpp_utils import calc_dist as cpp_calc_dist
というようにモジュールをインポートして使用します。(mastermind/utils/common.pyを参照)
コンパイル
コンパイルはCMakeLists.txtを用いて行いました。上記のリンク参照。python setup.py install を行うとcppファイルのコンパイルとmaster_mindのパッケージ化が行われて、mastermindというコマンドが使えるようになります。
実験
colorとpinを変えながら実行速度を比較してみます。 10回実行してその平均実行時間を下記の表にまとめました。
(COLOR, PIN) | pybind11 | python only |
---|---|---|
(6, 3) | 0.075s | 0.155s |
(5, 4) | 0.510s | 1.276s |
(6, 4) | 2.132s | 6.290s |
(7, 4) | 9.327s | 27.377s |
(4, 5) | 1.985s | 4.445s |
(5, 5) | 17.248s | 44.275s |
ということで大体3倍ほどの高速化に成功していますね。(ここまでうまくいくとは...)
実行コード
function execute () { for i in {0..9}; do # echo $COLOR $PIN $(python master_mind.py $COLOR $PIN --benchmark | grep "Running time") echo $COLOR $PIN $(mastermind $COLOR $PIN --benchmark | grep "Running time") done } COLOR=6; PIN=3; execute COLOR=5; PIN=4; execute COLOR=6; PIN=4; execute COLOR=7; PIN=4; execute COLOR=4; PIN=5; execute COLOR=5; PIN=5; execute
まとめ
pybind11を用いてpythonで書かれたmaster_mind求解コードの一部の関数をc++コードによる関数に置き換えることで、大体3倍ほど高速化されました。