はじめに
TL;DR
C++向けにCUDAを使ったライブラリを開発している中で、pybind11を使えばpythonからでも比較的簡単に使えることに気が付いたまでは良かったのですが、いざsetup.pyを作ろうとしたら苦労したので、備忘録として記載しております。殆ど自分用のメモなので需要はあまりないと思いますが、もし同じようなことをされる方の参考になれば幸いです。
なお、私は普段はC++メインでpythonは初心者なのでもっと良いやり方があるのかもしれません。
やりたいこと
下記に対応できるsetup.pyを作るのが目標です。
- Windows と Linxu(ubuntu) に対応したい
- nvccのインストールされていない環境ではCPU版がビルドされるようにしたい
ということで、状況に応じて4パターンに対応できれば良いということになります。
参考にしたもの
結局いろいろ見ましたが、下記のソースコードを読むのが一番理解に繋がりました。
unixccompiler と msvccompiler は共に ccompiler を継承したクラスで、それぞれunix(gcc)用とwindows(cl.exe)用となっているようです。
また、探した限りこちらの記事や、CuPyのソースコードも参考になりました。
ただし、前者はWindowsに未対応で、後者はやや大掛かりでした。今回は自分の必要な部分だけなるべくシンプルに動くように作成を試みて、PyPIに登録するところまで試すことができました。
事前準備
まず始めに nvcc が動く状態になっていることは前提とします。
また、Pythonは Python3 を前提として、windows環境には64bit版がインストールされているものとします。
Windows においては、msvc (Microsoft Visual C++) の cl.exe の 64bit 版がコマンドラインから使えるようになっている必要があります。
これは例えばコマンドラインから
"C:\Program Files (x86)\Microsoft Visual Studio\2017\Community\VC\Auxiliary\Build\vcvarsall.bat" x64
のように実行することで設定できます(実際のインストール先は環境によって異なると思いますので読み替えてください)。
このとき x64 オプションを忘れないことが重要です。
windowsの場合、pip を使ってインストールする場合であっても、バイナリではなくソースコード(sdist)からのインストールとなる場合は事前設定しておかないといけないので注意が必要です。
setup.py の作成
結果的に私が作成したのは、このソースコードとなります。
hook_compilerメソッドの箇所が今回のメインの箇所となります。
(プロジェクト特有の無関係の処理や、解析過程のprintなどが残っていますがご容赦ください)
方針
基本的に ccompiler クラスのメソッドをオーバーライドしてフックすることで、'.cu' ファイルが来たときのみ処理を切り替えることで対応しました。
pybind11のsetup.pyサンプルがここにあるので、これをC++コードのコンパイルの基本形として対応を行っていきました。
unix と msvc の切り分け
ccompiler クラスの compiler_type メンバ変数を見れば、派生クラスで 'unix' や 'msvc' と設定されるようなのでここで判定が出来ます。
compile処理
次にコンパイル時のフックなのですが、unixccompiler では compileメソッド -> _compileメソッド の順で処理されます。compileメソッドでpython用に必要なincludeパスやライブラリパスが追加されて、_compileメソッドが呼ばれるので、_compileメソッドをフックするのが簡単なようです。
一方で、msvccompiler では、_compileメソッドは呼ばれず compileメソッド の中で完結しています。
もともと cl 用のオプションが得られても、nvcc の形式に変換するのが逆に大変なので、msvc 時はcompileメソッドをフックして、代替処理を作成することを考えます。
このとき、親クラスである ccompiler の _setup_compile メソッドを活用することで楽が出来ます。 _setup_compile メソッドを呼び出すとその中でpythonに必要な項目を追加して出力先のディレクトリの作成などをしてくれるようです。
あとは、cl.exe 用であれば /D"(マクロ定義)" や /I"(インクルードパス)" のように、nvccやgccであれば -D(マクロ定義) や -I(インクルードパス) のようにの加工を行っていけばよい事になります。
Windowsでは空白のパス名の対策で "(ダブルクォーテーション)で囲う必要があるので注意ください。
link処理
リンクに関しては _linkメソッドのようなものはなく、メタデータの解釈が必要になりそうです。
linkはオリジナルのまま使うことも考えたのですが、nvcc を使う場合、CUDA関連のライブラリ関連の設定を考えるとリンクにも nvcc を使う方が楽そうです。
私の場合は cuBLAS や OpenMP を使っていたのですが、-lcublas オプションを nvcc に、-fopenmp は -Xcompiler オプションを使って、gcc に投げるようにしたことで、それぞれに適切なライブラリをリンクしてもらえました。
link に関しても同様に、親クラスの ccompiler の _fix_lib_args メソッドを使うことで楽が出来るようです。
その他
今回基本的に4パターンに応じたオプションを手動で用意しておくことでいろいろと簡略化しています。
コンパイラクラス自体を複数生成させるわけではないので、先人に習って extra_compile_args に、辞書として両方入れておき、フック先で切り替えています。
また、今回は nvcc を探すのにもいろいろコードを入れていますが、結果的に nvcc 経由で呼び出せば不要な情報もあり、もっと省略できそうにも思います。
おわりに
今後、setuptoolsなどが進化した場合に同じ手法が使えるかどうかわかりませんが、このあたりを活用したい人の参考になれば幸いです。
nvcc自体は unix 用とwindows用で使い方は近いのですが、unixccompiler と msvccompiler とで、意外に共通項が少なく思っていたより苦労しました。