はじめに

前回までに、OpenACCやC++17で規格化された標準仕様に則った並列化を通じてCUDA GPUの利用について紹介しました。既存のシミュレーションソフトウェアでは、Fortran/C/C++といったコンパイルを必要とするプログラミング言語が中心であるため、これらの機能の必要性はもちろん十分にあるのですが、言語にこだわらずよりGPU friendlyなプログラミング環境へ移行するというのも手だと考えます。

数値計算、特に機械学習においてはPythonが圧倒的に広範なユーザーに支持されています。個人の意見ですが以下の理由が主流ではないかと考えられます。

• スクリプト言語のため開発と実行のサイクルが比較的早く回せる
• 数値計算やデータ処理に関する豊富なライブラリがOSSコミュニティによって開発されている
• 上記のライブラリがGPUで容易に高速化できるように設計されている

特に2つ目のライブラリの豊富さは、他のスクリプト言語に対し大きな強みです。実際の計算はライブラリ内で行われるように実装されているため、GPUへのメモリコピーや計算カーネルの実装といったGPUを用いる上で必要なコードはユーザーランドには見えません。そのため、ユーザーは本来の目的であるシミュレーションやデータ分析、機械学習タスクといった問題に集中できるようになっています。高速化だけでなく開発自体が高度化している現在、HPC分野においてもPythonに代表されるスクリプト言語を用いたシミュレーションの開発の高速化や性能向上は1つのトピックとして注目されています。

ではどの程度使えるのか、簡単にではありますが、今回はその一旦をご紹介できればと思います。Pythonの書き方を詳細に説明しておりませんので、ご承知おきください。

今回使用したCPUとGPU

今回使用したCPUとGPUを表１に示します。使用するコンピュータには以下のCPUが2基、GPUが1基搭載されています。

取り扱うライブラリについて

今回は数値計算に着目し、以下のライブラリを取り上げます。

• NumPy
  • ベクトルや行列計算などの高速な実装を提供
• SciPy
  • NumPyをベースにしたより高水準な科学計算を提供
• CuPy
  • NumPyおよびSciPy互換なインターフェイスでGPUで高速化された実装を提供

今回使用しているバージョンは以下の通りです。

• NumPy 1.25.1
• SciPy 1.11.1
• CuPy 12.1.0

NumPyとSciPyはPythonの数値計算ライブラリの中でも特に広く使われているものの代表例です。NumPyはシンプルなインターフェイスでベクトルや行列計算を容易に記述できるため、他のライブラリは設計や実装のベースとして取り入れたり、その思想を強く受けています。SciPyはNumPyをベースに離散フーリエ変換、数値積分、疎行列計算などより高度な科学計算を提供しており、例えば共役勾配（CG）法に代表される疎な連立一次方程式の解法も提供されています。

NumPyとSciPy互換なインターフェイスで、計算をGPUで高速化するライブラリがCuPyです。CuPyはPythonとCUDAの橋渡しを行い、独自のCUDAカーネルやNVIDIAが提供するCUDAライブラリを組み合わせて、NumPyとSciPyの計算を高速化しています。前述の通り互換性のあるインターフェイスを提供しているため、NumPy/SciPyとCuPyを合わせることで、ソフトウェアを真にコードの変更なしにCPUでもGPUでも動作可能になります。

その魅力は次の節からご説明します。

Pythonでのライブラリ管理は色々な流派がありますが、ここでは標準的なコマンドを紹介いたします。

$ python --version
Python 3.11.4
$ pip install numpy scipy cupy-cuda12x

NumPyとCuPyによるベクトルと行列の計算

まずは非常に簡単なベクトル計算のコードが以下です。Pythonを読んだことがない方向けにコメントを記載していますので、参考にしてください。

import numpy as np  # numpy packageをnpという名前で参照、よくnpと略されます

def main() -> None:
    n = 12
    x = np.random.seed(n).astype(np.float64)  # [0.0, 1.0] で乱数ベクトルを生成
    y = np.zeros(n, dtype=np.float64)         # ゼロベクトル
    a = 0.5
    y = a * x + y
    print(y)
    """ 想定される出力
    [0.24414441 0.04736708 0.45870277 0.25938323 0.10383455 0.2711616
     0.15058139 0.4090193  0.31455517 0.321877   0.08782681 0.49883217]
    """

if __name__ == "__main__":  # プログラムのエントリポイント
    main()

非常に簡単なコードなので、メリットが今のところは見えないと思います。ではCuPyに置き換えましょうと言いたいのですが、これは先頭の1行を以下の通り書き換えるだけです。

import cupy as np  # cupy packageをnpという名前で参照

実際にプログラムを走らせてみると、Windowsの場合はタスクマネージャーで確認するとGPUが動いていることを確認できると思います。ですが、これだけだとまだ嬉しさが見えてこないため、実行時にCPUとGPU実行を切り替えるコードに書き換えます。

import argparse
def main(use_gpu: bool) -> None:
    if use_gpu:
        print("Computation with CUDA GPU")
        import cupy as np
    else:
        print("Computation with CPU")
        import numpy as np
    n = 12
    x = np.random.rand(n).astype(np.float64)
    y = np.zeros(n, dtype=np.float64)
    a = 0.5
    y = a * x + y
    print(y)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()  # オプション引数を解析するためのパーサー
    parser.add_argument("-g", "--gpu", action="store_true", default=False, dest="use_gpu")
    args = parser.parse_args()
    main(use_gpu=args.use_gpu)

オプション引数で-gまたは--gpuを与えると、GPUで実行されます。このときmain関数でライブラリの読み込みを制御し、フラグ引数に応じて読み込むライブラリを切り替えています。これを実際に実行してみると、以下のような出力が得られます。

$ python axpy.py
Computation with CPU
[0.24414441 0.04736708 0.45870277 0.25938323 0.10383455 0.2711616
 0.15058139 0.4090193  0.31455517 0.321877   0.08782681 0.49883217]
$ python axpy.py --gpu
Computation with CUDA GPU
[0.0066791  0.42153075 0.19287358 0.17575687 0.47828022 0.48310485
 0.35802358 0.3128375  0.46295458 0.44317445 0.16289805 0.10339087]

CPU（NumPy）とGPU（CuPy）を動的に切り替えるというのは、FortranやCなどでは（できなくもないですが）簡単にはできないと思います。

上記のコードではGPUの高速化メリットを感じられないため、次に行列の計算を示します。ここではN次正方行列と既知ベクトルを作り、未知ベクトルをNumPyで求めます。

import argparse
import time
def main(use_gpu: bool) -> None:
    if use_gpu:
        print("Computation with CUDA GPU")
        import cupy as np
    else:
        print("Computation with CPU")
        import numpy as np
    n = 8000
    A = np.random.uniform(size=(n, n)).astype(np.float64)  # 全要素 [0, 1) のn次正方行列
    b = np.ones(n, dtype=np.float64)                       # 全要素 1 の既知ベクトル
    """
    warmup: 初期化などのオーバーヘッドを計測から除外
    CuPyは最初のカーネル実行でCUDAコードの最適化といった性能に関する初期設定が走るようです
    """
    x = np.linalg.solve(A, b)  # https://numpy.org/doc/stable/reference/generated/numpy.linalg.solve.html
    tbeg = time.clock_gettime(time.CLOCK_MONOTONIC)
    x = np.linalg.solve(A, b)  # Ax = b を解く、計算にはLAPACKの?gesvが使われる
    tend = time.clock_gettime(time.CLOCK_MONOTONIC)
    print(x)
    print(f"computation time: {tend - tbeg:.2f} sec.")
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-g", "--gpu", action="store_true", default=False, dest="use_gpu")
    args = parser.parse_args()
    main(use_gpu=args.use_gpu)

--gpuフラグの有無で実行すると、GPUを用いて高速化されていることが確認できます。使用しているコンピュータは2CPUが搭載されており、1CPU（12コア）と1GPUを比較するためにCPUでの実行ではMKL_NUM_THREADS, NUMEXPR_NUM_THREADS, OMP_NUM_THREADSを12スレッド並列で指定しnumactlコマンドを使用しました。?gesvの計算はLU分解を使って解きますが、今回の場合はGPUが非常に高速に計算できているようです。

$ export MKL_NUM_THREADS=12
$ export NUMEXPR_NUM_THREADS=12
$ export OMP_NUM_THREADS=12
$ numactl --cpunodebind=0 --localalloc python linalg_solve.py
Computation with CPU
[ 0.03040997  0.01440329  0.06772684 ...  0.00026573 -0.0293215
  0.01875962]
computation time: 0.94 sec.
$ python linalg_solve.py --gpu
Computation with CUDA GPU
[ 0.0675002   0.03450186 -0.03794135 ...  0.01511969 -0.01006322
 -0.0252226 ]
computation time: 0.01 sec.

SciPyを用いた連立一次方程式の計算とCuPyでの高速化

もう少し実践的なコードとして、SciPyを使った係数が疎行列になる連立一次方程式の計算をCuPyで高速化します。SciPyコードの場合はNumPyと組み合わせますが、NumPy/ScipyともにCuPyで置き換えることで、GPUで実行できるようになります。

SciPy互換インターフェイスを持つCuPyの実装はcupyxというパッケージ名で提供されており、概ねscipyパッケージ名の前にcupyxをつければ読み込めるようです。

import argparse
import time
def main(use_gpu: bool):
    if use_gpu:
        print("Computation with CUDA GPU")
        import cupy as np
        from cupyx.scipy.sparse import dia_matrix  # パッケージ単位だけなく関数やクラス単位で読み込み可能
        from cupyx.scipy.sparse.linalg import cg
    else:
        print("Computation with CPU")
        import numpy as np
        from scipy.sparse import dia_matrix  # https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.dia_matrix.html
        from scipy.sparse.linalg import cg   # https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.cg.html
    nx = 1024
    ny = 1024
    n = nx * ny
    ex = np.ones(n)
    """
    2次元中心差分の帯行列を構築
       -1
    -1  4 -1
       -1
    帯行列になるためDIA (CDS) formatで疎行列を表現します
    https://netlib.org/linalg/html_templates/node94.html
    """
    datas = np.array([-ex, -ex, 4 * ex, -ex, -ex])  # 疎行列の要素を帯ごとに持つ
    offsets = np.array([-nx, -1, 0, 1, nx])         # 各帯の対角要素からの相対位置
    A = dia_matrix((datas, offsets), shape=(n, n), dtype=np.float64)  # DIA (CDS) matrix classに変換
    # 未知ベクトルをランダムに設定し、既知ベクトルを作成
    x = np.random.uniform(size=n).astype(dtype=np.float64)
    b = A @ x  # matrix-vector multiplication
    x, info = cg(A, b, tol=1.0e-6)  # warmup
    tbeg = time.clock_gettime(time.CLOCK_MONOTONIC)
    x, info = cg(A, b, tol=1.0e-6)  # CG法で Ax = b を計算
    tend = time.clock_gettime(time.CLOCK_MONOTONIC)
    if info == 0:  # 収束 = 解が得られた
        print("solved: absolute error ", np.linalg.norm(A @ x - b) / np.linalg.norm(b))
    else:
        print(f"not solved: number of iters {info}")
    print(f"computation time: {tend - tbeg:.2f} sec.")
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-g", "--gpu", action="store_true", default=False, dest="use_gpu")
    args = parser.parse_args()
    main(use_gpu=args.use_gpu)

こちらも先程と同様にオプションをつけることでCPU実行とGPU実行を切り替えられます。疎行列の計算の場合、演算性能よりもメモリバンド幅が性能に影響しますので、先程の密行列の例に比べると性能差は縮んでいるようです。演算誤差の影響により反復回数が若干増減しているかもしれませんが、大きな差を与えるものではないと考えられます。

$ export MKL_NUM_THREADS=12
$ export NUMEXPR_NUM_THREADS=12
$ export OMP_NUM_THREADS=12
$ numactl --cpunodebind=0 --localalloc python sparse_solve.py
Computation with CPU
solved: absolute error  9.910306939584394e-07
computation time: 41.03 sec.
$ python sparse_solve.py --gpu
Computation with CUDA GPU
solved: absolute error  9.94122243817905e-07
computation time: 1.53 sec.

おわりに

今回はPython + GPUを取り上げ、数値計算ライブラリであるNumPy/SciPy、その互換インターフェイスを提供するGPU実装のCuPyを用いたコードを紹介しました。CuPyが互換インターフェイスを持つため、プログラマはNumPy/Scipyのコードを書けば、それをCuPyへのパッケージ置換だけでGPUに計算させられます。GPU化そのものにはほとんど労力を使っていないので、計算コードの実装に集中できます。

互換インターフェイスを持つ実装としてCuPyを取り上げましたが、例えばデータ解析用ライブラリであるPandasとそれをGPU化するcuDFのような完全互換ではないものの元のCPU実装をGPUに変換する機能を提供するライブラリや、機械学習ライブラリであるPyTorchは単体でNVIDIA以外のGPUも含むマルチアーキテクチャをサポートしており、Pythonはより身近にGPUの性能を感じられる環境と言えそうです。

補足: Python + GPU環境について

多くのライブラリがGPUで提供されており魅力的である一方、Pythonの環境構築は様々な方法があり苦労が多いようです。機械学習やデータサイエンスメインではありますが、アカウントがあれば環境構築の手間なしに、かつGPUも使える環境を提供するウェブサービスがあります。このようなウェブベースのPython実行環境では、Jupyter Notebookという対話型環境が広く利用されています。

有償無償が混在しておりますのでここではサービス名だけ紹介し、皆様の事情に合わせて利用可否を調査いただければ幸いです。

• Google Colab
• Amazon SageMaker
• Paperspace Gradient