SOLAR:AI駆動の光速性能分析
SOLARは、PyTorchおよびJAXのソースコードから検証済みの光速(SOL)性能限界を自動的に導出するフレームワークです。生成型大規模言語モデルと決定論的コンポーネントを組み合わせ、非融合、融合、キャッシュ対応のSOL限界を提供し、多忠実度分析をサポートします。KernelBench、JAX/Flaxモデル、ロボティクスワークロードでの実験により、複数レベルのヘッドルーム分析、最適化機会の発見、クロスプラットフォーム探索、逆ルーフラインハードウェア構成の4つのユースケースが実証されました。
ディープラーニングモデルの最適化において、「モデルはターゲットハードウェア上でどれだけ高速に実行できるのか?」「現在の実装は理論限界からどの程度離れているのか?」という問いは、ソフトウェア、ハードウェア、アルゴリズムの最適化全体にわたる中心的な課題です。光速(Speed-of-Light, SOL)分析は、与えられたアーキテクチャにおけるワークロードの理論的最小実行時間を計算することでこれらの問いに答えを提供しますが、従来のSOL限界の導出は手動で行われ、エラーが発生しやすく、急速なモデル開発から切り離されていました。このギャップを埋めるために、研究者らはSOLAR(Speed-of-Light Analysis Runner)フレームワークを提案しました。SOLARは、PyTorchおよびJAXのソースコードから検証済みのSOL性能限界を自動的に導出する完全自動化システムです。
SOLARのフローは、生成AIと決定論的手法を巧みに融合しています。まず、大規模言語モデル(LLM)ベースのフロントエンドが任意のソースプログラムを実行可能なAffine Loop中間表現(IR)に変換し、出力比較によりゼロ違反で検証します。次に、決定論的フローがそのIRをアインシュタイン総和グラフ(einsum graph)に昇華し、最後に分析バックエンドが非融合、融合、キャッシュ対応のSOL限界を計算します。この設計により、SOLARは広範なオペレーターと言語をカバーし、観測された違反ゼロの検証済み限界を生成するとともに、多忠実度分析によって限界を厳密化し、最適化の洞察を明らかにします。
研究チームはKernelBench、JAX/Flaxモデル、ロボティクスワークロードを含む複数のベンチマークでSOLARを評価しました。実験では、4つの主要なユースケースが実証されました:異なる忠実度レベルでのヘッドルーム分析(headroom analysis)、最適化機会の特定、クロスプラットフォーム探索、および逆ルーフライン(inverse-roofline)によるハードウェア構成。これらの機能により、手動でエラーがちだった分析プロセスが自動化され、再現可能なパイプラインに変わります。
SOLARの登場は、性能分析の自動化における重要な進歩を示します。手動導出の不確実性と遅延を排除することで、研究者やエンジニアは最適化の効果を迅速に評価し、ハードウェアリソースをより賢く配分できるようになり、ディープラーニングモデルの展開と反復を加速します。論文の詳細はarXiv(番号2606.26383)で公開されており、コードとデータは近くオープンソース化される予定です。