[Python][WSL] PyCUDA で GPU プログラミング

  • 2022年8月21日
  • 2022年9月17日
  • Python, CUDA
  • 45View
  • 0件

はじめに

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 upgrade

CUDA Toolkit

次に CUDA の Toolkit をインストールします。
これを入れておかないと、PyCUDA のインストール時に cuda.h が見つからない旨のエラーが発生します。

$ sudo apt install nvidia-cuda-toolkit

Ubuntu リポジトリではなく NVIDIA 公式の最新版が利用したい場合は、こちらの記事の「手順②」を参照です。

Python パッケージ

Python パッケージをインストールするために pip3 をインストールします。
WSL Ubuntu には python3 は標準でインストールされていますが、pip3 は初期状態では含まれていません。

$ sudo apt install pip3

pip3 のインストールができたら PyCUDA をインストールします。

$ pip3 install PyCUDA

ホスト ⇔ デバイスの値やり取りで numpy を利用しますので、これもインストールします。

$ pip3 install numpy

サンプルコードの実行

準備できたらサンプルコードを実行します。

まず長さ 10 の配列を生成し、これを変数 a に代入します。
a の各要素を GPU 上で2倍して結果を変数 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.]