AI Engine を用いた修正コレスキー分解のマルチコア実装

概要

本記事では、Xilinx (現在は AMD に統合) 開発の AI Engine (AIE) という高い演算性能をもつ演算プロセッサを用いた修正コレスキー分解のマルチコア並列処理の実装、評価キットによるシミュレーションによる計算にかかるサイクル数の測定、ライブラリ実装の結果との比較について説明を行っています。 AIE は、 ACAP (Adaptive Compute Acceleration Platform) と呼ばれるプラットフォームに搭載されているプロセッサの 1 つです。 ACAP は、CPU 等を用いたスカラー演算、 DSP や GPU 等を用いたベクトル演算、 FPGA 等を用いた特定の処理に特化した演算の 3 つを統合的に扱い高速な処理を行うことを目的として開発されたものであり、その実現のために高い演算性能をもつ AIE が搭載されています。

現時点で実用化の例があまり多くない AIE について、アーキテクチャに関する基本的な説明や、AIE を用いた処理を行う際に考慮すべき点など、AIE を用いた開発に関する情報を提供することを目的としています。

背景

バンドル調整という、画像からの3次元復元において3次元点群とカメラ姿勢を推定する手法があり、この手法は自動運転技術の実現の際に用いられている手法の 1 つです。この手法は問題の規模が大きい場合には演算量が多くなり、処理時間も長くなってしまうという性質があります。バンドル調整についての詳細や、行う処理とそれにかかる時間についての詳細は別記事「CUDAによるバンドル調整の高速化とソースコード公開」をご覧ください。

自動運転技術はリアルタイム性が求められるものであるため、バンドル調整の処理を高速化しレイテンシを小さくすることが重要となっており、ハードウェア実装を行うことでこの要件を満たすことができると考えられます。 ハードウェア実装の手法としては FPGA を用いたものがありますが、より演算性能の高いプロセッサを用いて実装を行い効率を上げることもできます。 そこで今回は、演算性能の高い AIE を用いて、バンドル調整においてボトルネックとなっているコレスキー分解の実装を行いました。

修正コレスキー分解とアルゴリズム

修正コレスキー分解は、正定値対称行列 \(A\) を下三角行列 \(L\) と対角行列 \(D\) を用いて

\(A = LDL^\top\)

と分解することをいいます。 ここで、\(L\) の対角成分は \(1\) と定められているため、ここを無視して対角成分が \(D\)、非対角下三角成分が \(L\) である行列 \(C\) を計算結果として表現することができます。 \(C\)は一要素ずつ計算することができ、各要素の計算時には自身より行/列番号の小さい \(C\) の要素の一部を使用します。したがって、下図のように row-major order で \(C\) の要素を計算することによってそれぞれの要素の計算時に必要な要素が揃っている状態となります。

具体的な計算手順は以下のようになります。 非対角成分 \(C_{ij}\) の計算に必要な \(C\) の要素は

• 第 \(i\) 行の \(j-1\) 列目までの要素
• 第 \(j\) 行
• \(j\) 番目までの対角成分

であり、以下のように計算されます。

buf = A_{ij}
for k = 0, ... , j-1:
    buf -= C_{ik} * C_{kk} * C_{jk}
C_{ij} = buf / C_{jj}

対角成分 \(C_{ii}\) の計算に必要な \(C\) の要素は

• 第 \(i\) 行の \(i-1\) 列目までの要素
• \(i-1\) 番目までの対角成分

であり、以下のように計算されます。

buf = A_{ii}
for k = 0, ... , i-1:
    buf -= C_{ik} * C_{kk} * C_{ik}
C_{ii} = buf

非対角成分、対角成分について、計算に必要な要素を図示したのが下図となります。

マルチコアで修正コレスキー分解を実装するにあたり、各要素の計算時に必要となるデータについて、計算方法や図を見ながら考えておきます。

• 計算結果 \(C\) の対角成分は常に計算に必要である。
• 計算しようとしている要素と同じ行の要素も必要となるが、これは直前で計算している要素なので一時的な保持で問題ない。
• 左の図の青色の部分、すなわち \(C_{ij}\) の計算時に必要となる \(j\) 行目の要素というのは長期的に保持しておく必要がある。

これらの観察は、マルチコアで修正コレスキー分解のアルゴリズムを実装する際のデータの保持の仕方を考える際に重要となります。

AI Engine と Versal AI Core シリーズ

本節では AI Engine (AIE) のアーキテクチャやそれを搭載しているプロセッサについて説明します。AIE については「実践的！FPGA開発セミナー vol.7」にてそのアーキテクチャや開発例について説明されていますので、適宜こちらもご参照ください。

AIE は、Xilinx (現在は AMD に統合) によって開発された、高い演算性能を持つ演算プロセッサです。 Versalシリーズ に内蔵されているプロセッサであり、1つのチップに数百の AIE コアが内蔵されています。

AIE のアーキテクチャは下図のように複数の AIE タイルを2次元状に配列した構造となっています。

引用: https://japan.xilinx.com/products/technology/ai-engine.html

各エンジンは VLIW SIMD プロセッサを内蔵しており、高性能なベクトル演算を提供しています。 また命令プログラムやローカルデータを格納するための専用のメモリを保持し、各コア間でデータのやり取りを行うためのインターフェース (Interconnect) も提供しています。

データ送受信の手法

各コアにおけるデータの送受信についてもう少し詳しく説明します。 コア間でデータをやり取りする方法は大きく分けて buffer (ver. 2022.2 までは window という名称) と stream の 2 種類があります。

buffer はその名の通りバッファを確保してデータを送受信する方法です。 上の図では赤矢印の Memory Interface が該当し、隣接コア間でしか使えず、バッファの読み書きの際にロック機構を用いる分のオーバヘッドがありますが、データ送受信のスループットは stream に比べて高いです。

stream は各コアにある入力ポートと出力ポートを用いてデータを送受信する方法です。 上の図では黒矢印の Stream Interface、青矢印の Cascade Interface が対応します。 黒い四角で表された AIE Interconnect を通して、隣接していないコア間でもデータを送受信することが可能です。

buffer と stream の性能比較を下表に示します。
引用: https://docs.amd.com/r/2023.1-English/ug1079-ai-engine-kernel-coding/Buffer-vs.-Stream-in-Data-Communication

今回実装したアルゴリズムは計算結果を少しずつ送受信していくやり方であり、buffer を用いて実装しようとすると、リードロック、ライトロックの制御頻度が高くなってしまうために、却って効率の低下を招きました。 したがって、今回は stream を用いてデータの送受信を行うことにしました。

実験に用いたプロセッサおよび評価キット

Versal シリーズはその用途や性能によっていくつかのグループに分けられます。 Versal AI Core シリーズ はVersalの中でも高い演算性能を有するプロセッサであり、今回はこのプロセッサの1つである VC1902 を有する評価キット Versal AI Core Series VCK190 Evaluation Kit を用いて修正コレスキー分解の実装および性能測定を行います。

AI Engine を用いた修正コレスキー分解の自前実装

以下の実装で性能測定を行いました。

1. シングルコア・スカラー演算
2. シングルコア・SIMD演算
3. マルチコア・スカラー演算
4. マルチコア・SIMD演算

シングルコア・スカラー演算

前述のアルゴリズムを素直に実装したものです。ただし、時間のかかる除算の回数を減らすために、計算結果の対角成分についてその逆数を計算して保持しておき、除算を逆数との乗算としています。 この変更により、除算の回数は、行列サイズを \(N\) として \(N(N-1)/2\) 回から \(N\) 回へと減少します。

シングルコア・SIMD演算

前述の通り AIE コアは SIMD 演算をサポートしているため、これを用いて計算の並列化を図ります。 修正コレスキー分解のアルゴリズムをみると、非対角成分において \(C_{i+1, j}\) の計算には \(C_{ij}\) を必要としないことが分かるので、縦に要素を並べてベクトルとし、このベクトルに対して SIMD 演算を行うことで並列に結果を計算できることが分かります。

各要素を 4 Byte の浮動小数点数とすると、AIE では一度に 8 要素まで計算を行うことができます。 したがって、縦の 8 要素をベクトルとしてまとめて計算を行います。

図の大きな正方形は 8 x 8 ブロックを表し、初めの図と同様に計算する要素や計算に必要となる要素を色付けしています。

対角ブロックの計算方法について補足します。 計算する必要のない部分（非下三角成分）は適切に捨てることで特に支障なく必要な部分を計算することができます。 また、対角成分の計算は、非対角成分の計算から対角成分で割る操作を除いたものです(非対角成分の計算で \(i = j\) とすれば buf から引く操作は対角成分の計算方法と一致します)。 よって、対角・非対角関係なく同様の操作で計算すれば対角成分も問題なく計算することができます。そして得られた対角成分の計算結果を用いて非対角成分を割れば非対角成分の計算も完了します。

マルチコア・スカラー演算

ここからはマルチコアで計算する手法について説明します。

前述のように AIE の各コアは独立にプロセッサや命令・ローカルデータ用メモリをもつので、それぞれが並列に処理行うことが可能です。 保持が必要なデータの選択や計算を行う要素の分担を適切に行うことで演算強度を上げ、効率的な計算を行えることが期待されます。 今回は、以下のようなデータの流れ・計算コアの割り当てで計算を行いました。

データの流れ・保持

計算を行うコアが 4 つ (コア #0, #1, #2, #3 と表記します) だとして説明を行います。

コア #0 が PL (Programmable Logic) から入力を受け取り、最後のコアが PL に出力を行います。 各コアで計算した \(C\) の要素は出力のため最後のコア #3 に送る必要があり、また基本的にほかのコアでの計算でも必要なので、図の緑矢印に沿ってデータを一周させます。 計算結果 \(C\) の計算には以前に計算した要素を保持しておく必要がありますが、今回は行ごとに保持するコアを割り振ります。 すなわち、流れてくる \(C\) のうち決められた行のデータを保持しておくことにします。 具体的には、保持を行うコアは一番上の行から順に #0, #1, #2, #3, #0,… とします。 また、計算結果の対角成分は全てのコアが保持し、ある行について計算を行っているときはその行の計算結果を保持しておくことにします。 このようにすることで、計算に必要な要素が常にコア内に存在するようにしています。

処理を行うコアの割り当て

図は各要素について、\(C\) の計算を行うコアの分担を示したものです。 一番左の列は上から #0, #1, #2, #3, #0, … と要素ごとに計算コアを割り当て、その他の列については左から #1, #2, #3, #0, #1, … と列ごとに計算コアを割り当てます。

一番左の列については計算負荷を分散させる目的で分担させていますが、その他の列については計算結果のデータの保持の持ち方から決定されます。例えば、\(C\) の第 2 列の要素の計算には第 2 行の要素が必要ですが、これはコア #2 が計算結果を保持しているため、このコアが計算を行うようにすれば他のコアから新しくデータを受け取る必要はなく効率的に計算できます。

各要素の計算が終わったら、緑の矢印に沿ってデータを次のコアに転送します。最後のコアは PL への出力も行います。

各コアについて、ある \(C\) の要素を計算する際、そのコアで 1 つ前に計算した \(C\) の要素が同じ行にある場合には、その時に利用した \(C\) の要素を再利用することができます。 よって、新たに必要となる \(C\) の要素が計算されて送られてくる前に計算を前もって行っておくことができます。 これにより、コアが並列で計算を行うようになり、全体の計算にかかる時間が短くなることが期待されます。

マルチコア・ SIMD 演算

上記 2 つの計算方法を組み合わせたのがこの手法です。

データの流れはスカラー演算の場合と同じですが、データの保持・処理を行うコアの割り当てを少し変更します。 SIMD 演算で処理を行う都合上、割り当ての分担は 8×8 ブロック単位がやりやすいためそのようにします。 具体的には、\(C\) の保持は上から 8 行ごとに #0, #1, #2, #3, #0, … とし、処理を行うコアは左から 8 列ごとに #0, #1, #2, #3, #0, … とします。 ここで、マルチコア・スカラー演算で行っていた一番左の列の計算を分担する手法は簡単のため廃止しました。

ライブラリ実装

Vitis統合ソフトウェアプラットフォームでは Vitis Accelerated Libraries というライブラリが提供されており、その中にコレスキー分解を行うライブラリが存在します(実装)。

このライブラリは行列サイズ（行列の行数または列数）に等しい数のコアを用いて計算を行うもので、(PL からの入力) → (コア) → … → (コア) → (PL への出力) のように一方向にデータを流して計算を行います。

このライブラリが行うのはコレスキー分解であり“修正”コレスキー分解ではないため厳密には今回実装したものと等価な計算を行うものではないのですが、アルゴリズムとして大きな違いはなく実行時間に大きな変化はないと考えられるため、こちらに関してもシミュレーションを行いサイクル数の計測を行いました。

結果と考察

アルゴリズム間の比較

自前実装 4 つにライブラリ実装を加えた 5 つの実装について、Xilinx Versal AI Core VCK190 で実行することを想定してシミュレーションを行い、計算にかかったサイクル数を記録しました。 マルチコア処理の自前実装ではコア数を 4 としました。 結果を下表に示します。

まず自前実装について結果を比較します。4つの実装の中で一番サイクル数が少なかったのはマルチコア・ SIMD 実装で、計算を工夫したことの効果が出ていることが読み取れます。 また、シングルコア・マルチコアそれぞれについてスカラー演算・ SIMD 演算のサイクル数を比較すると倍程度の計算速度となっており、SIMD 演算でまとめて計算することの効果が読み取れます。

シングルコア実装は行列サイズが 128 のときの結果が記載されていませんが、これはコア内蔵のメモリに行列が乗りきらないためです。 今回使用した Xilinx Versal AI Core VCK190 には各コアに 32 KB のローカルデータ用のメモリが内蔵されています。 行列の各要素を float (4 Bytes) で表現する場合、 128 x 128 の行列の保持には 4 Bytes * 128 * 128 = 64 KB となり、単一のコアに乗りきりません。 そのため、大きな行列に対して AIE でコレスキー分解の計算を完結させたい場合、必然的にマルチコアでの処理が必要となってきます。

ライブラリ実装の結果と、最速の自前実装であるマルチコア・ SIMD 演算実装の結果を比較します。 サイクル数の観点ではライブラリ実装はとても早く、 行列サイズ 128 の場合ではマルチコア・ SIMD 演算実装に比べて約 18 倍の速度となっています。 ただし、演算に用いるコア数が異なることには留意する必要があります（ライブラリ実装はコア数 128、自前実装はコア数 4）。 サイクル数にコア数をかけた値をのべサイクル数として比較すると、自前実装は 5164496 であるのに対しライブラリ実装は 9015040 であるため、この指標での比較では自前実装に軍配が上がります。ただし、データ待ちストール等により演算コアが動いていない時間もあるため、演算コアが実際に動作しているのべサイクル数を示すものではないことには留意が必要です。

コア数とサイクル数の比較

次に、同じ行列サイズについて、コア数を変化させたときのサイクル数の変化を見ます。

行列サイズ 128 の行列に対して、コア数 4, 6, 8, 16 の 4 種類でマルチコア・ SIMD 演算実装のシミュレーションを行い、サイクル数を測定しました。結果を下表に示します。

コア数を多くするとサイクル数も多くなるという結果となってしまいました。 行列サイズがまだまだ小さく、各要素の計算に必要な演算の数が少ないため、データ転送がボトルネックとなりコア数の増加に伴ってサイクル数が増えてしまっていると考えられます。

コア内蔵メモリの制約

行列サイズを大きくすればコア数を多くすることのメリットが現れそうであるため、行列サイズを大きくしてシミュレーションを回したいところですが、 AIE の各コアに内蔵されているローカルデータ用のメモリのサイズが 32 KB と小さめであることにより、そもそもビルドが行えない、といった状況となってしまいました。

例として、行列サイズ 512、コア数 64 の場合を考えます。各コアは計算結果の行列のうち 8 行分を持つ必要があります。 各要素は 4 Byte であるため、このデータの保持に必要なサイズは 4 Bytes * 8 * 512 = 16KB となります。 他にも持つべきデータはいくつかあり、それらを合算すると大きなデータサイズとなってしまい、ビルドに失敗します。 詳しい内訳は不明ですが、最終的なスタックサイズ（これにローカル変数のサイズが含まれる）は 54 KB 程度になってしまっています。

ERROR: Core 1 0: Estimated stack size 53792 bytes exceeds the allocated stack size 28672

今回用いたプロセッサ VC1902 がもつ AIE の数は 400 であるため、 400 を超える行列サイズの行列のコレスキー分解をしたい場合ライブラリ実装ではそもそも計算が行えません。 自前実装ではコア数を自分で設定することが可能であるため、行列サイズが大きくなっても適切にコア数を設定すれば計算できることを示し自前実装の優位性を示そうとしましたが、スタックサイズの問題によりビルドができないという問題が発生してしましました。

おわりに

AIE を用いた修正コレスキー分解のマルチコア実装について、データの保持方法や転送方法、シミュレーション結果をご紹介しました。 AIE はそれぞれが専用の命令用・データ用メモリをもち、 SIMD 演算など高速な演算をサポートしているため、上手く使うことができれば効率的な処理を実現することができます。 実際、行列サイズ 64 の行列のコレスキー分解について、マルチコア・ SIMD 演算実装はシングルコア・スカラー演算実装に比べて4倍程度早く計算することができています。 一方で、効率的な処理のため、あるいはデータメモリのサイズの制約からデータの保持方法や転送方法についてしっかりと考えてから実装する必要があり、その点に関しては難しいと感じました。

なお今回実装は行っていませんが、データ転送に関してよりシンプルな構造となっているコレスキー分解のシストリックアレイアルゴリズムが存在しており、このアルゴリズムを実現することでより高速なコレスキー分解の計算が実現できる可能性があると感じました。

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。