vulkpy: Vulkan でGPGPUするPythonライブラリを作り始めた
0. はじめに
スクラップで情報収集しながら開発してたライブラリが、やっと形になってきたので記事にまとめます。
1. 開発内容と動機
1.1 vulkpy
今回開発を始めた vulkpy は、ベンダー・ニュートラルなGPU APIであるVulkanを利用して数値計算をするライブラリです。
import vulkpy as vk
gpu = vk.GPU()
a = vk.Array(gpu, data=[10, 10, 10])
b = vk.Array(gpu, data=[5, 5, 5])
c = a + b
c.wait()
print(c)
# [15, 15, 15]
1.2 なぜ開発したのか？
元々GPGPUに興味があったのですが、ちゃんとしたGPUを積んだPCを持っていなくて、何年もそのまま放置していました。
そんな中、Intel CPU内蔵のHD Graphicsでも（性能はともかく）Vulkanが利用可能であるという情報を見かけました。
動く環境があるのであれば、触ってみようかと開発をはじめました。
個人的な興味の観点から、深層学習をGPGPUで実現することを目標としています。
2. 構成
公式のC++バインディングである Vulkan-Hpp を呼ぶコア部分をC++で記述し、pybind11でPythonから呼べるようにモジュールを作成しています。
CIは以前記事を書いたように、Dockerfileで可能な限りCIプラットフォーム非依存を維持しつつGitHub Actions上で実行しています。
2.1 (検討したけど)採用しなかったもの
-
fynv/VkInline
- ⭕ Python内で、文字列でシェーダーを記載することで動的にコンパイルできて、簡単
- ❌ Vulkanは (WebAssembly等のように) SPIR-V中間バイナリの採用で、コンパイルを削減する設計だが、そのメリットを享受できない
- ❌ ライセンスがOSI Approvedなライセンスではなく、Anti 966ライセンスなる（中国の？）人権活動に関連する政治的なライセンス
-
realitix/vulkan
- ⭕ VulkanのAPIを全てそのままPythonに出力して、Pythonから直接触れるようにしている
- ❌ Pythonのバージョン非互換に起因するバグがあって、PRも取り込まれているけどPyPIにリリースされていない
- ❌ VulkanのAPIをPythonに出しているだけでそれ以上の支援はない。
-
nanobind
- ⭕ pybind11の作者が作成している後継ライブラリで、パフォーマンスやサイズの面で改善されている
- ❓ ヘッダーonlyだったpybind11と違い、動的ライブラリ(
libnanobind)を利用。（たぶんうまくやってくれるのだろうが）pipやcondaでインストールする動的ライブラリに正しくリンクできるのか、LD_LIBRARY_PATH周りで懸念があった。 - ❌(当時→今⭕) ドキュメントがpybind11のドキュメントを見て、うまく読み替えて使ってくれだった（が、いつの間にか専用のドキュメントサイトが立ち上がっていた）。
3. 苦労した点・課題
3. 1 苦労: リソース管理が大変
Vulkan本体は (たぶん) C言語のAPIで、vkCreateXXX() や vkDestroyXXX() のような独立した関数群でリソースを作成したり、削除したりします。(XXXにはDevice/Pipeline等のリソース名が入ります。)
また引数には VkXXXCreateInfo{} のような構造体を使い、メンバーに別の構造体を含むなど入れ子構造になっています。更にコピーしたくない（と思われる）構造体の場合は、そのアドレスを渡す設計になっている部分も多々ありました。
C++向けのVulkan-Hppは、上記のAPIをラップしており、vk::createXXX() や vk::destroyXXX() を提供しており、さらにvk::createXXXUnique()により、vk::UniqueXXXのようなデストラクタで自動的にリソースを解放してくれるスマートポインタのオブジェクトを作れます。
さて苦労した本題ですが、Pipelineが実行している間は参照しているリソースは変更してはいけない、解放してはいけないと、オブジェクト間には生成・解放の制約があるのですが、それはプログラマー側が責任を持って管理する必要があります。制約を守らないと、Segmentation faultが発生したり、壊れた値が出力されたりとバグが生まれます。
結局、一緒に使うリソース（のスマートポインタ）を保持するクラスを作ってそのクラスを std::shared_ptr で管理する方式に落ち着きました。依存するリソースが他にある場合は、依存リソースの std::shared_ptr をクラスに保持して、自身が生きている間は、依存リソースを維持できるようにしました。
3.2 苦労: C++の静的なtemplateとPythonの相性の悪さ
初めにC++で試行錯誤してある程度実装を固めてからPythonにI/Fを出していく順序で開発を進めました。
計算に応じて利用するバッファの数やパラメータが異なるため、C++ではそれをtemplateで実装しました。コンパイル時にバグを検出できるし、異なる型の引数を取ることができるしC++内では問題はありませんでした。
ですが、いざPython に出すにあたって、pybind11で動的にディスパッチするわけですが、コンパイルするわけでは無いので、必要な型を先に実体化して登録しておく必要がありました。
このぐらいだったら、ちょっと記述が増えるだけじゃないかと感じるかもしれませんが、テンプレートメタプログラミングの黒魔術っぽいヘルパー関数でバッファ数のディスパッチを作っています。
結論を言うと、(私はtemplate好きなのですが、) pybind11で外に出す部分は継承と仮想関数によるポリモーフィズムの方が設計として適していそうです。
3.3 課題: コンピュート・シェーダーの管理が大変
Vulkanのパイプラインはそれぞれ対応したコンピュート・シェーダー (SPIR-V中間バイナリ) が必要で、vulkpy では GLSL を glslc でコンパイルする方式を採用しています。
私が調べた限り、GLSLには import や include に対応するようなライブラリシステムはなく、tempalate や generics のような複数型をサポートする仕組みもなさそうです。
そのため、ほとんど同じで一部だけ異なるような GLSL ファイルが多数存在し、またfloat (32bit 浮動小数点) 以外の型をサポートしようとするとサポートするオペレーション全てを新たに（ほぼ）複製して追加する必要がある状態です。
次の例は、和と差が違うだけの２つのファイルです。
これは管理できないなと別の方式の検討を始めたのですが、今のところ「これだ！」と言える決め手に出会えていない状態です。
4. 今後の方向性
目標としてはじめに書いたように、深層学習ができるようにモジュールを開発していきたいと思っています。
一部は既に書き始めているのですが、いまいち設計が固まらなくて、書いては消しの試行錯誤を繰り返しています。
というのも、固定のネットワークを書いて学習させることは何とでもなるのですが、ライブラリとしてネットワーク構造を自由に組み替えたり、拡張できるようにしたりしようとすると途端に色々悩ましくて。。。
PyTorchやTensorFlowの実装を研究して見る必要があるのかもしれませんね。
4. まとめ
Vulkan でGPGPUするライブラリ vulkpy を作って公開しました。
もし興味を持ってもらえたら、使ってみたりレポジトリにスターを付けてもらえたりすると嬉しいです。
Discussion