🤖

Numbaって何者やねん ― PythonをCUDAに変える魔法

に公開
GPU
Numba
tech

Numbaって何者やねん ― PythonをCUDAに変える魔法

前回のあらすじ

前の記事で「NumPy + NumbaでGPU計算できる」って書いた。

でも冷静に考えたら「Numbaって何してんの？」って話になる。NumPyはただの配列やし、CUDAはNVIDIAのGPU用言語やし。その間を埋めてるNumbaって何者なん？

Numbaの正体：PythonをネイティブコードにするJITコンパイラ

Numbaは PythonコードをLLVM経由でネイティブコード（機械語）に変換する ライブラリ。

普通のPython：

Pythonコード → インタプリタが1行ずつ解釈 → 実行（遅い）

Numba使用時：

Pythonコード → LLVM IR → ネイティブコード → 実行（速い）

LLVMってのはコンパイラの中間表現。ClangとかRustとかSwiftも使ってる、業界標準のコンパイラ基盤。

@jit ― CPU高速化の魔法

まずCPU向けの話から。

from numba import jit
import numpy as np

@jit(nopython=True)
def sum_array(arr):
    total = 0.0
    for i in range(arr.shape[0]):
        total += arr[i]
    return total

data = np.random.rand(10000000)
result = sum_array(data)  # 初回：コンパイル＋実行、2回目以降：キャッシュから実行

@jit つけるだけで、このPythonのforループがC言語並みの速度で動く。

なんで速くなるん？

普通のPythonが遅い理由：

  1. 動的型付け: 毎回「この変数の型なんやっけ」って確認してる
  2. インタプリタ実行: 1行ずつ解釈しながら実行
  3. オブジェクトのオーバーヘッド: Python的なint型はただの数字やなくてオブジェクト

Numbaがやってること：

  1. 型推論: 初回実行時に「この変数はfloat64やな」って特定
  2. コンパイル: 型が決まったらLLVM経由でネイティブコードに変換
  3. キャッシュ: 2回目以降は機械語を直接実行

要するに 「Pythonの皮を被ったC」 を作ってる。

@cuda.jit ― GPU向けの魔法

で、本題。GPU向けには @cuda.jit を使う。

from numba import cuda
import numpy as np

@cuda.jit
def gpu_add(a, b, result):
    idx = cuda.grid(1)
    if idx < a.size:
        result[idx] = a[idx] + b[idx]

# NumPy配列を作る
a = np.array([1, 2, 3, 4, 5], dtype=np.float32)
b = np.array([10, 20, 30, 40, 50], dtype=np.float32)
result = np.zeros(5, dtype=np.float32)

# GPUに転送
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_result = cuda.to_device(result)

# カーネル実行（1ブロック、5スレッド）
gpu_add[1, 5](d_a, d_b, d_result)

# 結果を取得
result = d_result.copy_to_host()
print(result)  # [11. 22. 33. 44. 55.]

Numbaが裏でやってること

  1. Pythonコードを解析: @cuda.jit がついた関数を見つける
  2. CUDA PTXに変換: PythonコードをGPU用の中間言語（PTX）にコンパイル
  3. カーネル登録: NVIDIAのドライバにカーネルを登録
  4. 実行時に呼び出し: gpu_add[1, 5](...) でGPU上でカーネル実行

つまり、PythonがCUDA Cに変身してる。

nvccとの違い

普通にCUDAやろうと思ったら：

// これがCUDA C
__global__ void gpu_add(float *a, float *b, float *result, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        result[idx] = a[idx] + b[idx];
    }
}

これを nvcc でコンパイルして、C/C++のホストコードから呼び出す。

めんどくさい。

Numbaなら Pythonのまま書ける。裏でPTXに変換してくれる。

なんでNumPyと相性がいいのか

NumPyの配列（ndarray）は内部的には 連続したメモリ領域 + 型情報 + 形状情報。

arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)

これ、メモリ上では：

[8バイト][8バイト][8バイト]  ← float64が3つ並んでるだけ

C言語の配列と同じ構造。だからNumbaが型推論しやすいし、GPUに転送するときもそのままコピーできる。

Pythonのリストだとこうはいかない：

lst = [1.0, 2.0, 3.0]

これ、内部的には オブジェクトへのポインタの配列。型もバラバラになりうる。GPUに送るには変換が必要。

NumPy配列は「Cっぽいデータ構造」だから、Numbaと相性がいい。

CuPyとの違い

じゃあCuPyは何が違うのか。

CuPy: NumPyのAPIをGPU上で再実装したもの

import cupy as cp

a = cp.array([1, 2, 3])
b = cp.array([4, 5, 6])
c = a + b  # GPU上で計算される

NumPyと同じ書き方で、裏でGPU計算してくれる。便利。

Numba CUDA: 自分でカーネル書く

@cuda.jit
def my_kernel(arr):
    # 自分でスレッドIDとか管理
    idx = cuda.grid(1)
    ...

細かい制御ができる。

使い分け

やりたいこと 選択
NumPyのコードをそのままGPU化 CuPy
細かいメモリ制御したい Numba CUDA
カスタムカーネル書きたい Numba CUDA
とりあえず速くしたい CuPy

俺の場合、メモリレイアウトを細かく制御したかったからNumba CUDAにした。

まとめ：Numbaがやってること

  1. Pythonコードを解析
  2. 型を推論
  3. LLVM IR（CPU）またはPTX（GPU）に変換
  4. ネイティブコード / CUDAカーネルを生成
  5. 実行時にキャッシュから高速実行

要するに「Pythonを書いてるつもりで、裏でCやCUDA Cに変換されてる」。

だからNumPyだけでGPU計算できる。CuPyは要らん。Numbaが 翻訳者 になってくれてる。

余談：なんで「Numba」って名前なん

"Number" + "Mamba"（速い蛇）らしい。Pythonが蛇やから。

...だと思ってたけど、公式には "Numba" の由来は明言されてないっぽい。まあ、数値計算速くする蛇、みたいなイメージなんやろな。

Numbaの正体を理解した日の記録

Discussion