[論文レビュー] GauS: Differentiable Scheduling Optimization via Gaussian Reparameterization

Yaohui Cai, Vesal Bakhtazad+2 more|arXiv (Cornell University)|Feb 23, 2026

Parallel Computing and Optimization Techniques被引用数 0

ひとこと要約

GauS はガウス再パラメータ化を導入し、演算子スケジューリングを微分可能な連続問題としてモデル化することで、2|V| 個のパラメータとパレート最適なスピード-品質トレードオフを実現し、スケーラブルでGPU対応の最適化を達成します。

ABSTRACT

Efficient operator scheduling is a fundamental challenge in software compilation and hardware synthesis. While recent differentiable approaches have sought to replace traditional ones like exact solvers or heuristics with gradient-based search, they typically rely on categorical distributions that fail to capture the ordinal nature of time and suffer from a parameter space that scales poorly. In this paper, we propose a novel differentiable framework, GauS, that models operator scheduling as a stochastic relaxation using Gaussian distributions, which fully utilize modern parallel computing devices like GPUs. By representing schedules as continuous Gaussian variables, we successfully capture the ordinal nature of time and reduce the optimization space by orders of magnitude. Our method is highly flexible to represent various objectives and constraints, which provides the first differentiable formulation for the complex pipelined scheduling problem. We evaluate our method on a range of benchmarks, demonstrating that Gaus achieves Pareto-optimal results.

研究の動機と目的

• 演算子スケジューリングの最適化空間の複雑さとサイズを削減する。
• 正規化およびパイプライン化（モジュロ）スケジューリングをモデル化できる微分可能なフレームワークを提供する。
• 既存のソルバーおよび以前の微分可能な方法よりも高速に高品質でパレート最適なスケジュールを達成する。
• 潜時、リソース、メモリ、再帰制約を扱える柔軟性を維持する。

提案手法

• 各演算子のスケジュールを平均 mu_i と標準偏差 sigma_i を持つガウス確率変数として表現する。
• カテゴリ的アプローチに比べて最適化パラメータを 2|V| に削減し、大規模グラフを可能にする。
• ガウス分布の期待値として differentiable な目的関数と制約違反を計算する（例：Dep, Lat, Mem, Rec）。
• 連続的な mu_i, sigma_i を離散的なステップ割り当てに結びつけるために確率積分 P_i^d を用いて丸めでリンクする。
• Augmented Lagrangian Method（ALM）を適用して目的最小化と制約充足の動的バランスを取る。
• mu_i を丸めて実現可能性を検証し、衝突を修正する離散スケジュールを抽出し、さらなる最適化のために再初期化する。
• ASAP/ALAP 境界またはリストスケジューリング信号を用いて初期化し、時間的近接性を活用する。

Figure 1 : Impact of optimization objectives on scheduling — Impact of optimization objectives on scheduling. Nodes $v_{i}$ (operators) show resource requirements $r_{i}$ (number of dots) and bitwidths of storage units required $b_{i}$ (number of arrows). Max depth ( $D$ ) sets to $3$ . Left schedul

実験結果

リサーチクエスチョン

• RQ1 ガウス再パラメータ化はスケジューリングの序数的・時間的近接構造をカテゴリ的緩和よりもうまく捉えられるか。
• RQ2 GauS はノード tens-of-thousands の大規模グラフにスケールし、解の品質を維持できるか。
• RQ3 GauS は正規スケジューリングとモジュロスケジューリングの遅延、リソース、メモリ、再帰制約を同時に最適化できるか。
• RQ4 GauS は厳密ソルバーおよび以前の微分可能なアプローチと比較して速度とパレート最適性でどうか。
• RQ5 ALM ベースの制約対応が解の実現可能性と収束に与える影響は何か。

主な発見

• GauS は評価済みベンチマークでパレート最適な結果を達成する。
• GauS は線形のパラメータ空間（2|V|）を維持し、以前の O(D|V|) アプローチに比べて大規模グラフを可能にする。
• GauS はベースライン（GS-Schedule や商用ソルバーを含む）よりも一貫して高品質なスケジュールを、15分の制限内で高速に提供する。
• GauS は大規模なグラフで見られる GPU メモリの爆発を回避し、パラメータ効率が良いことを示す。
• GauS はレイテンシ制約付きリソース・通信最適化、メモリフットプリント最適化、再帰制約付きモジュロスケジューリングを含む複数の形式をサポートする。
• GauS は複雑な問題でほぼ100% の GPU 利用率を達成し、並列性の利用において以前の微分可能な手法を上回る。

Figure 2 : Illustration of Gaussian reparameterization — Each operator $v_{i}$ is parameterized as an independent continuous random variable $X_{i}\sim\mathcal{N}(\mu_{i},\sigma_{i}^{2})$ over the scheduling timeline. The means $\mu_{i}$ represents the expected execution step, while the standard dev