Exo — exocompilation

languagecompileroptimizationhardwareformal-methods

高性能カーネル(GEMM・畳み込み等)を人間が手で最適化しながら、その最適化が元の素朴定義と等価であることを機械が保証する、Python 埋め込みの user-schedulable 言語(USL)。MIT CSAIL(Jonathan Ragan-Kelley グループ)発。BLIS が新アーキごとに手書きする microkernel 特殊化を桁違いに楽にし、ベンダーライブラリ/手書きアセンブリ級の性能を狙う。Halide/TVM と同じ アルゴリズム / schedule 分離を継ぐが、他の USL と分かつ柱が2本ある:

  1. exocompilation — ハードウェアの定義(命令・メモリ・構成レジスタ)をコンパイラ本体でなくユーザーライブラリに置く。コンパイラは小さく信頼できるコアに保つ。
  2. rewrite-based scheduling — schedule の各操作を LoopIR → LoopIR の書き換えとして実装し、各書き換えの意味保存を effect analysis で機械検証する(Halide/TVM の lowering-based とは別系統)。

設計原則

2言語構造: object code と scheduling

Exo のプログラムは性質の違う2つのコードから成る。

object code(= アルゴリズム)

@proc を付けた Python 関数。構文は Python 風だが Python ではなく、パーサが Exo IR(LoopIR)へ変換する。素朴で読める「正しい仕様」をここに固定する。

@proc
def rank_k_reduce_6x16(K: size, A: f32[6,K] @ DRAM, B: f32[K,16] @ DRAM, C: f32[6,16] @ DRAM):
    for i in seq(0, 6):
        for j in seq(0, 16):
            for k in seq(0, K):
                C[i, j] += A[i, k] * B[k, j]

解析を成立させる3つの制約(ここが Exo の効きどころ):

制約内容なぜ要るか
affine 添字のみdata[i*j] 等の非アフィン添字は不可(定数 div/mod は可)添字推論・依存解析を決定可能に保つ
値依存制御フロー禁止if data[i] < 3.0: は書けない。制御値とバッファ値を分離データの動きを値から切り離し、等価性を静的に解けるようにする
引数エイリアス禁止foo(y, y) 不可(C の __restrict__ 相当)解析の精度を保つ

値依存制御フローの禁止は almide-kernel の着想の土台でもある。データの動きが値から独立する → permutation なら写像が値に依存しない、という性質が言語レベルで保証される。almide-kernel を参照。

scheduling(= 最適化)

普通の Python コードとして書く。各 scheduling primitive は proc を受け取り新しい proc を返す。アルゴリズムは読みやすいまま固定され、最適化は別プログラム(書き換えの列)として分離される。

avx = rename(rank_k_reduce_6x16, "rank_k_reduce_6x16_scheduled")
avx = reorder_loops(avx, 'j k')          # ループ入れ替え
avx = divide_loop(avx, 'for j in _: _', 8, ['jo','ji'], perfect=True)  # = Halide/TVM の split
avx = stage_mem(avx, 'for k in _:_', 'C[0:6, 0:16]', 'C_reg')         # レジスタへステージング
avx = set_memory(avx, 'C_reg:_', AVX2)   # メモリ空間を AVX2 レジスタに
avx = replace_all(avx, mm256_loadu_ps)   # ループを AVX2 命令へ置換

各ステップ後に print(avx) で途中状態が読める proc として出る(rewrite-based の利点)。primitive は6カテゴリ — buffer / loop & scope / config / subproc / memory・precision・parallelism / other。代表例: divide_loop・reorder_loops・fission・fuse・unroll_loop(loop)、stage_mem・lift_alloc・expand_dim・bind_expr(buffer/mem)、replace/replace_all(命令置換)、高階コンビネータ repeat・simplify。

exocompilation: ハードをユーザーコードで定義する

新アクセラレータ対応にコンパイラを fork しない。命令・メモリ・構成状態をライブラリで宣言する。

命令 @instr — proc の本体に命令の意味(Exo object code)を書き、デコレータ文字列に吐く C/intrinsic を書く({arg_data} がプレースホルダ):

@instr("{dst_data} = vld1q_f32(&{src_data});")
def neon_vld_4xf32(dst: [f32][4] @ Neon, src: [f32][4] @ DRAM):
    assert stride(src, 0) == 1
    assert stride(dst, 0) == 1
    for i in seq(0, 4):
        dst[i] = src[i]          # ← この本体が「命令の仕様」

replace(proc, cursor, instr) が、命令本体とコード片を unification modulo linear equalities で照合する: subproc の引数を未知数、コード片の自由変数を既知とし、整数制御式の等価性を線形方程式系として解く。AST が(整数制御以外)厳密一致し方程式が解ければ、そのコード片は命令の意味と等価と証明されたことになり、安全にハード命令へ差し替えられる。「専用バックエンド不要 — 重い処理は Exo 内で完結し、命令の意味そのものが仕様になる」。

メモリ Memory サブクラス — alloc/free/window/can_read/global_ を C 文字列を返すメソッドとして実装。ベクタレジスタも「独自の確保・アクセス codegen を持つメモリ」として表現(AVX2, AVX512, Neon, GEMM_SCRATCH 等を同梱、ユーザー定義が推奨)。

構成状態 config — アクセラレータの構成レジスタをモデル化。delete_config で冗長な構成書き込みを除去。config を跨いだ推論は PLDI 論文の中心貢献の一つ。

正しさ: 2層の機械保証

層何を保証手段実装
frontendassert 充足・配列範囲外なしSMT ソルバ(PySMT + Z3/CVC4)boundscheck.py
rewrite各 schedule 操作が意味を保存(全入力で等価)effect analysis(read/write 効果・依存の解析)Check_*(new_eff.py)

schedule は人間が書くが、各 primitive の正しさはコンパイラが担保する — LoopIR_scheduling.py の各 primitive が Check_* 安全検査を呼び、効果解析で書き換えが等価かを判定する。この rewrite 等価性の層こそ almide-kernel が「permutation は値非依存だから恒等インデックスを1回通せば effect analysis/SMT 無しで全証明できる」と崩した着想元。Exo は一般の変換を相手に効果解析を回すが、permutation という制約を入れると 1 実行で decidable な断片に落ちる。

コンパイルパイプラインとリポジトリ

Python AST → UAST(無型) → typecheck → LoopIR(主 IR)
   → [rewrite] scheduling primitives: LoopIR → LoopIR(不変・各段 printable)
   → [backend] mem/prec/win/parallel 解析 → C 生成

Exo は独自レクサを持たず Python の構文解析を乗っ取る。フル proc は UAST、コード片照合用は穴 _ を持つ PAST にパースする。CLI は exocc file.py → file.c + file.h。主要ファイル:

Exo 2(ASPLOS 2025): scheduling 言語そのものを育てる

PLDI 版の primitive は低水準で冗長(上の AVX2 例も約20ステップ)。Exo 2 は schedule 操作自体をユーザー拡張可能にした — trusted な細粒度 primitive を合成して自前の scheduling ライブラリを書き、tile2D・auto-vectorize、果ては Halide の compute_at 風操作まで再現できる。

鍵は2つ:

結果: scheduling ライブラリが 80以上のカーネルにわたって労力を償却し、scheduling コード量を桁違いに削減、3 プラットフォームで SOTA に伍する性能。

立ち位置

Halide / TVMExo
algorithm/schedule 分離ありあり
schedule の実装lowering-based(コンパイラ内で段階的に降ろす)rewrite-based(IR→IR、各段 printable・検査可能)
ハード対応コンパイラに組み込みユーザーライブラリ(exocompilation)
命令選択コンパイラの裁量replace をユーザーが明示+ 等価性を機械検証
正しさ主にテストeffect analysis + SMT で書き換え等価を保証

出典の区別

関連

Graph

Backlinks