論文情報

著者: C Hong, S Bhatia, YS Shao
論文概要リンク: https://arxiv.org/abs/2505.18574
論文PDFリンク: https://arxiv.org/pdf/2505.18574

要約

本論文では、Tensor Accelerator（テンソル処理用ハードウェアアクセラレータ）向けのコード最適化を大規模言語モデル（LLM）を用いて自動化する「Autocomp」を提案している。Tensor Acceleratorのコードは専用DSLで記述されることが多く、性能を最大限引き出すには高度な最適化技術とハードウェアの詳細な理解が必要である。Autocompは最適化の計画フェーズとコード生成フェーズの二段階からなるLLM駆動の探索フレームワークを用い、ハードウェアのフィードバックを参照しながらコードを反復的に改善することにより、ベンダー提供のライブラリや手動チューニング済みコードを上回る高性能コードを自動生成できる。

主要なポイント

Tensor Accelerator向けの低リソース特殊DSLに対して、学習不要のLLMベースの反復的な最適化探索手法を初めて提案。
最適化の計画をメニューから選択し、その計画に基づくコード生成を二段階プロンプトで実施。ハードウェアの性能・正当性フィードバックを活用しビームサーチで探索。
GEMM（行列積）、畳み込み、細粒度線形代数など多様な代表的ワークロードで、ベンダー提供ライブラリに対し最大5.6倍、手動チューニング済みコードに対しても1.1～1.4倍の高速化に成功。
生成した最適化スケジュールは類似の演算に再利用可能で、探索コストを削減しつつ最高速度の24%向上を実現。

メソッド

二段階最適化プロンプト

計画生成フェーズ
LLMにアクセラレータのISA、現在のコード、ハードウェア性能フィードバック（レイテンシ、スクラッチパッドの使用率など）、および「ループ展開」「ループ並べ替え」「ダブルバッファリング」などの最適化メニューを与え、一つの最適化計画を自然言語で生成。
コード生成フェーズ
生成した計画に基づき、元コードを変換して最適化コードを生成。チリングなど複雑な変換にはインコンテキスト学習で例示を利用。

ビームサーチ探索
ビーム幅B=6ほどで複数の最適化計画を並列検討し、多様性を確保。候補コードは正当性検証（シミュレータで動作検証、結果の正しさ確認）および性能評価を実施し、良好な候補のみ次の探索に継続。
多様性向上技術
最適化メニューの選択肢をランダムに削除（dropout）することで最適化計画の多様性を促進し、LLMのアンサンブル利用によってコード生成の多様性と品質を向上。
スケジュール再利用
最適化済みスケジュールを類似したテンソル演算の初期探索として用いることで、最適化の計算コストを大幅に削減。

意義・影響

本研究は、これまで専門家の高度な知識や大規模なトレーニングに依存していたテンソルアクセラレータコードの最適化をLLMの推論力を活用し自動化した初の試みである。
Autocompはハードウェアの詳細なISP知識をプロンプトに含めることで、従来困難だった低リソースDSLのコード生成と複雑な最適化を実現。
ハードウェア性能を用いたフィードバックループと多様な探索手法により、現状最高性能のソフトウェアコードを上回る成果を示した。
生成された最適化スケジュールの再利用性は実用的な高速化開発に寄与し、新たなアクセラレータへの適応や展開コストを低減。
今後はより汎用的なメタスケジュール化やさらに広範なスケジュール共有の研究により、さらに効率的な最適化とアクセラレータ設計プロセスの自動化が期待される。

以上より、AutocompはTensor Acceleratorに特化した高度なコード最適化をLLMで手軽かつ効果的に実現する新たな自動化技術として、ハードウェアとソフトウェアの協調設計に重要なインパクトを持つことが示された。