「GPU上の行列乗算」.

「GPU上での行列乗算」

CUDAで最先端の行列乗算性能を実現する方法

このブログは、GPU上で行列乗算がどのように機能するかについての自分の無知に気付いた瞬間から生まれました。私は多くの機械学習プロジェクトを行ってきたので、機械学習における最も重要な操作がどのように機能するかを理解する必要があると感じました：この「テンソルコア」とは何ですか？なぜ「データの移動がボトルネック」と言われるのですか？GPUは実際にどれくらい速いのですか？

これらの質問に答えるために、私はPyTorchのバブルを抜け出し、CUDAの深淵に冒険することにしました。私はこのブログを執筆し、私が学んだすべてを記録することにしました。そして、この記事を読んでいる人たちは、私が行ったようにCUDAのドキュメントやコードを探す苦労をしなくてすみますようにと願っています。

この旅で学んだことがあるとすれば、並列行列乗算は難しいということです。効率的な行列乗算は、使用している具体的なハードウェアと解決しようとしている問題のサイズに大きく依存します。ワンサイズフィットオールの解決策はありません。

長話はそれだけにして、さあ始めましょう！

• 「コードの機能と読みやすさを強化するために、パンダのデータフレームをカスタムメソッドで拡張する方法」
• 「OSMネットワークでの移動時間によって重み付けされたダイクストラのアルゴリズム」
• 「初心者のためのバイトペアエンコーディング」

GPUアーキテクチャのおさらい

NVIDIAのGPUがどのように機能するかを思い出しましょう。GPUは多くのスレッドを実行することにより並列性を実現します。各スレッドは単一のCUDAコアで実行されますが、ある時点では、アクティブなスレッドのサブセットが存在するため、CUDAコアよりも多くのスレッドが存在することがあります。アクティブでないスレッドであっても、各スレッドは独自の一連のレジスタを持っています。

32個のスレッドのグループをワープと呼びます。ワープ内のすべてのスレッドは一緒に実行される必要があります（または一緒に非アクティブにする必要があります）。ほとんどの場合、アクティブなワープよりも非アクティブなワープの数が多くあり、ワープスケジューラは与えられた時間に実行するワープを選択する役割を担っています。これにより、ワープがデータを待っている間に他のワープを実行することで、GPUはメモリアクセスのレイテンシを隠すことができます。

ワープのグループをスレッドブロックと呼びます。スレッドブロック内のすべてのワープは同じStreaming Multiprocessor（SM）で実行されます。各スレッドブロックには、スレッドブロック内のすべてのスレッドからアクセスできる共有メモリがあります。

注意：新しいアーキテクチャ

Voltaアーキテクチャ以降、各スレッドにはそれぞれ独自のプログラムカウンターや呼び出しスタックなどがあります。つまり、ワープ内の各スレッドは同時に異なる命令を実行できます。

Voltaアーキテクチャでは、特定のサイズの行列乗算を解決するためのテンソルコアも導入されました。各アクティブなワープは1つのテンソルコアにアクセスできます。

最新のHopperアーキテクチャでは、スレッドブロッククラスタという概念が導入されています。これはスレッドブロックのグループを表し、スレッドブロックのスケジューリングを細かく制御し、同じクラスタ内の他のスレッドブロックが共有メモリにアクセスできるようにします。

行列乗算の並列化

次の計算を行いたいとしましょう：

この場合、問題のサイズは(M, N, K)となります。この操作を並列化するために、AとBをより小さい行列に分割し、それぞれを個別に行列乗算し、結果を連結してCを形成することができます。

具体的には、Aを行方向に（つまり、MをM’サイズのチャンクに分割）分割し、Bを列方向に（つまり、NをN’サイズのチャンクに分割）分割することができます：

各サブ行列Cは互いに独立しているため、各サブ行列の計算を容易に並列化することができます。

実際には、Kが大きすぎて直接メモリにロードして計算することはできない場合、一般的な実装ではKをK’のサイズのチャンクに分割し、各チャンクに対して反復し、部分結果を蓄積（合計）します。これはシリアルK削減として知られています（後述のパラレルK削減と対比されます）。数学的には以下のようになります：

注：パディング

問題のサイズがパーティションのサイズで割り切れない場合、パディングを追加する必要があります。これは通常、パーティションされた入力（𝐴ᵢ,ₖと𝐵ₖ,ⱼ）を下位メモリにロードする際に暗黙的に行われます。その際、ロードされるパーティション（𝐴ᵢ,ₖの場合はM’×K’、𝐵ₖ,ⱼの場合はK’×N’）が常に「完全」であるように、必要な部分には常にゼロを追加します。結果をグローバルメモリに書き戻す際には、範囲外エラーを回避するために特別な注意が必要です。

大まかに言うと、GPU上の行列積を並列化するために、3つのネストされたパーティションが発生します：1. 最初のパーティションはスレッドブロックレベルで発生します。各スレッドブロックはCᵢ,ⱼ = Aᵢ Bⱼの計算を担当します。2. 2番目のパーティションはワープレベルで発生します。スレッドブロックレベルの問題Cᵢ,ⱼはさらにパーティションされ、各ワープはCᵢ,ⱼ⁽ᵐⁿ⁾ = Aᵢ⁽ᵐ⁾ Bⱼ⁽ⁿ⁾の計算を担当します。3. 3番目のパーティションは命令レベルで発生します。一部の命令は特定のサイズの入力を予期しています。たとえば、第2世代のTensor Coresはfp16のサイズ（16、8、8）で動作し、CUDAコアの直接の実装では単純にサイズ（1、1、1）で動作します。したがって、ワープレベルの問題はさらにパーティションされ、各チャンクは命令に適したサイズとなるようになります：Cᵢ,ⱼ⁽ᵐⁿ⁾⁽ᵃᵇ⁾ = Aᵢ⁽ᵐ⁾⁽ᵃ⁾ Bⱼ⁽ⁿ⁾⁽ᵇ⁾.

次のセクションで見るように、3つのパーティションレベルが必要な理由があります。

データの冗長性

行列の乗算は、計算を実行するたびにデータをグローバルメモリからレジスタに再取得すると、メモリに制約される可能性があります。重要な観察点は、多くのサブ入力AᵢとBⱼが異なるサブ行列の乗算に再利用されることです。たとえば、AᵢはCᵢ,₁、Cᵢ,₂などに必要であり、BⱼはC₁,ⱼ、C₂,ⱼなどに必要です。可能な限り冗長なデータ移動を最小化し、ロードされたデータをできるだけ再利用することで、最大のスループットを得ることができます。

CUDAでは、3種類のユーザーアクセス可能なメモリがあります：

次に、各メモリタイプの利用方法の概要を示します：

1. 各スレッドブロックは、まず必要な入力をグローバルメモリから共有メモリへロードします。その後のアクセスは、共有メモリによって提供されるため、遅いグローバルメモリではなくなります。
2. 各スレッドブロック内では、まず各ワープは必要な入力を共有メモリからレジスタにロードします。その後のアクセスは、高速なレジスタによって直接提供されます。

詳細について掘り下げる

スレッドブロックレベル

スレッドブロックレベルでは、問題はサイズ(M’、N’、K’)のサブ問題に分割されます。したがって、各スレッドブロックはCの断片を計算する責任があります。:

冗長なデータの移動は、サブ入力Aᵢ,ₖとBₖ,ⱼを共有メモリにロードすることで最小限に抑えられます。Aᵢ,ₖ Bₖ,ⱼの計算が完了したら、次のチャンク(Aᵢ,ₖ₊₁とBₖ₊₁,ⱼ)がロードされます。

ワープレベル

ワープレベルでは、サブ問題はさらにサブサブ問題のサイズ(M”、N”、K”)に分割されます。したがって、各ワープはCᵢ,ⱼの断片を計算する責任があります。Cᵢ,ⱼ⁽ᵐ ⁿ⁾として示されます:

冗長なデータの移動は、サブサブ入力Aᵢ,ₖ⁽ᵐ ˡ⁾とBₖ,ⱼ⁽ˡ ⁿ⁾をレジスタにロードすることで最小限に抑えられます。ワープ内でAᵢ,ₖ⁽ᵐ ˡ⁾とBₖ,ⱼ⁽ˡ ⁿ⁾へのアクセスは、高速なレジスタで処理されます。

注意：レジスタ間でデータを分配する

レジスタはスレッドレベルのみです。つまり、レジスタの入力はワープ内の他のスレッドからアクセスすることはできません。各スレッドのレジスタにAᵢ,ₖ⁽ᵐ ˡ⁾とBₖ,ⱼ⁽ˡ ⁿ⁾がどのように分割されるかは、使用される具体的な命令によって異なります。NVIDIAのWarp Level Matrix Multiply-Accumulate Instructionsには、各命令の詳細な説明があります。

テンソルコアレベル

実際の行列乗算を行うために、GPU上のテンソルコアを使用します。私のGPU（RTX 2060）は、第2世代のテンソルコアを備えており、サイズ(M”’、N”’、K”’) = (16, 8, 8)のfp16の問題を解決するために特化しています。したがって、Cᵢ,ⱼ⁽ᵐ ⁿ⁾をさらに分割して命令で期待されるサイズに合わせます:

ここでは、すべての入力が既にレジスタに格納されており、データ移動のオーバーヘッドは最小限です。

注意：テンソルコア

テンソルコアの操作はワープレベルの命令であり、ワープ内のすべてのスレッドが同時にテンソルコア命令を実行し、データを共同で準備します。

パーティションのサイズを選ぶ

したがって、データの移動を最小限に抑えたい場合、共有メモリとレジスタを最大限に利用するようなパーティションのサイズを選ぶべきですか？それは完全にはありません。

スレッドブロックのパーティションサイズ

漸近的に、問題のサイズが増加するにつれて、共有メモリとレジスタを最大限に利用したいと考えるでしょう。しかし、小さな問題サイズでは、2つの問題に直面する可能性があります:

1. 大きなパーティションサイズを持つと、スレッドブロックが少なくなる可能性があります。各スレッドブロックは1つのSMでのみ実行できるため、すべてのSMを利用できない可能性があります。
2. パーティションサイズで割り切れない問題サイズの場合、入力にさらにパディングを追加する必要があります。これにより、意味のある入力で行われる計算が少なくなり、効率が低下します。

典型的な実装では、パーティションサイズは(M’, N’, K’) = (128, 256, 32)のようなものが使われます。

ワープのパーティションサイズ

一般的に、大きなワープのパーティションサイズは冗長なデータの移動が少なくなりますが、ワープが少なくなるというコストが発生します。ワープが少なすぎると、メモリアクセスのレイテンシを隠せなくなります（現在のワープがデータを待っている間に他のワープをスケジュールする余裕がなくなる可能性があるため）。

典型的な実装では、パーティションサイズは(M”, N”, K”) = (64, 64, 32)のようなものが使われます。

命令のパーティションサイズ

これはGPUがサポートする命令によって完全に決まります。私のRTX 2060では、fp16 Tensor Core行列乗算（fp16の累積を使用）のptx命令はmma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16で、入力のサイズは(16, 8, 8)を期待しています。

さらなる最適化

上記の技術は、問題のサイズが大きい場合にはGPUの理論的なピーク性能に近づけます。しかし、問題のサイズが小さい場合には効率的ではありません。行列乗算のパフォーマンスをさらに向上させるためには、parallel-K reductionとsoftware pipeliningという2つの一般的なテクニックがあります。

Parallel-K reduction

MとNが小さい場合、スレッドブロックがほんの数個しかないかもしれません。例えば、(M’, N’) = (128, 256)で、元の問題のサイズがM ≤ 128およびN ≤ 256の場合、スレッドブロックは1つしかなく、したがってGPUの計算能力の一部しか利用していません（例えば、私のRTX 2060には30のSMがあるため、利用率を最大化するためには少なくとも30のスレッドブロックが必要です）。

MとNが小さい場合でも、Kが大きい場合（MとNが小さい場合でも）、parallel-K reductionを利用してより多くの並列性を活用することができます。直列-K reductionでは、各スレッドブロックは次の合計を繰り返し処理します：

そして、中間結果をCᵢ,ⱼに蓄積します。一方、並列-K reductionでは、各スレッドブロックは合計の要素のうち一つだけを計算するように割り当てます（つまり、Aᵢ,ₖ Bₖ,ⱼ）。これにより、スレッドブロックの数をK/K’の比率で増やすことができ、より多くのSMを利用します。

ただし、今度は各スレッドブロックからの結果を格納するためにより多くのメモリを割り当て、部分的な結果に対して最終的なreductionを実行するために2度目のカーネルを呼び出す必要があります。

ソフトウェアパイプライン処理

通常、CUDAはメモリアクセスのレイテンシを隠すために、ワープがデータを待っている間に他のワープを実行するようにスケジュールします。これには、レイテンシをマスキングするために十分な数のワープが必要です。

しかし、GEMMを行う際のワープの数は通常比較的少ないです。これは、「スレッドブロックごとの利用可能なレジスタの数をワープごとに必要なレジスタ数で割った値が制限されるため」という理由です。行列乗算では、データ再利用を最大化するためにワープごとに多くのレジスタが使用されます。その結果、レイテンシをマスキングするのに十分な数のワープがないかもしれません。

「アキュムレータの要素は通常、1つのスレッドの総レジスタ予算の半分以上を占める」- CUTLASS Docs

この問題を軽減するために、software pipeliningを使用することができます。本質的には、特殊な命令を使用してループの次のイテレーションの入力を非同期に事前にロードすることができます。入力がロードされている間、現在のイテレーションで演算を続けることができます。次の図によって要約されています：

これは、GPUが現代のCPUと同様のメモリアクセスと算術演算をパイプライン化できるために可能になっています。データ間の依存関係がない限り、これは命令レベルの並列処理として知られています。

行列の乗算のアクション

これらの概念が実際の実装でどのように結びつくかを見たい場合は、CUDAを使用してからスクラッチでMNISTのトレーニングを行った私の実装をチェックしてください。そこでは、128のユニットサイズの複数層パーセプトロンをMNISTでトレーニングし、最適化されたPyTorchに対して6倍の高速化を達成しました。

GitHub – andylolu2/cuda-mnist

GitHubでアカウントを作成してandylolu2/cuda-mnistの開発に貢献してください。

github.com

参考文献

1. CUTLASSドキュメント2. CUDAドキュメント3. CUTLASSの例

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!