DARPA MOCHA のコンパイラをどのハードウェアで試すか

現代のコンピュータはもはや単一のプロセッサを中心に組み立てられてはいない。今日まともなマシンといえば、異種要素の寄せ集めだ。CPU、GPU、そして機械学習、信号処理、ネットワークそれぞれに向けた専用アクセラレータの増え続ける一族。この寄せ集めから良い性能を引き出すのは容易ではない。まして、コンパイラが一度も見たことのないハードウェアでそれを素早くやるのは、なおさら難しい。DARPA MOCHA——DARPA の「異種アーキテクチャ向け機械学習および最適化誘導コンパイラ」プログラム——は、まさにこの隔たりを埋めるために設けられた。SEI 人工知能部門の先端コンピューティングラボはここ数か月、このプログラムの今後2年の方向を決めるある問いを考え続けてきた。プログラムのコンパイラは、どのハードウェアでテストすべきか。本稿では、この問いにどう取り組んだかを紹介する。MOCHA が何を目指しているのか、SEI がそこでどんな役割を担うのか、候補ハードウェアをどう絞り込んだのか、候補の一覧とそれぞれにどう負荷をかけて試したのか、そして最終的にどこへ落ち着いたのか。

コンピュータ最適化の自動化

MOCHA は DARPA 情報処理技術局のプログラムで、Dr. Howard Shrobe が統括している。この取り組みを後押ししているのは、よく知られた痛点だ。従来のコンパイラは、CPU、GPU、アプリケーションアクセラレータが混ざった異種構成のために効率的な機械語を生成することを想定して設計されていない。新しいアクセラレータを使いこなすには、開発者はたいてい専用コードを手書きし、ベンダーがチューニングしたライブラリに頼る。この道は機能するが、遅く、高くつき、そして気づかないうちにベンダーロックインを助長する。アプリケーションをある専有ライブラリ向けに書いてしまえば、別のハードウェアへ移すには書き直しが要る。コンパイラを拡張して真に新しい計算ユニットに対応させるのは、コンパイラの専門家にしかできない手作業だ。時間がかかり、間違いやすく、新しいシリコンが登場するテンポに追いつかない。MOCHA の仮説はこうだ。データ駆動型の手法、機械学習、先進的な最適化によってこのプロセスを加速し、最小限の人手でコンパイラを新しいハードウェアへ素早く適応させられる。鍵となる洞察がある。対象ハードウェアの性能モデルがコンパイルの各段階を駆動しており、そのモデルを手作業で構築することが主要なボトルネックになっている。もしこのモデルを、生成したコードを実ハードウェア上で測定し、アーキテクチャ文書を掘り下げることで自動生成できれば、新しいデバイスを支えるコストは大きく下がる。DARPA MOCHA の指針の一つが ALARA——人の関与を合理的に達成可能な最低水準に保つこと——だ。ALARA は MOCHA の二重の狙いを表している。コンパイラを新しいハードウェアに素早く対応させること、そして異種ハードウェアをまたいでコンパイルできるようにすること。ある新しいチップで速く動くだけでは足りない。このプログラムが問うのは、大きく異なる計算ユニットを幅広く同時にカバーするのに、どれだけの人手が要るかだ。

SEI の役割

SEI の先端コンピューティングラボは AI Division に属し、政府チームを支援しており、重点はテストと評価にある。政府チームは DARPA のプログラムマネージャと数名のシステムエンジニアリング技術支援者（SETA）で構成される。具体的には、私たちは選択肢を集めて評価し、プログラムマネージャが決定を下すのに必要な情報を届ける。どの計算ユニットをプログラムに採り入れるかは彼が決める。その後、私たちは評価マシンを組み立てて維持し、選んだ計算ユニットを統合し、MOCHA の測定手法を確立して、参加チームの結果を公平に比較できるようにする。参加チームはコンパイラ技術を開発する。私たちの仕事は、特徴がはっきりしていて代表的、かつ難易度の適切なターゲットを彼らに与え、評価そのものの妥当性をできる限り確保することだ。DARPA はプログラム終了時に、MOCHA が6種類の異なる計算タイプをカバーしていることを計画している。計算タイプの区別はハードウェアアーキテクチャだけでなく、命令セットとプログラミングモデルが異なるかどうかでも決まる。この定義に従えば、データセンター GPU、フィールドプログラマブルゲートアレイ（FPGA）構造と空間 AI エンジンアレイを融合させたデバイス、長ベクトルプロセッサ、そして RISC-V とデータフローを組み合わせた AI アクセラレータは、4つの異なるタイプに当たる。ぱっと見では後ろの3つがすべて「アクセラレータ」に分類されかねないとしても、だ。このプログラムの要点は、アーキテクチャと命令セットアーキテクチャ（ISA）の境界をまたいだ、迅速で低コストな再ターゲティングを示すことにある。だからこそ、ターゲット集合におけるアーキテクチャの多様性は欠かせない。現実的な制約もある。選んだハードウェアは実際に評価マシンに収まらなければならない。このマシンは Intel Core Ultra 9 285K を中心に組んだワークステーション級タワーで、それ自体がいくつかの計算タイプを備えている。CPU 上の AVX2 SIMD、統合 Xe GPU、そしてニューラルプロセッシングユニット。加えて PCIe 5.0 接続があり、ベースライン参照として NVIDIA RTX 4500 Ada カードがすでに装着されている。したがって候補アクセラレータは PCIe カードの形で提供され、単一タワーの筐体、電力、冷却の条件に合う必要がある。

候補アクセラレータを選ぶ5つの要因

候補リストを決めたのは5つの要因だ。入手性、成熟度、価格の手頃さ、プログラマビリティ、そしてデバイスを評価マシンに収められるか。入手性と設置条件によって、それ自体は興味深い選択肢がいくつも落ちた。サーバ一体型でしか出荷されないウェハスケールエンジンや再構成可能データフローシステム、そして入手も設置もできないクラウド専用のアクセラレータだ。価格の手頃さは、マシン本体の残り全部を合わせたより高いチップに対して冷静でいさせてくれた。最も影響が大きかったのはプログラマビリティ、具体的には各ターゲットプラットフォームにおけるコンパイラとマルチレベル中間表現（MLIR）サポートの現状だ。MOCHA の参加者の大半は LLVM と MLIR のエコシステム上で構築している。LLVM の鍵となる革新は拡張可能な IR と、開発者がそれを扱うためのツールチェーンだった。MLIR はより新しい革新で、異なる抽象レベルで IR を定義することでこの能力を拡張した。これは現代のコンパイラ基盤をつなぐ組織のようなものになり、コンパイラが複数の抽象レベルで計算を表現し、それを段階的に具体的なデバイスへ下ろしていけるようにする。ターゲットプラットフォームにすでに MLIR か LLVM の経路があれば、参加者はそこに接続し、本当に面白い研究に力を注げる。再ターゲット可能なコード生成、学習型コストモデル、最適化の選択、そして最終的には異種ユニット間での作業の分割だ。ターゲットプラットフォームが閉じたブラックボックスで、ベンダーが事前にチューニングした高水準の推論スタックを通してしか触れないなら、どれだけ機械学習を注ぎ込んでも、MOCHA のコンパイラができることはごく限られてしまうだろう。

開放的な低レベルアクセスと閉鎖的なベンダーライブラリの間の緊張は、プロジェクト全体を貫いている。プロプライエタリなライブラリ向けにプログラムするのは楽だが、それは新しいハードウェアを素早くサポートするという目標と根本的に衝突する。ライブラリは、ベンダーがすでにサポートすると決めたハードウェアにしか存在しないからだ。各候補を評価する際には、そのプログラミングモデルがどれだけ開かれているか、そして MLIR ベースでハードウェアに直接到達する経路が存在するか、あるいは現実的に開通させられるかどうかを、注意深く調べた。

候補リスト

これらの基準でふるいにかけた結果、目標のショートリストは GPU、FPGA と AI エンジンを組み合わせた空間アーキテクチャデバイス、ベクトルプロセッサ、それぞれ性格の異なる AI アクセラレータ数種、そしてネットワークチップをカバーするものになった。およそ6種類の真に異なる計算タイプを支えるのに十分な広がりである。

AMD Instinct MI350P —— データセンター向け GPU（CDNA 4）。ROCm/HIP でプログラムし、rocMLIR と Triton の経路を通じて成熟した MLIR サポートを持つ。注目すべきは、新発表の PCIe フォームファクタ製品であり、OAM クラスの Instinct 計算能力を標準スロットに持ち込んでいる。
AMD Versal ACAP —— FPGA ロジック、空間アーキテクチャの AI エンジンアレイ、ARM コアを1つのデバイスに統合したヘテロジニアスデバイス。オープンな mlir-aie/IRON ツールチェーンとその Peano LLVM バックエンドにより、AI エンジンアレイはハードウェアに近い、真に興味深い MLIR ターゲットになっている。
Intel Data Center GPU Max 1100 —— Xe-HPC GPU。oneAPI/SYCL でプログラムし、Intel Triton XPU バックエンドと SPIR-V を経由して MLIR に接続できる。Max シリーズで PCIe カードとして提供される唯一のモデルである。
Intel Gaudi 3 —— 行列エンジンと VLIW-SIMD テンソルエンジンを備えた AI アクセラレータ。カスタム演算子はオープンな TPC-LLVM コンパイラで TPC-C を用いて記述し、グラフコンパイラが MLIR ベースの融合演算子を構築する。ただしグラフコンパイラ自体は依然としてプロプライエタリである。
Intel Xe iGPU —— 評価機の CPU にすでに統合されている内蔵 GPU。Max 1100 と oneAPI/SYCL インターフェースおよび MLIR 経路を共有するため、コストゼロの移植ターゲットになる。
NEC SX-Aurora TSUBASA —— PCIe カード上の古典的な長ベクトルプロセッサで、アップストリームの LLVM バックエンドを持つ。ベクトル化、自動チューニング、帯域律速の HPC 演算子にアーキテクチャの多様性をもたらす。
Qualcomm Cloud AI 100 —— 推論アクセラレータ。主な経路は ONNX/PyTorch だが、「閉鎖的」という評判とは裏腹に、アップストリーム LLVM ベースのオープンなコンパイラを同梱し、低レベルのカスタム演算子の登録にも対応している。
Tenstorrent Blackhole（p100a/p150a） —— 手頃な価格で異常にオープンな AI アクセラレータ。Tensix コアと RISC-V を組み合わせ、ネイティブで完全にオープンソースの MLIR コンパイラ（tt-mlir/tt-forge）を備え、PyTorch、JAX、ONNX からのモデルを受け取れる。
MangoBoost BoostX DPU と GPUBoost RNIC —— ネットワーク向けの部品（SmartNIC/DPU 1枚と RDMA 網カード1枚）。網羅性のために挙げたが、汎用計算向けの MLIR や LLVM の演算子経路はない。

実行チームへの意見聴取と、こちらが出した難題

上の候補リストを確定する前に、草稿を各実行チームに送り、2つ質問した。我々が見落としている、リストに入れるべきターゲットは何か。そのうちどれが、あなたたちのツールチェーンでは到底サポートできないものか。回答は率直で、役に立った。各チームは、どこに本物の LLVM/MLIR バックエンドがあり、どこにないかを指摘し、価値のすべてをクローズドな推論スタックに依存するターゲットには反対を唱え、ネットワーク関連は自分たちが取り組む計算ターゲットではないと明言した。候補の選別は一貫して1つの原則に従った。ターゲットは難しくあるべきだが、不可能であってはならない。易しすぎるターゲット——技術スタックが成熟し、ベンダーがすでに最適化を済ませているもの——は、MOCHA の中核的主張である「迅速で低コストな適応」を真に検証できない。難しい部分はとっくにベンダーが片付けているからだ。難しすぎるターゲット——ドキュメントのないブラックボックス、低レベルプログラミングインターフェースなし、マイクロアーキテクチャのモデル化も不可能——は作業を行き詰まらせるだけであり、実行チームは研究を進める代わりにツールチェーンと格闘する時間を費やすことになる。理想的なターゲットとは、到達して分析できるだけの開放性を持ち、かつ実行チームがすでに慣れ親しんだものと十分に異なるため、方向転換が彼らのコンパイラ、コストモデル、カーネル生成器を真に試すことになるものである。

ハードウェア選定：AMD RDNA 4 GPU と Tenstorrent Tensix コア

あの候補リストを起草した時点と本稿執筆の時点の間に、実際に入手可能なハードウェアへとさらに絞り込んだ結果、範囲は狭まり、プロジェクトの最初の調達は2枚のカードに集中した。AMD Radeon AI PRO R9700 と Tenstorrent Blackhole p150a である。AMD ROCm/HIP 技術スタックに基づく GPU は、最初から一席を占めることが決まっていた。実行チームの間で、これは満場一致の第一候補だった。真面目な、NVIDIA 以外の GPU であり、rocMLIR と Triton を経由する本物の MLIR 経路を持ち、複数の実行提案の中核にあるテンソル、スパース、グラフ、コストモデリングの作業と直接関係する。当初は新発表の MI350P、すなわち AMD のフラッグシップ Instinct 製品の PCIe 形態を選んでいた。しかし AMD Research のエンジニアから、彼ら自身も MI350P のチップを待っており、入手は2027年春になる見込みだと告げられた。2年目の評価を始めるのに必要な時期をはるかに超えている。そこで Radeon AI PRO R9700 に切り替えた。今すぐ買えるワークステーションカードである。標準の PCIe スロットに挿さり、同じ ROCm/HIP プログラミングモデルを使い、同じ MLIR と Triton の経路を通る。狙っていた AMD GPU という計算タイプを保ちつつ、実行チームが今年実際にマシンへ組み込めるものになっている。R9700 は思いがけず優れた MOCHA ターゲットになったが、その理由はまさにこのプロジェクトの核心に触れている。新しいアーキテクチャ（RDNA 4）に基づくため、AMD 自身が手作業でチューニングしたアセンブリライブラリがまだ完全には対応していない。いくつかのライブラリはサポートアーキテクチャのリストをハードコードしており、このカードを除外し、遭遇すると黙って低速な経路にフォールバックする。実際に機能するのはコンパイラの経路だ。Triton と MLIR の経路は、新しいアーキテクチャ向けのネイティブカーネルを実行時に生成する。まさにプリコンパイル済みのベンダーライブラリが破綻する箇所である。これは MOCHA の主張を縮小した見本だ。コンパイラが生成したコードは、手作業でチューニングされたライブラリが手も足も出ない場所で新しいチップに適応し直せる。そして MOCHA コンパイラには、ベンダーがほぼ最適まで磨き上げたベースラインと向き合うのではなく、現実的で測定可能な性能の伸びしろがあるということでもある。アーキテクチャ上の対照として選んだのが Tenstorrent Blackhole p150a だ。R9700 が成熟した LLVM バックエンド上で動く GPU なのに対し、Blackhole はまったく異なる。Tensix コア群と汎用 RISC-V コアを組み合わせ、Tenstorrent の完全にオープンで MLIR を中心としたコンパイラスタック（tt-mlir と tt-forge）でプログラムする。このスタックは StableHLO を経由して PyTorch、JAX、ONNX からのモデルを受け取る。その MLIR の話は、我々が評価した中でおそらく最も強い。Blackhole はオープンソースコンパイラの名指しのターゲットであり、開発は完全に公開で進み、実行チームのコンパイラが接続するための文書化された StableHLO の入口点がある。ここでの困難は門をくぐることではなく、慣れ親しんだ LLVM ターゲットモデルではなく、カスタム dialect の塔を学ぶことにある。

MOCHA が計測し、最終的には自動化しようとしているのは、まさにこうしたリダイレクションの課題だ。最初のバッチには、3 つ目のアーキテクチャ種別として AMD Versal も加えたいと考えていた。AI engine アレイは空間データフロー構造で、GPU とも Tensix アレイとも大きく異なり、mlir-aie には魅力的なオープン MLIR ツールチェーンもある。だが、AI engine を外部に公開しつつ PCIe カードとして出荷され、評価機に収まる Versal デバイスは結局見つからず、AI engine という種別はプロジェクト後半まで持ち越しとなった。2 枚のカードがプロジェクトにもたらすのは、重複ではなく補完的なリダイレクションの問題だ。R9700 が試すのは成熟したエコシステム内でのリダイレクションであり、ただしそのエコシステムのライブラリが、この新しい具体的なアーキテクチャにはまだ成熟していないというだけのこと。Blackhole が試すのは、まったく新しいアーキテクチャ種別へのリダイレクションで、技術スタックは若いが完全にオープンだ。これに評価機もともとの計算能力、すなわち AVX2 SIMD、Xe iGPU、NPU、そして NVIDIA のベースラインが加わることで、6 つの計算種別という目標に向けて実質的な一歩を踏み出した。3 枚の PCIe カードを同時にこのマシンに収めることはできない。6 つの計算種別という指標を満たすために、ほかにどのカードを同時に載せるかは、プロセスの後半で決める。

次世代ハードウェア：数か月のチューニングから数日の計測へ

最初のハードウェアを選ぶのは、始まりにすぎない。この先もプロジェクト全体につきまとう問題がいくつかある。重要なもののひとつは抽象化のレベルだ。十分に不透明なデバイスでは、ISA レベルでのプログラミング効率が、ベンダー製の高水準ツールにいつまでも及ばない可能性がある。だが、そうしたプロプライエタリなツールに依存すれば、新しいハードウェアを素早くサポートするという目標と衝突する。どのレベルを狙うべきかを見極め、ブラックボックスのマイクロアーキテクチャをモデル化する能力を高めることが、今後どのデバイスを対象に含める価値があるかを決める。最初の 2 つの目標はすでに選定済みで、作業の重心は選択から実行へ移った。Radeon AI PRO R9700 と Tenstorrent Blackhole p150a を評価マシンに統合し、その特性を把握し、各実装者のコンパイラがこれらに対して何を達成できるかを比較する計測パイプラインを組む。さらに、6 つの計算種別から恩恵を受けるワークロードを定義し、その 6 つの計算種別向けに手作業で実装し、自動での分解に各実装者が挑む。MOCHA が成功すれば、次に登場する新しいアクセラレータを採用するのに、専門家による数か月の手作業チューニングではなく、数日の計測で済む世界が手に入る。ハードウェアを正しく選ぶことが、そのすべてを解明する第一歩だ。

その他のリソース

DARPA MOCHA（Machine learning and Optimization-guided Compilers for Heterogeneous Architectures）をダウンロード - https://www.darpa.mil/research/programs/mocha-machine-learning