PythonからCUDAを利用してベクトル加算を行う方法

TL;DR

本記事では、PythonからCUDAを利用して大規模なベクトル加算を行う方法について解説します。

CUDAのコードはC++と似ているため、C++の基礎知識があると比較的理解しやすいです。PythonからCUDAを利用するためには、ctypesを使って共有ライブラリをロードし、データを渡す必要があります。

コードはこちらです：makinzm/cuda-from-python

1. CUDAコードの解説（vector_add.cu）

まず、CUDAで実装されたベクトル加算のコードを見ていきましょう。

ヘッダファイルのインクルード

#include <cuda_runtime.h>

• #include <cuda_runtime.h>:
  • CUDAのランタイムAPIを使用するためのヘッダファイルをインクルードしています。これにより、CUDAのメモリ管理やカーネルの起動に関する関数が使用可能になります。
  • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#default-stream

extern "C"による関数のエクスポート

extern "C" void vector_add(float *a, float *b, float *c, int n);

• extern "C":

  • この宣言は、C++の名前修飾を防ぎ、関数をC言語形式でエクスポートするためのものです。これにより、Pythonなどの他言語からこの関数を正しく呼び出すことができます。
  • Ref: extern "C"をなんとなく付けている人へ #C++ - Qiita

• void vector_add(float *a, float *b, float *c, int n);:

  • ベクトル加算を行う関数のプロトタイプ宣言です。
  • 引数:
    • float *a: ベクトルAのデータ
    • float *b: ベクトルBのデータ
    • float *c: 結果を格納するベクトルC
    • int n: ベクトルの要素数
  • *は配列の先頭を指すポインタを表します。
    • Ref: C++ポインタまとめ #メモリ管理 - Qiita

CUDAカーネル関数の定義

__global__ void vector_add_kernel(float *a, float *b, float *c, int n) {
    int idx = threadIdx.x + blockDim.x * blockIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

• __global__:

  • この修飾子は、関数がCUDAカーネルであることを示します。GPU上で並列に実行されます。
    • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#kernels

• void vector_add_kernel(float *a, float *b, float *c, int n):

  • GPU上で実行されるベクトル加算カーネル関数です。

• 変数の解説:

  • int idx: 現在のスレッドが処理する要素のインデックスを計算します。
    • threadIdx.x: ブロック内のスレッドのインデックス
    • blockDim.x: ブロック内のスレッド数
    • blockIdx.x: グリッド内のブロックのインデックス
      • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#thread-hierarchy

• ロジック:

  • 各スレッドは一つの要素を処理します。
  • if (idx < n)で、範囲外アクセスを防ぎます。
  • c[idx] = a[idx] + b[idx];で対応する要素の加算を行います。

ホスト関数の実装

void vector_add(float *a, float *b, float *c, int n) {
    float *d_a, *d_b, *d_c;
    size_t size = n * sizeof(float);

    // デバイスメモリの割り当て
    cudaMalloc((void**)&d_a, size);
    cudaMalloc((void**)&d_b, size);
    cudaMalloc((void**)&d_c, size);

    // ホストからデバイスへデータを転送
    cudaMemcpy(d_a, a, size, cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, b, size, cudaMemcpyHostToDevice);

    // カーネルの起動
    int threadsPerBlock = 256;
    int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock;
    vector_add_kernel<<<blocksPerGrid, threadsPerBlock>>>(d_a, d_b, d_c, n);

    // デバイスからホストへ結果を転送
    cudaMemcpy(c, d_c, size, cudaMemcpyDeviceToHost);

    // デバイスメモリの解放
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);
}

• void vector_add(...):

  • ホスト側で呼び出される関数で、GPU上での計算を管理します。

• デバイスメモリの割り当て:

  • cudaMalloc((void**)&d_a, size);:
    • GPUメモリ上にベクトルAの領域を確保します。
    • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#device-memory

• データ転送:

  • cudaMemcpy(d_a, a, size, cudaMemcpyHostToDevice);:
    • ホストメモリからデバイスメモリへデータを転送します。
    • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#device-memory

• カーネルの起動:

  • vector_add_kernel<<<blocksPerGrid, threadsPerBlock>>>(d_a, d_b, d_c, n);:
    • CUDA特有のシンタックスでカーネルを起動します。
    • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#kernels

• 結果の取得:

  • cudaMemcpy(c, d_c, size, cudaMemcpyDeviceToHost);:
    • デバイスメモリからホストメモリへ結果を転送します。
    • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#device-memory

• メモリの解放:

  • cudaFree(d_a);などで、GPUメモリを解放します。
  • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#device-memory

2. CUDAコードのコンパイル方法

nvcc --shared -Xcompiler -fPIC -o libvector_add.so vector_add.cu

• nvcc:

  • NVIDIAのCUDAコンパイラです。
  • Ref: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#compilation-with-nvcc

• オプションの解説:

  • --shared:
    • 共有ライブラリ（.soファイル）を生成します。
  • -Xcompiler -fPIC:
    • コンパイラにポジション独立コード（PIC）を生成するよう指示します。共有ライブラリを作成する際に必要です。
  • -o libvector_add.so:
    • 出力ファイル名を指定します。
  • vector_add.cu:
    • コンパイル対象のCUDAソースファイルです。

Ref: C++ユーザーの為のリンクの話２ #C++ - Qiita

3. Pythonコードの解説（main.py）

PythonからCUDAで作成した共有ライブラリを呼び出すコードを見ていきましょう。

共有ライブラリのロード

import numpy as np
import ctypes

# 共有ライブラリをロード
lib = ctypes.cdll.LoadLibrary('./libvector_add.so')

• ctypesモジュール:

  • PythonからC言語で書かれた共有ライブラリを呼び出すための標準ライブラリです。

• LoadLibrary関数:

  • 指定した共有ライブラリをロードします。

関数の引数と戻り値の型指定

lib.vector_add.argtypes = [
    np.ctypeslib.ndpointer(dtype=np.float32, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float32, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float32, flags='C_CONTIGUOUS'),
    ctypes.c_int
]
lib.vector_add.restype = None

• argtypes:

  • 関数の引数の型を指定します。これにより、ctypesが正しくデータを渡すことができます。
  • Ref: https://docs.python.org/ja/3/library/ctypes.html#ctypes._CFuncPtr.argtypes

• np.ctypeslib.ndpointer:

  • NumPy配列をCのポインタとして渡すためのヘルパー関数です。
  • 引数:
    • dtype=np.float32: データ型を32ビット浮動小数点数に指定。
    • flags='C_CONTIGUOUS': メモリが連続していることを保証。
      　 - Ref: https://numpy.org/doc/2.1/reference/routines.ctypeslib.html#numpy.ctypeslib.ndpointer

• restype:

  • 関数の戻り値の型を指定します。Noneの場合、戻り値はないことを示します。
  • Ref: https://docs.python.org/ja/3/library/ctypes.html#return-types

データの準備と関数の呼び出し

# データの準備
N = 1000000
a = np.random.rand(N).astype(np.float32)
b = np.random.rand(N).astype(np.float32)
c = np.zeros(N, dtype=np.float32)

# CUDA関数の呼び出し
lib.vector_add(a, b, c, N)

• データの準備:

  • N = 1000000:
    • 要素数100万のベクトルを生成します。
  • np.random.rand(N):
    • 0から1の間の乱数を生成します。
  • .astype(np.float32):
    • データ型を32ビット浮動小数点数に変換します。

• 関数の呼び出し:

  • lib.vector_add(a, b, c, N):
    • 先ほど定義したCUDA関数を呼び出します。

結果の検証

# 結果の検証
if np.allclose(c, a + b):
    print("計算結果は正しいです。")
else:
    print("計算結果に誤りがあります。")

• np.allclose:

  • 2つの配列が近似的に等しいかを判定します。
  • https://numpy.org/doc/2.1/reference/generated/numpy.allclose.html

• ロジック:

  • CUDAで計算した結果cと、NumPyで計算したa + bを比較し、一致すれば成功とします。

4. まとめ

本記事では、PythonからCUDAを利用してベクトル加算を行う方法について解説しました。CUDAコードでは、GPUメモリの管理やカーネル関数の起動方法など、CUDA特有の文法とロジックを説明しました。Pythonコードでは、ctypesを用いてCUDAで作成した共有ライブラリをロードし、データを渡して計算結果を取得する方法を示しました。

5. 感想

久々にCUDAを触りましたが、ポインタなどの低レベルな操作が必要なため、C++の基礎知識があると理解しやすいと感じました。また、Pythonとの連携も意外とかんたんにできることを知ったので、今後も活用していきたいです。