[2023] PyO3でPythonを100倍高速化する

最近 free-threading Python を勉強しながら PyO3 に興味が湧いたので、2年前の記事ですが共有します。

Making Python 100× Faster with <100 Lines of Rust – 要約

背景

• 社内の3-D処理パイプラインの中核Pythonライブラリが、同時利用者の増加によりボトルネックになっていた。
• 全体をRustで書き直すにはリスクも時間も大きいため、部分最適化を選択。

アプローチ

1. まず計測: py-spy のサンプリングプロファイラでボトルネックを特定。
2. 段階的なRust導入
   • PyO3 + maturin で Python ↔ Rust を接続。
   • まず find_close_polygons 関数だけを Rust に移植。
   • 続いて Polygon データ構造も Rust に移し、Python からサブクラス化。
3. プロファイリングと改善を反復
   • 不要な NumPy → Rust 変換を最小化。
   • 割り当て・コピーを減らし、距離を直接計算してマイクロ最適化。

性能の変化

段階	平均実行時間 (ms)	改善倍率
初期の純粋な Python	293.41	1×
v1 – 関数のみ Rust (--release)	23.44	12.5×
v2 – Polygon も Rust	6.29	46.5×
v3 – 割り当て削減・直接計算	2.90	101×

主要技術

• PyO3 : 安全な Python ↔ Rust FFI。
• maturin : ビルド・配布の自動化。
• ndarray / numpy crate : Rust側の配列・線形代数。
• py-spy : ネイティブスタックまで見えるプロファイラ。

学び

• まずプロファイリングすれば、小さなコード変更で大きな効果を得られる。
• Python API を維持したまま Rust モジュールだけを置き換えても、実運用サービスにすぐ適用できる。
• Rust は「性能が重要な領域」に薄く導入するだけでも十分に効果的。

[allinux]

C/C++ で Python 拡張を作るのは生産性が低すぎますが、PyO3 はまず maturin や cargo があるのでとても便利です。
また、Python モジュールではクロスコンパイルも必須ですが、Rust はクロスコンパイルもしやすいです。

[lamanus]

maturin... 苦痛...

[aer0700]

できるだけNumPyのベクトル化で粘って、だめならGPUを挿してCuPyやtorchに切り替えて、それでもだめならCythonでネイティブを書く、という感じですが……ネイティブはよほどでなければやらないほうがいい気がします。大変です。