Python高速化！数値計算向けライブラリ比較【Numba, Taichi, Warp, JAX】

## はじめに

Pythonで数値計算や物理シミュレーションを行う際、処理速度が課題になります。

本記事では、Pythonコードを大幅に高速化できる4つの主要フレームワークの特徴と性能を比較し、用途に応じた最適な選択肢を紹介します。

マンデルブロ集合計算のベンチマークを通じて、各フレームワークのCPU/GPU性能を検証した結果、用途によって最適なフレームワークが異なることがわかりました。特に**WarpはGPU処理で最高速**を示し、**NumbaはCPU処理と既存コード互換性に優れている**ことが確認できました。

[![Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/YanagiAyame/python-mandelbrot-benchmark/blob/main/mandelbrot-benchmark.ipynb)

![download.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665667/edace7ed-2227-4386-aa0f-45e2d81f458c.png)

## 各種フレームワークの特徴

### [Numba](https://numba.readthedocs.io/en/stable/index.html)
**シンプルで使いやすいJITコンパイラ**

:::note info
NumbaはPython関数をLLVMコンパイラを用いて機械語に変換し、処理を大幅に高速化するフレームワークです。特に既存のNumPyコードへの適用が容易で、単純なデコレータの追加だけで効果を得られる点が大きな特徴です。
:::
  - デコレータ（@jit）で関数を簡単に最適化
  - NumPyとの高い互換性により既存コードの移行が容易
  - CPUマルチコア並列化およびCUDA GPUアクセラレーションに対応
  - 純粋なPythonコードと比較して劇的な速度向上が可能
### [Taichi](https://docs.taichi-lang.org/)
**クロスプラットフォーム対応の高性能フレームワーク**

:::note info
Taichiは物理シミュレーション向けに設計された言語で、複数のハードウェアプラットフォームをサポートしています。CPUとGPUの両方で効率的に動作するコードを生成する能力を持っています。
:::
  - CPU/GPU向けに自動最適化されたコードを生成
  - CUDA, Vulkan, Metal, OpenGLなど複数のGPUバックエンドをサポート
  - 独自の階層的データ構造（SNode）により複雑なシミュレーションに対応[^SNode]
  - 物理ベースのシミュレーションに最適化された機能を提供
[^SNode]: https://docs.taichi-lang.org/docs/sparse
### [Warp](https://nvidia.github.io/warp/)
**NVIDIA謹製の高性能計算エンジン**

:::note
Warpは、NVIDIA社が開発したフレームワークで、特にNVIDIA GPU上での高速計算に最適化されています。物理シミュレーションと機械学習の統合に強みを持っています。
:::
  - NVIDIA GPUで特に高いパフォーマンスを発揮
  - CPUおよびGPUで実行可能なコード
  - PyTorchやJAXと連携可能な微分可能カーネルを提供
  - シミュレーションのための特化した関数ライブラリを内蔵
### [JAX](https://docs.jax.dev/en/latest/index.html)
**自動微分と並列計算に特化したフレームワーク**

:::note
JAXはGoogleが開発した、NumPyの関数をGPU/TPUで加速しつつ自動微分機能を提供するフレームワークです。機械学習の研究開発でよく使用されています。
:::
  - NumPyの関数をGPU/TPUで高速実行できるよう拡張
  - 自動微分機能により機械学習モデルの最適化に適している
  - [XLA](https://openxla.org/xla?hl=ja)（Accelerated Linear Algebra）によるJITコンパイルで高速化
  - SPMD[^SPMD]（単一プログラム、複数データ）モデルによるマルチデバイス並列処理機能

[^SPMD]: 「Single Program, Multiple Data」の略称。複数のプロセッサが同じプログラムを実行しながら、異なるデータに対して処理を行う並列計算のモデル。

## ベンチマークコード

[マンデルブロ集合](https://ja.wikipedia.org/wiki/%E3%83%9E%E3%83%B3%E3%83%87%E3%83%AB%E3%83%96%E3%83%AD%E9%9B%86%E5%90%88)の計算は、反復計算と条件分岐を含む典型的な数値計算タスクであり、各フレームワークの性能を評価するのに適しています。[^mandelbrot]
[^mandelbrot]: 例えば、[The Computer Language Benchmarks Game](https://benchmarksgame-team.pages.debian.net/benchmarksgame/)で使われています。

以下のGitHubリポジトリで完全なコードを公開しています。

[![Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/YanagiAyame/python-mandelbrot-benchmark/blob/main/mandelbrot-benchmark.ipynb)

https://github.com/YanagiAyame/python-mandelbrot-benchmark/blob/main/mandelbrot-benchmark.ipynb

<details><summary>ベンチマークコードのサンプル（一部抜粋）</summary>

```python
# Numbaのコード例
@numba.vectorize([numba.int32(numba.complex64)], target="parallel")
def run_numba_cpu_ufunc(c):
    counter = 0
    z = 0j
    for _ in range(1000):
        z = z**2 + c
        if z.real**2+z.imag**2 >= 4:
            break
        counter += 1
    return counter

x, y = np.meshgrid(np.linspace(-2, 1.2, 7680, dtype=np.float32), np.linspace(-1.2, 1.2, 4320, dtype=np.float32))
complex_array = x + y * 1j
z = run_numba_cpu_ufunc(complex_array)

# Taichiのコード例
@ti.func
def taichi_func(c: tm.vec2) -> ti.i32:
    counter = 0
    z = tm.vec2(0.0, 0.0)
    for i in range(1000):
        z = tm.vec2(z.x**2 - z.y**2, 2*z.x*z.y) + c
        if z.x**2+z.y**2 >= 4:
            break
        counter += 1
    return counter

@ti.kernel
def run_taichi_kernel(field: ti.template(), out: ti.template()):
    for i, j in field:
        out[i, j] = taichi_func(field[i, j])

def calculate_taichi_mandelbrot(complex_array):
    field = ti.Vector.field(n=2, dtype=ti.f32, shape=(4320, 7680))
    field.from_numpy(complex_array)
    out = ti.field(dtype=ti.i32, shape=(4320, 7680))
    run_taichi_kernel(field, out)
    return out.to_numpy()
ti.init(arch=ti.cpu)

x, y = np.meshgrid(np.linspace(-2, 1.2, 7680), np.linspace(-1.2, 1.2, 4320))
complex_array = np.stack([x, y], axis=2)

z=calculate_taichi_mandelbrot(complex_array)

# Warpのコード例
@wp.func
def warp_func(c: wp.vec2) -> wp.int32:
    counter = wp.int32(0)
    z = wp.vec2(0.0, 0.0)
    for i in range(1000):
        z = wp.vec2(z[0]*z[0] - z[1]*z[1], 2.0*z[0]*z[1]) + c
        if z[0]*z[0]+z[1]*z[1] >= 4.0:
            break
        counter += 1
    return counter

@wp.kernel
def warp_kernel(field: wp.array2d(dtype=wp.vec2), out: wp.array2d(dtype=wp.int32)):
    i, j = wp.tid()
    out[i, j] = warp_func(field[i, j])

def calculate_warp_mandelbrot(complex_array, device):
    field=wp.array(complex_array, dtype=wp.vec2, device=device)
    out=wp.zeros(shape=(4320, 7680), dtype=wp.int32, device=device)
    wp.launch(kernel=warp_kernel, dim=(4320, 7680), inputs=[field], outputs=[out], device=device)
    return out.numpy()

device = "cpu"

z=calculate_warp_mandelbrot(complex_array, device)

# JAXのコード例
def jax_func(c):
    def body(i, carry):
        counter, z = carry
        z = jnp.where(z.real**2 + z.imag**2 >= 4, z, z**2 + c)
        counter = jnp.where(z.real**2 + z.imag**2 >= 4, counter, counter + 1)
        return counter, z
    final_counter, _ = jax.lax.fori_loop(0, 1000, body, (0, 0.+0j))
    return final_counter

x, y = jnp.meshgrid(jnp.linspace(-2, 1.2, 7680, dtype=np.float32), jnp.linspace(-1.2, 1.2, 4320, dtype=np.float32))
complex_array = x + y * 1j

calculate_jax_mandelbrot = jax.jit(jnp.vectorize(jax_func), backend="cpu")
z = calculate_jax_mandelbrot(complex_array).block_until_ready()
```
</details>

## 速度比較結果

Kaggleノートブック上で、 NVIDIA Tesla P100のアクセラレータを選択して実行しました。

- **CPU最速**: Numba
- **GPU最速**: Warp
- **最遅**: JAX（CPU/GPUともに）

|     | Numba | Taichi | Warp  | JAX  | 
| --- | ----- | ------ | ----- | ---- | 
| CPU | 7.80s  | 13.6s  | 31.8s  | 63.5s | 
| GPU | 0.837s | 0.708s  | 0.310s | 2.34s | 

![download.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665667/edace7ed-2227-4386-aa0f-45e2d81f458c.png)

## フレームワーク選択の指針

1. 既存コードがNumPyベースの場合→**Numba**
    - Numbaが最適です。簡単な変更で大幅な高速化が可能です
1. NVIDIA GPU専用環境の場合→**Warp**
    - Warpが最も効率的です。特にPyTorchやJAXとの連携が必要な場合にも有効です
1. ポータビリティ重視の場合→**Taichi**
    - Taichiは複数GPUバックエンドをサポートしており、幅広い環境で利用可能です
1. 巨大テンソル計算や機械学習モデルの場合→**JAX**
    - JAXは今回のベンチマークでは遅かったものの、自動微分やSPMD並列処理が強みとなるため、特定用途では有効だと思います
    - ただし、本ベンチマークのようなステートフルな計算（かつ遷移回数も多い場合）には不向きと考えられます

:::note
物理シミュレーションプラットフォームで選ぶのもいいかもしれません。
[Omnivese](https://www.nvidia.com/ja-jp/omniverse/)⇔Warp
[Genesis](https://genesis-embodied-ai.github.io/)⇔Taichi
（資本的にはNVIDIA Omniverseがよさそうだけど、NVIDIAへの依存度が...）
:::

## おまけ（実装上の注意点）

共通の最適化ポイント
- データ転送の最小化: CPU-GPU間のデータ転送は極力減らす
- メモリアクセスパターン: 連続したメモリアクセスを心がける

フレームワーク別のヒント
- Numba: parallel=Trueオプションでマルチコア活用、fastmath=Trueで数値精度と引き換えに速度向上（するかも）
- Taichi: 適切なデータ構造（sparse/denseなど）選択が重要
- Warp: NVIDIAの最新GPUドライバー利用でパフォーマンス向上
- JAX: 関数型スタイルを採用し、状態変化を最小限に