Python互換の高速言語MojoでGPUを叩いて、Pythonと速度を比べてみた話

Last updated at 2026-08-12Posted at 2026-08-12

TL;DR

「Python互換の高速言語」を謳うMojoが、GPU向けの実行基盤「MAX」を経てNVIDIA GPUに対応しているという話をPublickeyの過去記事で知り、実際にGPU計算をやらせて素のPythonとの速度差を体感してみました
題材は800〜2,000個の粒子が互いの重力で引き合うN体シミュレーション。表の数字だけだとピンと来ないので、実際に動く様子をGPUで計算してGIFにもしました
2,000粒子・10ステップ分の重力計算で比較したところ、純Pythonのforループが約12.7秒、NumPyのベクトル化で約382ms、MojoのGPUカーネルは約1.0msという結果に。カーネル計算だけならNumPyの約380倍、純Pythonの約12,500倍でした
ただしMojoは実行するたびにコンパイルが走る言語で、コンパイル込みのプロセス全体では約4.8秒。粒子数が少ないうち(実測でN≈17,000くらいまで)はNumPyの方が総合的に速いという、看板の速さだけでは見えない一面も確認しました
検証時点の最新版はMojo 1.0.0b2(2026年6月リリース)で、まだ正式版ではなくベータ版です
「Python互換」を鵜呑みにせず実際に確認したところ、GPUカーネル自体は型注釈・comptimeパラメータ・スレッド/ブロックIDなどPython本体にはない概念が必須で、「Pythonの書き味を借りた新しいGPU向け言語」というのが実態でした
GPUが本当に使われているかも DeviceContext の情報とnvidia-smiの実測で確認し、デバイス名が実機のRTX 3050と一致すること、実行中にGPU使用率・メモリ使用量が上がることを確かめました

やろうとしたこと

きっかけはPublickeyの記事「Python互換の高速言語Mojo、GPU向けの実行環境「MAX GPUs」を発表」でした。2023年のModCon '23で発表された当時の内容ですが、「単一の言語でCPUとGPUの両方を、しかもPython互換の書き味で扱える」というコンセプトは今も変わらず面白いなと思い、今の版のMojoで実際に試してみることにしました。

最初は要素ごとの数式(sin・cos・sqrt)をひたすら計算するだけのベンチマークを書いていたのですが、それだと**「表の数字が大きい」以上の実感が湧きにくいなと思い直しました。GPUの並列計算のすごさが一番わかりやすく伝わるのは、大量の物体が互いに影響し合いながら動く様子を目で見られる物理演算だろうということで、題材をN体(Nボディ)重力シミュレーション**に変更しました。すべての粒子同士が万有引力で引き合う、計算量 $O(N^2)$ の古典的な題材です。

構成/環境

検証環境はこんな感じです。

OS: Windows 11(WSL2 Ubuntu 24.04.3 LTS上で実行)
GPU: NVIDIA GeForce RTX 3050 Laptop(VRAM 6GB)
Mojo: 1.0.0b2(2026年6月リリース、正式版ではなくベータ版)、pixi経由でインストール
Python: 3.12.3 + NumPy 2.5.1 + Matplotlib/Pillow(可視化用、venv環境)

curl -fsSL https://pixi.sh/install.sh | sh
pixi init mojo-gpu-test -c https://conda.modular.com/max/ -c conda-forge
cd mojo-gpu-test
pixi add mojo

やってみた結果

環境構築: やっぱりWindowsではなくWSLだった

Mojoの公式ドキュメントには前提としてこう書かれています。

"You can use Mac, Linux, or Windows with WSL."

つまりWindows単体はサポート外で、WSL経由が公式ルートです。GPU自体はWSL2からも nvidia-smi でちゃんと見えるので、GPUパススルー周りのセットアップは特に何もしなくても機能しました。

本題: N体重力シミュレーションのGPUカーネルを書く

やることはシンプルです。各粒子について「自分以外の全粒子から受ける引力を合計して加速度を求め、速度と位置を更新する」という計算を、粒子1個 = スレッド1個として全粒子分並列に実行します。

from std.math import ceildiv, sqrt
from std.gpu.host import DeviceContext
from std.gpu import block_dim, block_idx, thread_idx
from layout import TileTensor, row_major

comptime float_dtype = DType.float32
comptime n = 2000
comptime layout = row_major[n]()
comptime g_const: Float32 = 1.0
comptime eps2: Float32 = 0.5   # 近距離で加速度が発散しないための緩和項
comptime dt: Float32 = 0.01

def nbody_step(
    x_in: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
    y_in: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
    vx: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
    vy: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
    mass: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
    x_out: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
    y_out: TileTensor[float_dtype, type_of(layout), MutAnyOrigin],
):
    var i = block_idx.x * block_dim.x + thread_idx.x
    if i < n:
        var xi = x_in[i]
        var yi = y_in[i]
        var ax: Float32 = 0.0
        var ay: Float32 = 0.0
        for j in range(n):
            var dx = x_in[j] - xi
            var dy = y_in[j] - yi
            var dist2 = dx * dx + dy * dy + eps2
            var inv_dist3 = 1.0 / (dist2 * sqrt(dist2))
            ax += g_const * mass[j] * dx * inv_dist3
            ay += g_const * mass[j] * dy * inv_dist3
        var nvx = vx[i] + ax * dt
        var nvy = vy[i] + ay * dt
        vx[i] = nvx
        vy[i] = nvy
        x_out[i] = xi + nvx * dt
        y_out[i] = yi + nvy * dt

1点ハマりどころだったのが、「今の位置」を全スレッドが読み終わる前に誰かが「次の位置」を書き始めてしまう、というデータ競合です。位置バッファを1本しか使わないと、あるスレッドが他の粒子の位置を読んでいる最中に、別のスレッドがその値を新しい位置で上書きしてしまうことがあります。これを避けるため、**位置バッファを2本(A/B)用意して、奇数ステップと偶数ステップで読み書きを入れ替える(ping-pong)**構成にしました。

for step in range(num_steps):
    if step % 2 == 0:
        ctx.enqueue_function[nbody_step](
            x_a, y_a, vx, vy, mass, x_b, y_b,   # Aを読んでBに書く
            grid_dim=num_blocks, block_dim=block_size,
        )
    else:
        ctx.enqueue_function[nbody_step](
            x_b, y_b, vx, vy, mass, x_a, y_a,   # Bを読んでAに書く
            grid_dim=num_blocks, block_dim=block_size,
        )

速度(vx/vy)は各スレッドが自分の担当インデックスしか読み書きしないので、位置のようなダブルバッファは不要でした。

検証: 「Python互換」は実際どこまで本当か

Publickeyの記事も含めてMojoは「Python互換の高速言語」と紹介されがちですが、実際に書いてみると額面通りには受け取れませんでした。

CPU側の def main() raises: や for step in range(num_steps):、print(...) はほぼそのままPythonの見た目です
一方GPUカーネル側は、var xi: Float32 のような型注釈が実質必須で、comptime パラメータ、TileTensor、block_idx / thread_idx といったPython本体には存在しない概念を使わないと1行も書けません
素のPythonスクリプトをコピペしてGPU上で動かす、という意味での「互換」ではなく、**「Pythonに近い構文で書けるが、GPU向けに型と並列実行の概念を足した別言語」**というのが実態でした

なので「Python互換」は「文法の見た目がPythonに近い」という意味で捉えるのが正確で、既存のPythonコード資産がそのままGPUで動くという話ではありません。

検証: GPUは本当に使われているのか

has_accelerator() が True を返すだけでは「CPUにフォールバックしていないか」までは確認できないので、DeviceContext から直接デバイス情報を取ってみました。

ctx = DeviceContext()
print("device:", ctx.name())
print("api:", ctx.api())

device: NVIDIA GeForce RTX 3050
api: cuda

実機に載っているGPUの名前(RTX 3050)とCUDA APIが返ってきたので、少なくとも「別のGPU」や「CPUエミュレーション」ではなく手元の実機GPUが使われていることは確認できました。念のため、N体カーネルを2万ステップ連続実行してGPUを1.5秒ほどビジーにしながら nvidia-smi を並行して覗いてみたところ、待機時は約1.2GB・0%前後だったGPU使用率/メモリ使用量が、実行中は使用率16〜25%、メモリ約1.5GBまで上がることも確認できました。数値が控えめなのは、粒子数2,000程度だとRTX 3050の数千コアを埋めきれていないからだと思いますが、「ちゃんとGPUで計算している」という裏は取れたと思います。

計測結果: 純Python・NumPy・Mojo(GPU)を比べる

2,000粒子で10ステップ分のシミュレーションを、std.time.perf_counter_ns (Mojo側)と time.perf_counter (Python側)で計測しました。3実装とも同じ初期条件・同じ物理式で、最終的な粒子座標が一致することも確認済みです。

実装	処理時間	1ステップあたり
純Python(forループ)	約12,679ms	約1,268ms
NumPy(ベクトル化、ブロードキャスト)	約382ms	約38.2ms
Mojo GPUカーネルのみ	約1.0ms	約0.10ms
✅ Mojoプロセス全体	約4,803ms	コンパイル・pixi起動込みの実測(`time`コマンド、キャッシュなしの初回コンパイル)

カーネル実行時間だけで見るとNumPyの約380倍、純Pythonの約12,500倍。 ただし、コンパイル時間も込みのプロセス全体で比べると話は変わります。純Pythonのループ単体(12.7秒)よりは約2.6倍速いものの、NumPy(382ms)の方がMojoのプロセス全体(4,803ms)より速いという結果でした。N体問題は計算量がO(N²)で伸びる重い処理なので要素ごとの単純な計算よりは差が際立ちやすい題材でしたが、それでも「コンパイル込みの1回勝負」ではNumPyに軍配が上がります。

コンパイル込みだと、NumPyとはどこで逆転するのか

前節の結果を見て「1msで終わる計算のためにコンパイルに4.8秒もかけるなら、NumPyの方がマシでは」と思うのは自然な疑問だと思います。実際その通りで、粒子数が少ないうちはNumPyの方が速いというのが正直なところです。

ポイントは、コンパイル時間は粒子数Nにほとんど依存せず一定(実測で約4,300〜4,800ms)だが、NumPyの計算時間はO(N²)で増え続けることです。つまりNが十分大きくなれば、いつかは逆転するはずです。そこで粒子数Nを500〜20,000まで振って、「NumPy(計算のみ)」「Mojo合計(毎回コンパイル)」「Mojoカーネルのみ(コンパイル不要と仮定した場合)」の3本を計測しました。

実測では、N≈17,000あたりでNumPyとMojo合計(コンパイル込み)が逆転しました。それより粒子数が少ない領域では、コンパイル待ちが発生しないNumPyの方が総合的に速く、N≈17,000を超えるとO(N²)で重くなり続けるNumPyをMojoのほぼ一定のコンパイルコストが逆転する、という構図です。「Mojoカーネルのみ」(コンパイルが要らないと仮定した場合、つまり同じプロセス内で何度も呼び出す場面に相当)は、測定した全域で3ms未満でした。この線がMojo/GPUの本来の実力で、コンパイルコストを償却できる使い方(巨大な計算、または同じプロセス内での繰り返し実行)であれば、NumPyとは比較にならない差がつきます。

実際に動かしてGIFにしてみた

数字だけだとやはり実感が湧きにくいので、GPU上で800粒子・200ステップのシミュレーションを実際に走らせて、途中経過を画像化してGIFにしました。初期状態は円盤状に粒子を配置し、崩壊しすぎないよう軽く公転方向の初速を与えています。

重力で集まったり広がったりを繰り返しながら渦を巻く様子が、GPUなら数百ステップ分でも一瞬で計算できます。同じ800粒子・200ステップを純Pythonのforループで計算すると、実測で1ステップ約182ms(800粒子換算)かかったので、単純計算で約36秒。体感できるほど待たされるかどうかの境界線上ですが、GPU版は待ち時間ゼロで一瞬で終わる、という差でした。

学んだこと

「Python互換」は「文法がPythonに近い」という意味で、既存Pythonコードがそのまま動くわけではない。 GPUカーネル側は型注釈・comptime・スレッド/ブロックIDなどPythonにない概念が必須で、実態は「Python風の新しいGPU向け言語」
一方でGPUが実際に使われていることはDeviceContextのデバイス情報とnvidia-smiの使用率/メモリ使用量の実測で裏取りできた。 「動いているように見えるがCPUで計算していた」という落とし穴は今回はなかった
速度そのものは看板通りで、カーネル実行だけならNumPyの約380倍、純Pythonの約12,500倍という桁違いの結果だった。O(N²)の重い計算ほど差が際立つ
ただしMojoは毎回コンパイルが走るので、「コンパイル込みの1回勝負」だとNumPyに負ける場面がある。 実測ではN≈17,000あたりまではNumPyの方が総合的に速かった。GPUの速さを活かすには、コンパイルコストを償却できるくらい計算量が大きいか、同じプロセス内で繰り返し実行する使い方が必要
抽象的な数式のベンチマークより、物理演算のように「見える」題材の方がGPUのすごさが伝わりやすい。 表の数字だけでなく、実際に動く映像にしてみる価値があった
並列カーネルではデータ競合に注意が必要。 今回のN体シミュレーションでは、位置バッファをダブルバッファにしないと「読んでいる最中に上書きされる」問題が起きる。GPUプログラミング特有の落とし穴だった
Mojo 1.0.0b2はまだベータ版。 正式版ではないことを踏まえて試す必要がある
WindowsからMojoのGPU機能を試すならWSL一択。 GPUパススルー自体はWSL2で素直に機能した

おわりに

「Python互換の高速言語」というキャッチコピーを鵜呑みにせず、実際に手を動かして「どこまでPythonでどこからGPU専用の別言語か」「本当にGPUで計算しているのか」を確認できたのが今回の収穫でした。数式のベンチマークだけで終わらせず、粒子が渦を巻く様子をGIFにしてみたことで、GPUの並列計算のすごさが自分でも実感できました。速度そのものは期待以上でしたが、まだベータ版ゆえに言語仕様が動いている最中なので、今後試す人はその時点の公式ドキュメントを確認しながら進めるのをおすすめします。GPUプログラミングを他の言語(CUDA C++やTritonなど)と比べた話があれば、ぜひコメントで教えてください。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up