現代のコンピュータは、その内部構造において拡大する分裂を抱えています。一方には、人工知能のために特別に設計された強力なエンジンがあり、これらは数十億もの単純な計算を極めて高速に実行することに長けています。これらのエンジンは、膨大なスピードを得るために、わずかな詳細さを犠牲にして、短く単純な数値を扱うことを得意としています。もう一方には、気象パターンのシミュレーション、原子の結合モデル、流体の流れの予測といった科学的発見の世界があり、そこでは依然として長く精密な数値が求められています。これらの科学的な計算は、安定性と正確性を維持するためにあらゆる細部を必要としますが、それらを処理する標準的なコンピュータ部品は、新しいAIエンジンほど高速で効率的ではないことがよくあります。これはジレンマを生み出しています。科学者は新しいハードウェアのスピードを必要としていますが、彼らの研究が求める精度を失うわけにはいかないのです。
中国のMaginfra Co., Ltd.の研究者たちは、Intel AMXと呼ばれる特定の種類のコンピュータチップを使用して、この溝を埋める方法を模索しました。彼らの目標は、高速で低精度のAIエンジンを、科学に必要な低速で高精度の数学を実行するように「騙す」ことができるかどうかを確認することでした。チップに直接難しい数学を行わせる代わりに、彼らは問題をより小さく単純な断片へと分解しました。非常に長い距離を、1インチ単位の目盛りしかない定規で測ろうとしている場面を想像してみてください。あなたはまず整数部分のインチを測り、次に残った端数を測り、さらに残った極小の破片を測り、それらすべてを足し合わせることで、正確な合計値を導き出すことができます。研究者たちはこの論理を数値にも適用しました。彼らは一つの複雑な数値を、高速なAIエンジンが容易に扱えるいくつかの単純なパーツに分割しました。そして、それらのパーツに対して多くの迅速な計算を行い、最終的に高度に正確な答えを再構成するために、結果を注意深く足し合わせたのです。
チームはこのアプローチを、2つの異なる精度レベルでテストしました。まず、多くの科学的アプリケーションの標準である単精度数学に取り組みました。彼らは、各数値を3つのパーツに分割し、6つの特定の計算を実行することで、既存の最高水準のソフトウェアと同等の精度を実現しつつ、大幅に高速化できることを見出しました。テストしたコンピュータチップにおいて、この手法は標準的な計算方法よりも1.14倍から2.56倍高速に動作しました。このスピードアップは、数値を分割して再組み立てする際のオーバーヘッドが、純粋な計算速度に比べて重要性を失うような、より大きなデータセットにおいて最も顕著に見られました。
さらに、より厳密で、最も要求の厳しい科学シミュレーションに使用される倍精度数学へと移行すると、課題は増大しました。ここでは、研究者は各数値を6つのパーツに分割する必要がありました。計算結果を極めて慎重に再組み立てする必要があるため、プロセスはより複雑になりました。彼らは、6個から21個の小さな計算パーツを保持するという、さまざまなバージョンのこの手法をテストしました。彼らは明確なトレードオフを発見しました。つまり、パーツを多く保持するほど答えはより正確になりますが、同時にプロセスも遅くなるということです。わずか6つのパーツを使用した場合、この手法は非常に大規模な問題に対して標準的なソフトウェアを打ち負かすのに十分な速さであり、最大1.7倍高速に動作しました。しかし、精度を高めるためにパーツを増やしていくと、それらを管理するための余分な作業がスピードの利点を食いつぶしてしまいました。最終的に、21個のパーツを保持しようとすると、たとえより正確であったとしても、標準的な手法よりも遅くなってしまいました。
この研究はまた、このテクニックがあらゆる状況に対する普遍的な解決策ではないことも明らかにしました。この手法は、計算される数値が特定の範囲内に留まっている場合に最も効果を発揮します。これは、限られた長さの定規では、特別な調整なしにはあまりに広大な距離や、あまりに微小な距離を測ることができないのと似ています。研究者たちは、彼らの手法がすべての種類の数値に対して機能するわけではなく、特に極端に大きい、あるいは極端に小さい数値には適さないことを指摘しました。また、既存のソフトウェアとビット単位で完全に一致することを保証するものでもありません。むしろ、これはデータが手法の制限内に収まっている場合に、高いスピードと高い精度を必要とする科学者に新たなツールを提供するものです。低精度のハードウェアを用いて高精度の問題を解決できることを示すことで、この研究は、人工知能のために構築された特化したエンジンが、科学的発見の重労働を加速させる未来を示唆しています。
技術要約：Intel AMXにおけるFP32およびFP64 GEMMのBF16コンポーネント積エミュレーション
1. 問題提起
現代のCPUは、低精度AIワークロード（例：BF16、FP8）に最適化された高スループットの行列エンジン（Intel Advanced Matrix Extensions、以下AMX）をますます統合しています。しかし、計算流体力学や量子化学などの分野における多くの科学計算アプリケーションは、数値的安定性と再現性のために依然としてFP32およびFP64の汎用行列乗算（GEMM）に依存しています。これにより、ハードウェアのミスマッチが生じています。すなわち、AI用に設計された低精度の行列スループットを、高精度な科学計算に直接利用することができないという問題です。
本論文が取り組む中心的な問いは、低精度の行列スループットを、分解、データ移動、および再構成による過度なオーバーヘッドを招くことなく、高精度な科学計算用GEMMを加速するためにアルゴリズム的に転用できるかどうかです。課題は、標準的なライブラリであるIntel oneMKLと同等の数値精度を維持しながら、これを実現することにあります。
2. 手法
本論文は、高精度なオペランドを複数のBF16行列に分解し、選択された低精度積をAMXを用いて計算し、その結果を高精度フォーマットに再構成する、コンポーネント積エミュレーションフレームワークを提案しています。
2.1 浮動小数点フォーマットとドメイン
- BF16: FP32の8ビットの指数範囲を保持しつつ、仮数部の精度を8ビットに削減しています。
- FP32エミュレーション: 標準的なFP32入力を対象とします。
- FP64エミュレーション: FP64の値が通常のBF16指数範囲内に収まる制限されたドメイン（FB64）で動作します。この制限により、パネルごとのスケーリング係数が必要なくなるため、パフォーマンスの余力を維持できます。
2.2 分解戦略
- AMX-FP32パス: 3コンポーネントBF16残差分解を使用します。各FP32オペランドは、直接残差丸めを用いて3つのBF16スライス（A0,A1,A2）に分割されます。これは、共有される指数範囲を利用してFP32における残差の正確な減算を保証する「Henry方式」に従っています。
- AMX-FP64パス: 簡略化された固定6スライスOzaki分解を使用します。各FP64オペランドは6つのBF16スライスに分割されます。動的なスケーリングを用いる完全なOzakiスキームとは異なり、この手法は抽出されたチャンクを（それらが正規値であると仮定して）直接BF16行列として格納するため、明示的なスケール・メタデータや再構成ウェイトの必要性を排除しています。
2.3 積のスケジューリングと再構成
- FP32再構成: 6つの選択されたコンポーネント積（コンポーネント相互作用行列の最初の3つの対角成分）を評価します。決定的なのは、これらの積がFP32 AMXタイル内で直接累積される点です。実装では、オペランド再利用スケジュール（例：A1B1→A1B0→…）を採用しており、これによりアキュムレータタイルをレジスタに留め、タイルへのロード/ストア・トラフィックを最小限に抑えています。
- FP64再構成: 変動する数の積（6、10、15、または21個）を評価します。FP32パスとは異なり、選択された各BF16積はFP32で計算され、メモリに格納され、FP64に拡張された後、FP64ブロックアキュムレータに累積されます。これにより、FP32の累積誤差が高精度な結果を汚染することを防ぎますが、繰り返されるストアとフォーマット変換のために大きなオーバーヘッドが発生します。
2.4 実装の詳細
実装には、256×256の出力ブロックと32×32のマイクロタイルを用いたブロック化GEMM構造が使用されています。メモリ・アクセス・パターンを最適化するために、事前計算されたブロック連続のコンポーネント・バッファに依存しています。Aオペランドは標準的な行優先レイアウトで格納され、BオペランドはAMX BF16ドット積命令に必要なVNNIレイアウトに事前パッキングされます。
3. 主な貢献
- アルゴリズムの架け橋: Intel AMX BF16命令を使用してFP32およびFP64 GEMMをエミュレートする実用的な手法を示し、AI最適化ハードウェアと科学計算の要件との間のギャップを埋めることを実証しました。
- FP32レベルの精度を持つFP32エミュレーション: 6つの積を用いたBF16スケジュールとFP32累積を用いることで、ビット単位の同一性は主張しないものの、oneMKL SGEMMと同等の精度（ノルム相対誤差およびGEMMスケール成分誤差によって測定）を達成できることを示しました。
- 制御可能なFP64エミュレーション: 固定の6スライス分解内において、保持するコンポーネント積の数を変える（6から21個）ことで、FP64エミュレーションにおける精度とパフォーマンスのトレードオフを調整可能にしました。
- 最適化されたデータ移動: オペランド再利用（FP32用）やデータレイアウト（VNNIパッキング、事前計算バッファ）といった特定のスケジューリング戦略を開発し、CPUマトリックスエンジンにおける分解と再構成のオーバーヘッドを軽減しました。
4. 実験結果
実験は、デュアルソケットIntel Xeon Platinum 8462Y+サーバー上で実施され、提案されたカーネルをIntel oneMKL SGEMMおよびDGEMMと比較しました。
- FP32パフォーマンス: AMX-FP32実装は、テストされたすべての正方行列サイズ（N=256からN=32768）において、oneMKL SGEMMを上回りました。加速率は1.14倍から2.56倍の範囲であり、分解のオーバーヘッドが相殺される大きな行列サイズにおいて最も高い利得が見られました。
- FP64パフォーマンス: 結果はより制約的です。
- AMX-FP64-6（6個の積）は、大きな行列（N≥16384）に対して最大1.72倍の加速を実現しました。
- AMX-FP64-10は、最大の行列においてのみわずかな利得を提供します。
- AMX-FP64-15およびAMX-FP64-21は、テストされた構成ではDGEMMを上回りませんでした。これは、追加の積を格納、変換、および累積するためのコストが、利用可能なパフォーマンスの余力を消費してしまうためです。
- 精度:
- AMX-FP32は、10−7の範囲の誤差を達成しており、これはFP32レベルの精度と一致しています。
- AMX-FP64バリアントは、積の数が増えるにつれて誤差が減少することを示しており、21個の積を用いるバリアントはDGEMMレベルの精度に近づきますが、加速を打ち消すほどのパフォーマンスコストが発生します。
5. 意義と主張
本論文は、本研究を完全な高精度ライブラリの代替ではなく、初期のアルゴリズム的研究として位置付けています。その意義は、以下のことを実証した点にあります：
- 注意深いアルゴリズム的分解とデータ移動の最適化を通じて、低精度CPUマトリックスエンジンが、高精度な科学計算に効果的に活用できること。
- エミュレーションには明確なパフォーマンスの天井が存在すること：FP32累積が効率的であるためFP32エミュレーションは非常に効果的ですが、FP64エミュレーションは、AMXタイルレジスタの外側での拡張と累積のオーバーヘッドによって制限されます。
- 本手法は調整可能なトレードオフを提供しており、アプリケーションが精度要件とパフォーマンスニーズのバランスを取るために、積の数を選択できるようにします。
著者らは、制限事項についても明記しています。現在のプロトタイプは正方行列のみをサポートし、入力は（FP64の場合）BF16の指数範囲内である必要があり、オーバーフロー/アンダーフローや例外値を処理しません。今後の課題として、矩形行列への拡張、スケーリング分解によるより広い入力ドメインへの対応、および他のCPUマトリックスエンジン（例：Arm SME）への適応が挙げられています。
毎週最高の computer science 論文をお届け。
スタンフォード、ケンブリッジ、フランス科学アカデミーの研究者に信頼されています。
受信トレイを確認して登録を完了してください。
問題が発生しました。もう一度お試しください。
スパムなし、いつでも解除可能。
週刊ダイジェスト — 最新の研究をわかりやすく。登録