MN-Core のデータ転送系命令の可視化に関する取り組み

PFN エンジニアの宮本です。今回は私が入社後に行っている MN-Core 2 のデータ転送系命令を可視化する取り組みについて紹介します。

MN-Core シリーズについて

MN-Core シリーズの第 2 世代である MN-Core 2 では、第 1 世代と同様に階層的なアーキテクチャを採用しており、メモリ構造も同様に階層化されています。

画像は第 1 世代のものですが、第 2 世代でも同様に、

• 2 つごとにグループ分けされた 8 つの L2B (Level 2 Block)
• それぞれの L2B の子である 8 つの L1B (Level 1 Block)
• それぞれの L1B の子である 16 個の MAB (Matrix Arithmetic Block)
• それぞれの MAB の子である 4 個の PE (Processing Element)

が存在し、L2B のグループごとに PDM と呼ばれる SRAM とメインメモリである DRAM が、L2B と L1B にはそれぞれ L2BM と L1BM と呼ばれる SRAM が付属しています。

メモリに関する操作の重要性

また、キャッシュメモリなどは存在せず、全てのデータの移動を事前にソフトウェア側で計画しておく必要があります。そのため、計算を行うテンソルのメモリ配置レイアウトや、レイアウトの変更に伴うデータの移動をどうするかが、MN-Core の性能を最大限に引き出すためにとても重要になってきます。

MN-Core 2 には MN-Core よりも多くのメモリユニットが搭載されており、データ転送に関する命令もより多様なものになっています。既に公開されている MN-Core 2 のソフトウェア開発者マニュアル では、データの転送に関する命令（放送・分配・縮約命令を含む）の説明が大部分を占めており、それぞれの命令の挙動も複雑なものが多いため、初学者にとってはこれらの命令の挙動の把握が大きなハードルの一つになっています。

また、ソフトウェアスタックの成熟とともに抽象化が進んでおり、転送命令の詳細を知ることなくコンパイラ開発ができるようになっています。結果として、ある程度 MN-Core 開発経験を積んだエンジニアであっても、「どの命令でどのデータがどこに移るか」というのを完全に把握しているわけではなく、詳細を確認するためにデータ移動に関して図を用いて記された追加の資料が重宝されている、という状況にありました。

そこで、実際に初学者である自分がそれぞれのデータ転送系命令について学ぶと同時に、どのデータがどこに送られるかをアニメーションとして可視化した資料を作成できれば、社内のエンジニア・リサーチャーにとって使いやすい資料となるだけでなく、これから MN-Core 2 のアーキテクチャを学ぶ人々にとっての学習ハードルを下げる効果があると考え、MN-Core 2 に含まれるデータ転送系の命令を網羅したアニメーション資料を作成することにしました。

アニメーション資料の作成

実際の MN-Core 2 にはボードあたり 4096 個という多数の PE が搭載されており、命令は SIMD で動くのでボード全体でデータの移動が発生するのですが、共通の動きをする部分をまとめて表示したり、データの移動がない部分の表示を省略することで、各命令の動作をコンパクトに表現することができます。

アニメーションの実装にあたっては、使いやすさと汎用性・拡張性の観点から、まずブラウザを用いた汎用的なアニメーションプレイヤーを作成し、次にそのプレイヤーで再生可能な形式の動画ファイルを出力するプログラムを別途作成しました。作成した動画ファイルは事前にアップロードしておくことでプレイヤーがまとめて読み込み、再生できるようにします。これにより、社内で共有されているプレイヤーの URL にアクセスすると、複数の命令をスムーズに切り替えながら挙動を確認することができ、かつ命令の種類が増えた際にも対応するプログラムを書いて得られた動画ファイルを追加するだけで新しいアニメーションを追加することができます。

また、アニメーションは通常の動画ファイルと違い、フレームごとにベクターグラフィックス形式で表現されているため、動画の作成プログラムで全体の解像度を気にする必要がなく、プレイヤーからも品質を損なうことなく拡大縮小を行うことができます。

作成したアニメーションの例をいくつか紹介します。

DRAM → L2BM 単独個別転送

この命令では、特定のグループの DRAM から、特定の L2BM へとデータを転送します。

MN-Core 2 の命令では、ホストマシンから供給される命令の帯域の都合上、4 サイクル分の動作をまとめて一つの命令として送信します（4 サイクル分の動作を 1 ステップと呼びます）。

データ転送に関するアドレスの最小単位は長語 (long word) と呼ばれるものであり、この命令では、1 サイクルごとに連続する領域の 16 長語のデータが、合計で 64 長語のデータが転送されていることが分かります。

先程の命令では、グループ 0 の DRAM から同じグループ内にある 1 番 L2B の L2BM にデータを転送しましたが、別のグループにある L2B の L2BM へとデータを転送することも可能です。

この命令では、グループ 0 の DRAM から、グループ 2 にある 1 番 L2B の L2BM へとデータを転送しています。先程とは異なり、サイクル単位で転送できる長語の数が 16 から 8 に減少しています。これは、グループを跨ぐデータの転送では、転送に用いるパスが異なり、スループットが低下してしまうためです。全体として、2 ステップ、8 サイクルかけて 64 長語の転送が完了しています。

L2BM → L1BM 放送

データを転送する際に、複数の転送先のメモリユニットの同じアドレスにデータを書き込む “放送” というタイプの命令があります。

この命令では、1 サイクルごとに 16 長語を奇数番目の L1B の L1BM に放送しています。データが分裂して送り込まれる様子を視覚的に表現しました。

L1BM → L2BM 縮約

また、異なるメモリユニットの同じアドレスに存在する異なるデータについて、加算、最大値、ビット論理和などの演算を用いて一つのデータにまとめる “縮約” というタイプの命令もあります。

先程とデータが逆方向に動いているのが分かります。一般に、データを縮約すると元のデータと異なる値になるため、縮約されたデータは色を濃くして表示しています。

L1BM 間内部マルチキャスト

中にはさらに複雑な動きをする命令もあります。内部マルチキャストと呼ばれている命令では、一つの命令で複数の L1BM から複数の L1BM へとデータを移動することが可能であり、命令に渡す引数によって読み出し元・書き込み先の L1BM が大きく変化します。

このように、単一の L1BM から他の全ての L1BM に放送を行うことや、

複数の L1BM からその他の L1BM にデータを転送・放送することもできます。

これらの挙動は前述のソフトウェア開発者マニュアルで説明されており、疑似コードによって挙動が定義されているのですが、条件が複雑であり、どのような引数に対してどのような動きをするのかを把握するのは簡単ではありません。こうした可視化資料を用いることで、実際に命令がどういう挙動をするかの理解に役立てることができます。

アニメーションはフルスクラッチで作成しているので、命令の種類に応じて補助的な情報を表示することもできます。この命令においては、引数から計算されるビット列を 3 桁で表示しています。

その他の命令

以上で紹介した以外にも、MN-Core 2 には様々なデータ転送に関する命令が存在し、全ての命令についてのアニメーションによる動作例が収録されています。

おわりに

本記事では、MN-Core 2 のデータ転送系命令を可視化する取り組みについて、実際の命令とアニメーションを交えて解説しました。

可視化やアニメーション作成は普段の業務ではあまり使うことのない技術ですが、意外な場所で自身のスキルが活かせるチャンスを発見することもできます。コンパイラ開発チームの周辺にはまだまだ難しい課題が転がっている状態で、PFN では一緒に開発を行う仲間を募集しています。興味のある方はぜひ応募を検討いただければと思います。

追記：MN-Core Challenge にてアニメーションを公開中です

2024 年 8 月現在、MN-Core 2 のアセンブリ言語を用いたコンテストが実施されており、コンテストのページ内において実際にアニメーションを確認することができます。色々動かしてみて、詳細な挙動の理解に役立ててください！