rapidsai/rmm

NVIDIA RMM is a library for allocating and managing GPU memory in C++ and Python.

NVIDIA RMM (RAPIDS Memory Manager)

概要 – RMM は、GPU メモリを割り当てるための共通かつ柔軟なインターフェースを提供する C++/Python ライブラリです。開発者は、さまざまな割り当て戦略(plain cudaMalloc、プール型アロケータ、マネージドメモリなど)をプラグインし、C++ コードと Python の両方から rmm パッケージを通じてそれらを使用できます。

重要性 – 高性能な GPU ワークロードでは、cudaMalloc/cudaFree を繰り返し呼び出すコストが実行時間を支配することがあります。RMM のプールおよびビニングアロケータはそのオーバーヘッドを劇的に削減し、ストリーム順序付き (stream-ordered) 割り当てをサポートしているため、追加の同期なしで同じ CUDA ストリーム上でメモリを再利用できます。

主な機能

  • 複数のメモリリソース実装 – cuda_memory_resource、managed_memory_resource、pool_memory_resource、fixed_size_memory_resource、binning_memory_resource。
  • ストリーム順序付き割り当て – 特定の CUDA ストリーム上で割り当てと解放を行い、高速な再利用を可能にします。
  • デバイスごとのデフォルトリソース – GPU ごとにグローバルなアロケータを設定したり、割り当てごとにオーバーライドしたりできます。
  • C++ API – CUDA の割り当て関数の薄いラッパーに加え、rmm::device_buffer、rmm::device_uvector、rmm::device_vector などの RAII コンテナを提供します。
  • Python バインディング – pip/conda 経由でインストール可能(librmm-cuXX および rmm-cuXX)。
  • ABI バージョニング – 同じプロセス内で複数のバージョンを共存させることができます。

典型的な使用例

rmm::mr::cuda_memory_resource cuda_mr;
auto pool = rmm::mr::pool_memory_resource{cuda_mr,
                                          rmm::percent_of_free_device_memory(50)};
auto prev = rmm::mr::set_current_device_resource(pool);
// Subsequent allocations use the pool automatically
rmm::device_buffer buf(1024, cuda::stream_ref{cudaStreamDefault}, pool);

Python では、import rmm; rmm.reinitialize(pool_allocator=True) を実行してプールアロケータに切り替えることができます。

入手方法 – 事前ビルドされた wheel をインストールするか(CUDA 13 の場合は pip install librmm-cu13 rmm-cu13 など)、CMake と提供されている build.sh を使用してソースからビルドします。リポジトリには、ドキュメント、チュートリアル、CMake 統合ヘルパーの完全なセットが含まれています。

採用事例 – NVIDIA の RAPIDS エコシステムの一部であり、cuDF、cuML、その他の GPU 加速データサイエンスライブラリなどの下流プロジェクトが、効率的なメモリ処理のために RMM に依存しています。

関連

  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト