Technology

GPU別カーネルを必要な分だけ選べるROCmの共通配布

この記事のポイント

  1. 実現したこと

    ROCmのPython配布では、単一、複数、または全対応GPU向けのデバイスコードを選んで導入できる。

  2. 実現の仕組み

    複数のGPUアーキテクチャをまとめてビルドし、パッケージ化の段階でGPU固有カーネルを共通のホストコードから分離する。

  3. 得られた結果

    GPU固有カーネルの取得範囲を単一GPUまで絞り、従来よりダウンロードを小さくできる。

  4. 従来との違い

    GPUファミリー別の配布先を選ぶ方式から、共通インデックスで対象GPUを指定する方式へ変わった。

共通ホストコードから必要なGPU向けカーネルだけを分離して配布する構成を抽象的に描いたイラスト
AI生成画像

ROCmのマルチアーキテクチャ配布は、GPU固有カーネルを共通のホストコードから分離する。利用者は共通インデックスから、必要なGPU向けコードを選んで導入できる。

一体型のリリース工程をコンポーネント単位へ分割

AMDは、GPU計算ソフトウェア基盤ROCmのビルドとリリースに、ROCm 7.14以降はオープンソースのビルドシステムTheRockを使用している。この移行により、従来の一体的なリリース工程は、コンポーネントを個別に扱えるモジュール型のワークフローへ置き換わった。

TheRockは、GPUプログラミング環境HIPとROCmの各コンポーネントを統合する最上位のビルドシステムである。分割したコンポーネントを、単一の製品としてビルド、テスト、リリースする役割を担う。

複数のGPU向けコードをまとめてビルド

マルチアーキテクチャリリースでは、対象となるGPUアーキテクチャをまとめてビルドする。GPUファミリーの部分集合ごとに成果物を作っていた従来方式とは、ビルド時に扱う対象の範囲が異なる。

まとめて生成した成果物は、パッケージ化の段階で共通部分とデバイス別部分に分かれる。GPU固有コードをカーネルパックとしてアーキテクチャ非依存のホストコードから切り離し、それぞれを別のパッケージとして配布できる構成にする。

device指定で導入するGPUを選択

マルチアーキテクチャ版のROCm Python配布は、OS別のコア、アーキテクチャ非依存ライブラリ、GPUターゲット別デバイスコード、OS別の開発ツールで構成される。コンパイラやユーティリティーを含む共通部分と、特定GPU向けのコードがパッケージ単位で分離されている。

利用者はPythonパッケージのdevice指定を使い、単一GPU、複数GPU、または全対応GPU向けのデバイスコードを選択する。複数種類のGPUを扱う環境では、必要な複数の指定を一つのインストール要求に含められる。

GPUファミリー別の配布先を共通インデックスへ統合

配布経路は、GPUファミリーごとに異なるインデックスURLを選ぶ方式から、共通インデックスで対象GPUを指定する方式へ変わった。利用者は配布先をGPUファミリー別に探す代わりに、共通の配布先へdevice指定を加える。

GPU固有カーネルを独立させたことで、取得範囲は単一GPUまで絞り込める。ファミリー全体または全対象のカーネルを含めていた従来方式と比べ、利用しないGPU向けコードを外し、ダウンロードを小さくできる。

依存関係の解決と実機動作を分けて確認

ナイトリー版の集約インデックスには、ROCm Coreに加えて、機械学習フレームワークのPyTorchとJAXのパッケージも収録される。共通のインデックスから取得することで、ROCmと各フレームワークにまたがる依存関係を解決できる。

ただし、パッケージを導入できることと、対象GPUで実行できることは別の段階である。GPU向けパッケージが生成、公開され、インストールに成功しても、動作確認を経ていないGPUではデバイス列挙、カーネル起動、ライブラリ読み込みが機能するとは限らない。