【CuPy入門 第3回】CuPyの真髄!カスタムカーネルで独自のGPU高速化に挑戦 🚀

「NumPyやCuPyの既存の関数(ufunc)だけでは表現できない、独自の複雑な計算をGPUで高速化したい!」
if/elseを含むような処理を、配列の全要素に対して一気に実行できないかな?」
「GPUプログラミングって聞くと難しそうだけど、Pythonから少しでも触ってみたい!」

こんにちは! CuPy探検隊、隊長のPythonistaです! 前回の第2回では、CuPy配列の基本的な操作がNumPyと驚くほど同じ感覚で行えること、そしてその上でGPUによる圧倒的な計算速度を体感しましたね。これで、あなたのNumPyスキルがGPU上でも強力な武器になることが分かったはずです。

シリーズ第3回の今回は、CuPyのさらに奥深く、そして真髄とも言える機能、「カスタムカーネル」の世界に足を踏み入れます。カスタムカーネルとは、あなた自身がGPU上で実行される処理(カーネル)を記述し、独自の高速な演算を定義するための仕組みです。これを使えば、CuPyが単なるNumPyの代替ではなく、より高度なGPUプログラミングへの入り口であることが理解できるでしょう。少し発展的な内容ですが、この記事を読み終える頃には、あなたも「GPUプログラマー」としての第一歩を踏み出しているはずです!


1. 既存の関数を超えるために:なぜカスタムカーネルが必要か?

cp.sin()cp.sum()といったCuPyのユニバーサル関数(ufunc)は、内部で高度に最適化されており非常に高速です。しかし、私たちが実行したい計算は、必ずしも既存の関数だけで表現できるとは限りません。

例えば、「配列の各要素に対して、もし値が0以上なら2倍し、負の値なら0.5倍する」といった条件分岐を含む独自の処理を考えます。

NumPyやCuPyの既存の関数を組み合わせることでも実現は可能です。

import cupy as cp
import numpy as np

x = cp.array([-2.0, -1.0, 0.0, 1.0, 2.0])

# np.where (CuPy版のcp.where) を使う方法
y = cp.where(x >= 0, x * 2, x * 0.5)
print(f"cp.whereを使った結果: {y}") # [-1.  -0.5  0.   2.   4. ]

この方法は簡潔で非常に良い方法ですが、条件分岐がさらに複雑になったり、複数のステップを含むような処理を各要素に対して行いたい場合、既存の関数を何度も呼び出すことになり、データのやり取りで非効率が生じる可能性があります。

カスタムカーネルは、このような「独自の要素ごとの処理」を一つのまとまった操作として定義し、それをGPU上で直接、何千ものコアを使って超並列に実行させてしまおう、という考え方です。これにより、柔軟性と最高のパフォーマンスを両立させることが可能になります。


2. あなたの最初のカスタムカーネル:cp.ElementwiseKernel

cp.ElementwiseKernelは、カスタムカーネルを作成するための最も簡単な方法です。これは、入力配列の各要素に対して独立した同じ処理を行い、結果を出力配列の対応する位置に書き込む、という「要素ごと(Element-wise)」のカーネルを簡単に定義するための機能です。

ループ処理(for i in range(len(x)): ...)の部分はCuPyが裏側で自動的に並列化してくれるので、私たちは「1つの要素に対して何をするか」だけを考えれば良いのです。

2.1. ElementwiseKernelの構造

基本的な構文は以下の通りです。

my_kernel = cp.ElementwiseKernel(
    '入力引数 (型と名前)',
    '出力引数 (型と名前)',
    '各要素に対する処理内容 (C言語風のコード)',
    'カーネル名'
)
  • 入力引数: 'float32 x, float32 y' のように、型と変数名を文字列で指定します。
  • 出力引数: 'float32 z' のように、出力用の変数を指定します。
  • 処理内容: 1つの要素に対する処理を、C言語に似た構文で文字列として記述します。セミコロン;で終わる必要があります。
  • カーネル名: デバッグ時などに使われるカーネルの名前を文字列で指定します。

2.2. 実践例:条件付きスケーリングを行うカーネル

先ほどの「値が0以上ならA倍、負ならB倍する」という処理を、ElementwiseKernelで実装してみましょう。

import cupy as cp

# カスタムカーネルを定義
conditional_scale_kernel = cp.ElementwiseKernel(
    'float32 x, float32 positive_scale, float32 negative_scale', # 入力: 要素x, 正の場合の倍率, 負の場合の倍率
    'float32 y', # 出力: 結果y
    '''
        if (x >= 0) {
            y = x * positive_scale;
        } else {
            y = x * negative_scale;
        }
    ''',
    'conditional_scale' # カーネル名
)

# GPU上でデータを用意
x_gpu = cp.array([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=cp.float32)
pos_scale_factor = 10.0 # 正の値には10を掛ける
neg_scale_factor = 0.1  # 負の値には0.1を掛ける

# 作成したカーネルを通常の関数のように呼び出す
y_gpu = conditional_scale_kernel(x_gpu, pos_scale_factor, neg_scale_factor)

print(f"\n--- ElementwiseKernelの実行結果 ---")
print(f"入力配列 x: {x_gpu}")
print(f"出力配列 y: {y_gpu}")

# C言語風の三項演算子を使うと、より簡潔に書けます
# 'y = (x >= 0) ? x * positive_scale : x * negative_scale;'

実行結果 (例):

--- ElementwiseKernelの実行結果 ---
入力配列 x: [-2. -1.  0.  1.  2.]
出力配列 y: [-0.2 -0.1  0.  10.  20. ]

このように、Pythonのループを書くことなく、独自の条件分岐を含む処理をGPU上で並列実行させることができました!


3. さらなる自由度へ:cp.RawKernelの世界を覗いてみる

ElementwiseKernelは非常に手軽ですが、その名の通り「要素ごと」の独立した処理に限定されます。もし、近傍の要素を参照したり、スレッド間でデータを共有したり、より複雑なメモリアクセスパターンを実装したりしたい場合はどうすればよいでしょうか?

その答えがcp.RawKernelです。これは、あなたが書いた生のCUDA C/C++カーネルコードを、ほぼそのままPythonから呼び出して実行できるようにするものです。

3.1. RawKernelとは?

cp.RawKernelを使うと、GPUプログラミングのより深いレベルにアクセスできます。スレッドのインデックス(threadIdx.x, blockIdx.xなど)を自分で計算し、どのスレッドが配列のどの要素を担当するかを明示的にコーディングします。これにより、最大限の柔軟性とパフォーマンスを引き出すことが可能になります。

注意: RawKernelを使いこなすには、CUDAプログラミングの基本的な知識(スレッド、ブロック、グリッド、共有メモリなど)が必要となります。ここでは、その雰囲気を感じてもらうための「入門の入門」として、非常に簡単な例だけを紹介します。

3.2. 簡単な例:2つのベクトルの足し算

import cupy as cp

# CUDA C/C++のカーネルコードを文字列として定義
vector_add_kernel_code = r'''
extern "C" __global__
void my_add(const float* x1, const float* x2, float* y, int N) {
    int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < N) {
        y[tid] = x1[tid] + x2[tid];
    }
}
'''

# RawKernelオブジェクトを作成
vector_add_kernel = cp.RawKernel(vector_add_kernel_code, 'my_add')

# データを準備
N = 10
x1 = cp.arange(N, dtype=cp.float32)
x2 = cp.arange(N, dtype=cp.float32) * 2
y = cp.zeros(N, dtype=cp.float32)

# カーネルを実行
# 実行するスレッドのグリッドとブロックのサイズを指定する必要がある
grid_size = (2,)  # 2つのブロックを使う
block_size = (5,) # 1ブロックあたり5つのスレッドを使う (合計10スレッド)
args = (x1, x2, y, N) # カーネルに渡す引数

vector_add_kernel(grid_size, block_size, args)

print(f"\n--- RawKernelの実行結果 ---")
print(f"入力 x1: {x1}")
print(f"入力 x2: {x2}")
print(f"出力 y : {y}")

実行結果 (例):

--- RawKernelの実行結果 ---
入力 x1: [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
入力 x2: [ 0.  2.  4.  6.  8. 10. 12. 14. 16. 18.]
出力 y : [ 0.  3.  6.  9. 12. 15. 18. 21. 24. 27.]

この例では、CUDAカーネルのコードを書き、それを実行するためのスレッドの構成(グリッドサイズ、ブロックサイズ)を指定しました。このように、RawKernelはより低レベルな制御を可能にしますが、その分CUDAに関する深い知識が求められます。しかし、これがCuPyの持つポテンシャルの頂点であり、究極のパフォーマンスを引き出すための道筋です。


まとめ:あなたはもうGPUプログラマー!

今回は、CuPyの真髄とも言える「カスタムカーネル」の世界に足を踏み入れました。

  • 既存のユニバーサル関数では表現できない、独自の複雑な処理をGPUで高速化するためにカスタムカーネルが必要であること。
  • cp.ElementwiseKernelを使えば、C言語風の簡単なコード片を記述するだけで、手軽に要素ごとのカスタム処理を定義できること。
  • cp.RawKernelを使えば、完全なCUDA C/C++コードをPythonから実行でき、最大限の柔軟性とパフォーマンスを引き出せること(ただしCUDAの知識が必要)。

カスタムカーネルは、CuPyを単なる「NumPyの高速な代替品」から、「PythonからGPUの並列処理能力を直接引き出すための強力なインターフェース」へと昇華させる機能です。ElementwiseKernelから始めて、少しずつGPUプログラミングの感覚を掴んでいくと、解決できる問題の幅が大きく広がるでしょう。

さて、CuPyの基本操作からカスタムカーネルまで、その使い方とパワーの一端を見てきました。次回からはいよいよ「性能比較編」に突入します! どのような条件下でCuPyがNumPyを圧倒するのか、そしてデータ転送コストを含めるとどうなるのか、具体的なベンチマークを通して、その秘密に迫っていきます。お楽しみに!

その他の投稿はこちら

コメント

このブログの人気の投稿

タイトルまとめ

これで迷わない!WindowsでPython環境構築する一番やさしい方法 #0

【Python標準ライブラリ完結!】11の冒険をありがとう!君のPython力が飛躍する「次の一歩」とは? 🚀