はじめに
Python で GPU プログラミングを行う手段として PyCUDA があります。
WSL Ubuntu 上で PyCUDA を動作させる手順をまとめます。
動作環境
- Windows 11
- WSL Ubuntu (20.04)
- GeForce RTX 2080
- nvcc 10.1
- python 3.8.10
事前準備
事前にいくつかインストールが必要なものがあります。
WSL ディストリビューションの複製
お試しで実行する場合は Ubuntu イメージを複製しておくのがお勧めです。
本手順は必須ではありません。
パッケージのアップデート
最初に各種パッケージをアップデートしておきましょう。
$ sudo apt update
$ sudo apt upgradeCUDA Toolkit
次に CUDA の Toolkit をインストールします。
これを入れておかないと、PyCUDA のインストール時に cuda.h が見つからない旨のエラーが発生します。
$ sudo apt install nvidia-cuda-toolkitUbuntu リポジトリではなく NVIDIA 公式の最新版が利用したい場合は、こちらの記事の「手順②」を参照です。
Python パッケージ
Python パッケージをインストールするために pip3 をインストールします。
WSL Ubuntu には python3 は標準でインストールされていますが、pip3 は初期状態では含まれていません。
$ sudo apt install pip3pip3 のインストールができたら PyCUDA をインストールします。
$ pip3 install PyCUDAホスト ⇔ デバイスの値やり取りで numpy を利用しますので、これもインストールします。
$ pip3 install numpyサンプルコードの実行
準備できたらサンプルコードを実行します。
まず長さ 10 の配列を生成し、これを変数 a に代入します。a の各要素を GPU 上で２倍して結果を変数 dest に還元します。
import numpy
import pycuda
import pycuda.autoinit
import pycuda.compiler
cuda_code = """
__global__ void double_vector(float *dest, float *a)
{
const size_t i = threadIdx.x;
dest[i] = 2.0 * a[i];
}
"""
module = pycuda.compiler.SourceModule(cuda_code)
double_vector = module.get_function("double_vector")
a = numpy.random.randn(10).astype(numpy.float32)
dest = numpy.zeros_like(a)
double_vector(
pycuda.driver.Out(dest),
pycuda.driver.In(a),
block=(10, 1, 1),
grid=(1, 1),
)
print(a)
print(dest)
print(dest - 2 * a)
$ python3 main.py
[ 1.7828326 -0.9319266 1.8419962 0.34073874 0.5151152 -1.7763407
-1.0569062 0.14110956 -2.1276314 -0.14621198]
[ 3.5656652 -1.8638532 3.6839924 0.6814775 1.0302304 -3.5526814
-2.1138124 0.2822191 -4.255263 -0.29242396]
[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]