AMD、RDNA 3/4向けFP16 GEMMを8段階で最適化する手法を解説

AMDは2026年10月16日、ROCm Blogsで、RDNA 3/RDNA 3.5/RDNA 4 GPU向けにWMMA命令でFP16 GEMMカーネルを8段階で最適化する解説記事を公開しました。最終カーネルは、アーキテクチャとGEMM形状によってROCmのチューニング済みライブラリhipBLASLtより数%速い場合も遅い場合もあります。著者はAdel Johar、Jan Stephan、Istvan Kissの3氏で、2部構成の第1部にあたります。

AMD ROCmは、AMDが提供するGPU計算向けのオープンなソフトウェア群で、実行基盤、コンパイラー、数値計算ライブラリなどを含みます。

何が新しいか

記事は本番ライブラリの提供ではなく、各ステップで解決する問題、コード変更、そして時間・レジスター・LDSへの影響を示すチュートリアルです。著者らは、ある最適化が一方のアーキテクチャでは効き、もう一方では逆効果になるケースを個別に示しています。第2部では、最終カーネルをRDNA 4向けにFP8へ移植する予定です。

8段階の内容は次のとおりです。

  1. 1 waveで小タイルを処理する最小カーネル
  2. タイル拡大、128ビットメモリコピー、ダブルバッファリング、staged commitの4手法を同時適用
  3. LDSバンク競合の回避(パディングまたはレイアウトの置換)
  4. グローバルロードにバッファリソースディスクリプターを使用(範囲外アドレスはゼロを返すため、タイルサイズが行列寸法を割り切れなくてもよい)
  5. Kスライスの深化
  6. キャッシュ再利用のためのブロックスケジューリング
  7. プリフェッチとWMMAの交互実行
  8. ダブルバッファリングの再検証(2バッファか、1バッファで常駐ワークグループを増やすか)

RDNA 3とRDNA 4のWMMAの違い

WMMAはwave内の32レーン全体で16×16×16の行列積和を実行する命令です。RDNA 3では2つの16レーン半分が同じオペランドを複製して保持し、各レーンが16個すべてのK値を持ちます。RDNA 4ではレーン0〜15がK値0〜7、16〜31がK値8〜15を持ちます。

項目 RDNA 3 (gfx11) RDNA 4 (gfx12)
フラグメント/アキュムレーター幅 16 FP16 8 FP16
レーンが読むK値 16個すべて(複製) レーン半分ごとに0-7または8-15
使用するアキュムレーター要素 1つおき すべて
命令 __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(OPSEL引数付き) __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12

このためRDNA 3のフラグメントはLDSからの読み込みに128ビットロードを2回要し、同じ出力タイルでアキュムレーターのレジスターを2倍使います。コードは両者の違いをアーキテクチャごとのpolicy構造体に隠し、コンパイル時に切り替えます。

ベンチマーク条件

項目 内容
GPU Radeon PRO W7800 (gfx1100, RDNA 3)、Radeon RX 9070 XT (gfx1201, RDNA 4)、Radeon 8060S (gfx1151, RDNA 3.5)
標準形状 M=N=K=8192
データ型・レイアウト A: FP16 row-major、B: FP16 column-major、C: column-major(leading dimension M、hipBLASLtのTN)
ソフトウェア ROCm 10
比較対象 hipBLASLt

Radeon 8060SはRDNA 3と同じWMMA命令を使うため同一バイナリで動きますが、専用GDDRではなくシステムRAM経由でメモリにアクセスするため、最適化後も絶対スループットは低くなります。正確性検証では、AとBの要素を-1または+1で埋め、部分和がFP16で正確に表現できる範囲に収まるようにしたうえで、hipBLASLtと許容誤差なしにビット単位で比較します。

追試に必要な環境

  • hipBLASLt開発パッケージを含むROCmのインストール
  • RDNA 3、RDNA 3.5、またはRDNA 4 GPU
  • HIPとGEMMタイル分割(ブロック、wave、Kスライスへの分割)の基礎知識

記事中のコードの多くは構造を示す疑似コードで、WMMA組み込み関数、バッファリソースディスクリプター、LDS置換のアドレス計算はHIPで示しています。全ステップの完全実装とベンチマークドライバーは付属リポジトリにあります。

この記事の固有名詞

コメント (0)

まだコメントはありません。

コメントを書く

2000文字以内