技術要約：コンパイラ・ワールドモデルに向けて

問題提起

テンソルプログラムの最適化は、現代の機械学習システムにおいて不可欠であるが、最も効率的なプログラムを特定するプロセスは、膨大な組合せ探索空間によって阻害されている。デバイス上での計測コストを回避するため、既存のオートスケジューラ（AnsorやAutoTVMなど）は、候補となるプログラムをランク付けするための学習済みコストモデルに依存している。しかし、これらのモデルは通常、候補を静的なコードのスナップショット（例：最終的なTensorIRテキスト）として扱う。

この静的な視点は、以下の2つの決定的な制限をもたらす：

1. 変換コンテキストの喪失： プログラムを生成したステップバイステップの変換プロセスを無視してしまう。その結果、評価器は、あるスケジューリング決定が、先行する決定によって確立されたコンテキストにどのように依存しているかを理解することができない。
2. 構文的ノイズへの敏感さ： 静的なエンコーダは、表面的なコードの差異によって容易に誤認される。例えば、構造的に等価なプログラム（例：ループサイズ $N=1024$ と $N=1023$ の行列乗算）であっても、境界処理（ループ・ピーリング）のためにTensorIRテキストが大きく異なる場合があり、これにより、ハードウェア上の挙動は同一であるにもかかわらず、モデルがそのパフォーマンスの潜在能力を誤分類してしまう原因となる。

手法

著者らは、スケジュールの評価を、プログラム状態に関する**アクション条件付き潜在ダイナミクス（action-conditioned latent dynamics）**として再定義する「コンパイラ・ワールドモデル」フレームワークを提案している。最終的な静的コードのアーティファクトを評価する代わりに、このフレームワークは、連続的な潜在空間を通じてプログラムの状態の進化をシミュレートする。

本フレームワークは、TVM AutoScheduler内に実装された3つの学習済みコンポーネントで構成される。

1. テンソルプログラム状態表現（エンコーダ）

• 入力： TensorIRテキスト（スケジューリング前およびスケジューリング後の状態）。
• アーキテクチャ： 長い入力に適応させた、学習可能なCodeBERTベースのアーキテクチャ。スライディングウィンドウを使用してトークン化とオーバーラップするチャンクのエンコードを行い、それらをマスク付き平均プーリング（masked mean pooling）によって集約し、高密度ベクトル表現 $z = E(s)$ を生成する。
• 学習： コントラスティブ学習（対照学習）を用い、同様のTensorIR状態（同じワークロードソースから派生したもの）を潜在空間内で近くに配置し、異なる状態を遠ざける。これにより、構文的な差異が存在しても意味的な等価性が保持される状態空間を構築する。

2. アクション条件付きマルチステップ状態遷移

• 目的： 中間的なASTを明示的に実体化することなく、一連のスケジューリングアクションによって誘発される終端状態の表現を予測すること。
• メカニズム： 初期状態の埋め込み $z_0$ から開始し、スケジューリングアクションのシーケンス $a_{1:T}$ の影響をステップバイステップでロールアウトする。
  • アクションはシリアル化され、高密度な埋め込みへとエンコードされる。
  • 遷移モデル $F(\cdot)$ は、TransH（元々は知識グラフの表現学習から派生したもの）ベースのアーキテクチャを用いて実装されており、次の潜在状態を予測する：$z_t = F(z_{t-1}, a_t)$。
  • これは、スケジューリングアクションをリレーション固有のハイパープレーンにおける幾何学的変換としてモデル化し、決定のプロベナンス（由来）を追跡し、参照の曖昧さ（例：融合後のループ変数が何を表しているか）を解決する。
• 学習： 状態一致項（グラウンドトゥルースの将来の状態との距離を最小化する）とマージンベースのランキング項（予測された状態がグラウンドトゥルースよりも負例に近いことを保証する）を組み合わせたマルチステップ・ロールアウト損失を用いる。

3. ランキングベースの候補評価

• 入力： 予測された終端状態の表現 $\hat{z}_T$、構造化されたアクション特徴量（例：アクション数、シーケンス長）、およびオプションのハードウェア特徴量。
• 出力： 同一ワークロード内における相対的な候補の優先順位付けのためのスカラー・スコア。
• アーキテクチャ： LambdaRank目的関数を用いて学習されたXGBoostランキングモデル。
• 学習戦略： モデルは絶対的な実行時間を予測するのではなく、測定された実行時間に基づき、ワークロードグループ内での候補のランク付けを学習する。これは、学習によるランキング（learning-to-rank）タスクとして問題を扱う。

主な貢献

1. コンパイラに対するワールドモデルの視点： 本論文は、候補となるスケジュールを静的なスナップショットではなく、アクション条件付きの潜在状態の進化を通じてモデル化するフレームワークを確立することで、テンソル最適化にワールドモデルのパラダイムを適応させた。
2. アクション・ステート・トラジェクトリ・データセット： マルチステップ予測のために、チューニングログを整列されたアクション・ステート・トラジェクトリ（事前スケジュールの状態、アクションシーケンス、中間状態、事後スケジュールの状態）として整理した、TVM/TenSetベースのデータセットを構築した。
3. 効率的な探索の統合： 本フレームワークをTVM AutoSchedulerに実装し、潜在ダイナミクスをモデル化することで、高い探索品質を維持しながらオートチューニングのコストを大幅に削減できることを示した。

実験結果

本手法は、Intel Xeon Gold 6430 CPUおよびNVIDIA RTX 4090 GPU上で、7つのニューラルネットワークモデル（ResNet、BERT、GPT、OPTを含む）と103個の抽出されたサブグラフに対して評価された。

• Performance vs. Ansor： 64回の試行予算において、本手法はAnsorと比較して、代表的なサブグラフのレイテンシをGPUで1.37倍、CPUで1.54倍改善した。
• サンプル効率： 本手法は、わずか1,000回の試行で、Ansor-10K（10,000回の試行）と同等の探索品質を達成した。具体的には、代表的なサブグラフにおいて、算術平均でAnsor-10Kをわずかに上回り、幾何平均において2.2%以内の差で同等の性能を示した。
• エンドツーエンドの加速： PyTorch/PyTorch-opt (cuDNN) と比較して、本手法は幾何平均で4.61倍（CPU）および3.67倍（GPU）の加速を実現し、ピーク加速率は最大58.61倍に達した。
• アブレーション研究： 学習済みの遷移モジュールを除去した場合（静的な初期または最終状態のエンコーディングに戻した場合）、レイテンシが大幅に増加（中央値で約56〜69%）した。これにより、中間状態の進化をモデル化することが最も重要なコンポーネントであることが確認された。

重要性と限界

重要性：
本論文は、コンパイラ最適化を静的なコード評価問題ではなく、状態遷移の動的なプロセスとして扱うことで、スケジューリング決定間の意味的な依存関係を捉えることが可能であると主張している。このアプローチは、構文的なノイズを効果的に抑制し、参照の曖昧さを解決し、より少ないハードウェア計測でより効率的な探索を実現する。

限界（著者による記載）：

1. 探索空間への依存： 本手法は候補の「評価」を改善するものであり、「生成」を改善するものではない。その有効性は、強力なスケジュールを提案できる基礎となるスケジューラの能力に制約される。優れた評価器であっても、一度も提案されなかったスケジュールを復元することはできない。
2. 誤差の累積： 潜在空間における長く複雑なスケジューリング・トラジェクトリにおいて、予測誤差が蓄積し、信頼性が低下する可能性がある。
3. ランキング vs 予測： モデルはワークロード内でのランキング用に訓練されており、絶対的なレイテンシ予測用ではない。そのため、オンラインの優先順位付けには理想的であるが、異なるワークロード間で較正されたパフォーマンス推定を必要とする設定には適していない。