Python ctypes で CUDA Driver API を呼び出すメモ

背景

CUDA Driver API を python からぺろっと叩きたい.

pybind11 とか cfffi だとネイティブコードのビルドが必要になりめんどい. pure python(標準 python 機能)でやりたい(CI ビルドとかだと CUDA SDK インストールとかめんどい)

ctypes を使います.

CUDA Driver API(ドライバが入っていれば使える API)を想定します.
CUDA Runtime API(CUDA SDK or ランタイムのインストールが必要になる)は扱いません(API 面倒いし)

動かす Kernel は PTX 形式で生成されているものとします.
python で CUDA を叩くことで, たとえば PTX コードをテンプレートにしておいて, python 側で動的に書き換えてコンパイル(ターゲット GPU に合わせて最適なパラメータを設定)するとかができるようになります.

Clang で CUDA コードを NVPTX に変換するメモ
https://qiita.com/syoyo/items/4e60543aded0210fde49

dll/so のありか

Windows

C:\Windows\System32\nvcuda.dll

通常は nvcuda.dll だけで良い
(CUDA は NVIDIA driver インストールすれば自動で入っているので, ロードできなかったら NV GPU が挿さっていない PC(e.g. Intel 内臓/Xe or AMD GPU) であろう)

Linux(Ubuntu)

Ubuntu ですと /usr/lib/x86_64-linux-gnu/ あたりにあります.

さわりだけ

from ctypes import *

cu = cdll.LoadLibrary("/usr/lib/x86_64-linux-gnu/libcuda.so")
print(cu)

ret = cu.cuInit(0)
assert ret == 0 # CUDA_SUCCESS

ver = c_int()
ret = cu.cuDriverGetVersion(byref(ver))
assert ret == 0 # CUDA_SUCCESS

print("CUDA version", ver)

CUDA version c_int(11020)

Voila!

ポインタのところは byref を使うといいようです.

あとはいろいろ API 呼び出していけば OK のはずです!

その他

PTX Compiler API

PTX Compiler API のメモ
https://qiita.com/syoyo/items/cfaf0f7dd20b67cc734e

static lib しか提供されていないので, 一旦なんらかの dll を作る必要がありますが, これによりドライバにある PTX compile では満足できない場合, クライアントサイドでの PTX compile も pure Python でできるようになります!

Runtime API

cuSparse など, runtime API の上に作られたライブラリを使いたいときもあります.

一応だいたいの dll/so は再配布可能なので,

必要であれば自身のアプリに同梱するなどして対応しましょう
(e.g. CI 環境で毎回 CUDA SDK インストールとかめんどい)

TODO

• numpy との連携を考える(Python buffer protocol で連携になるかしらん)