サブロウ丸

Sabrou-mal サブロウ丸

主にプログラミングと数学

pybind11を用いた高速化 (master_mind)

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++に置き換えます。簡単なやり方としては

  1. 既存のpythonで実装した関数と同様の処理を行うc++コードを作成
  2. そのさい、引数や返り値の型は元のpython関数のものに合わせて適当に選択
  3. c++コードをコンパイルpythonモジュール化
  4. python側からそのモジュールをimportして関数などを使用

関数だけでなくクラスなどもモジュール内に含めることができます。

pybind11によるc++の置き換え

utils/common.pyにあるcount_hitblowとcalc_dist関数をc++で書き直します。コードはutils/common.cppとして実装していきましょう。

count_hitblow関数

それぞれのリンクにPython版とc++版のコードを記載。

python , cpp

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;
}

コード全体

master_mind/common.cpp at 77534e385e6a09d9e90ac8888a3718c6846094a3 · nariaki3551/master_mind · GitHub

cppコードで実装した関数は下のリンクでPYBIND11_MODULEで定義してあるようにcpp_utilsという名前でモジュール化されます。

master_mind/common.cpp at 77534e385e6a09d9e90ac8888a3718c6846094a3 · nariaki3551/master_mind · GitHub

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倍ほど高速化されました。