Julia の SSA 形式 IR

背景

バージョン 0.7 以降の Julia のコンパイラは新しい SSA 形式の中間表現を利用します。これ以前のコンパイラは低水準形式の Julia AST から直接 LLVM IR を生成していました。低水準形式は抽象構文木から構文的な抽象化を取り除いたものですが、それでも見た目は抽象構文木とあまり変わりません。しばらくして最適化を助けるために SSA 値が IR に導入され、さらに IR は線形化 (関数の引数が定数と SSA 値だけからなる形に変形) されました。しかし IR に Phi ノード (後方辺や条件付き制御フローの再合流で必要なノード) が存在しないために SSA でない値 (スロット) が IR に残ってしまっており、SSA 形式を活かしたミドルエンドの最適化ができない状態でした。そういった最適化を SSA 形式の表現を使わずに可能にするために英雄的な労力が注がれましたが、最終的に SSA 形式の表現はどうしても必要であるという結論になりました。

新しい IR ノード

新しい IR 表現には Phi ノード・Pi ノード・PhiC ノード・Upsilon ノードという四つのノードが新しく加わります (最後の二つは例外処理でだけ使われます)。

Phi ノードと Pi ノード

Phi ノードは一般的な SSA という抽象化に含まれる概念です (分からなければ上記のリンクを参照してください)。Julia IR において Phi ノードは次の構造体で表されます:

struct PhiNode
    edges::Vector{Int}
    values::Vector{Any}
end

ここで edgesvalues は同じ長さのベクトルです。正準表現 (コード生成とインタープリタで使われる表現) では edges の要素がこのノードに向かう文の番号を表し、例えば edges15 が含まれるなら 15 番目の文はこの Phi ノードを指す gotogotoinnot・暗黙のフォールスルーのいずれかである必要があります。values は SSA 値もしくは定数であり、実行されたパスによっては変数に値が割り当てられない可能性もあります。ただしミドルエンドの最適化が終わると未定義変数の確認処理が挿入されて定義されているかどうかが真偽値として表されるようになるので、Phi ノードが使われるときコードジェネレータは割り当てられた値が対応するスロットに存在すると仮定して構いません。また正当なコードであっても対応関係が不完全になる、つまり Phi ノードに向かう辺の一部が値を持っていない可能性があります。その場合は対応する値が使われていないことを動的に保証する必要があります。

Pi ノードは自身が支配する基礎ブロックにおいて暗黙に仮定できると静的に証明された情報をエンコードします。これは ABCD: Eliminating Array Bounds Checks on Demand という論文で導入されたテクニック、あるいは LLVM の述語情報ノード (predicate info node) と概念的には等価です。Pi ノードの役割を示す例として、次のコードを考えます:

# %x は Union{Int, Float64} 型の SSA 値
%x::Union{Int, Float64}
if isa(x, Int)
    # x を使う
else
    # x を使う
end

このコードは、証明できる述語を挿入することで次のコードに変換できます:

# %x は Union{Int, Float64} 型の SSA 値
%x::Union{Int, Float64}
if isa(x, Int)
    %x_int = PiNode(x, Int)
    # %x_int を使う
else
    %x_float = PiNode(x, Float64)
    # %x_float を使う
end

Pi ノードは値に影響を及ぼさないので、通常はインタープリタからは無視されます。一方コンパイラでは Pi ノードがコード生成を引き起こす場合があります (型共用体の自動的な分割表現がボックス化されていないプレーンな表現に変更されるなど)。Pi ノードが有用なのは、値のパス条件が def-use チェーンの走査 (パス条件を利用する最適化の多くがいずれにせよ行う処理) で簡単に計算できるためです。

PhiC ノードと Upsilon ノード

例外処理があると SSA 形式が多少複雑になります。例外処理を示す制御フローの辺が IR に加わるので、それを追跡しなければなりません。これを行うアプローチの一つで LLVM によっても採用されているものが、例外を送出する可能性のある関数呼び出しを基礎ブロックの終端に入れ、例外を捕捉するハンドラへの明示的な制御フロー辺を追加するというものです:

@function_that_may_throw() を呼び出して、
%regular で起こった例外が %catch に戻るよう印を付ける。

regular:
# 通常の制御フローはここに続く。

catch:
# 例外があるとここに飛ぶ。

しかし Julia のように最適化パイプラインの先頭で例外が起こる関数がどれかが分からない言語では、このアプローチは問題です。全ての関数呼び出し (Julia では全ての文) が例外を送出すると保守的に仮定しなければならず、これにはいくつか欠点があります。一方では全ての基礎ブロックのスコープが事実上一つの関数呼び出しに狭められるために基礎ブロックのレベルで演算を行う意味が無くなってしまい、他方では全てのエラー捕捉ブロックが n*m 個の Phi ノードを引数に持つことになります (n はクリティカルリージョンに含まれる命令の個数、mcatch ブロック内で生存する変数の個数)。

これに対処するために用意されたのが Upsilon ノードと PhiC ノードです (PhiC の C は catch の C です。IR の装飾付き出力では Uncode の下付き文字が利用できないので φᶜ と表示されます)。二つのノードの捉え方はいくつかありますが、PhiC を複数回ストア/単一ロードのユニークスロットからのロードと考え、Upsilon を PhiC に対応するストアと考えるのがおそらく最も簡単でしょう。PhiC のオペランドは自身の内部スロットに値が格納された Upsilon ノードのリストですが、これに対して Upsilon ノードの引数には格納先の PhiC ノードが含まれません。

例外捕捉の仕組みがこうなっているのは、他の部分の SSA IR と自然に統合させるためです。例えば使われない PhiC ノードと Upsilon ノードは安全に削除できます。多くの IR パスは PhiC ノードを Phi ノードのように扱うことができ、必要に応じて use-def チェーンを追うことで PhiC ノードと Upsilon ノードを正しい位置に新しく作成できます。このスキームの結果として、Upsilon ノード (および PhiC ノードの引数) の個数は考えている変数に (SSA 変換前に) 割り当てられた値の個数に比例するようになり、クリティカルリージョン内に文の個数とは関係がなくなります。

このスキームの実際の動作例を示します。関数

@noinline opaque() = invokelatest(identity, nothing) # 不透明な処理
function foo()
    local y
    x = 1
    try
        y = 2
        opaque()
        y = 3
        error()
    catch
    end
    (x, y)
end

は、次の IR に対応します (関係ない型情報を省略しています):

1 ─       nothing::Nothing
2 ─ %2  = $(Expr(:enter, #4))
3 ─ %3  = ϒ (false)
│   %4  = ϒ (#undef)
│   %5  = ϒ (1)
│   %6  = ϒ (true)
│   %7  = ϒ (2)
│         invoke Main.opaque()::Any
│   %9  = ϒ (true)
│   %10 = ϒ (3)
│         invoke Main.error()::Union{}
└──       $(Expr(:unreachable))::Union{}
4 ┄ %13 = φᶜ (%3, %6, %9)::Bool
│   %14 = φᶜ (%4, %7, %10)::Core.Compiler.MaybeUndef{Int64}
│   %15 = φᶜ (%5)::Core.Compiler.Const(1, false)
└──       $(Expr(:leave, 1))
5 ─       $(Expr(:pop_exception, :(%2)))::Any
│         $(Expr(:throw_undef_if_not, :y, :(%13)))::Any
│   %19 = Core.tuple(%15, %14)
└──       return %19

クリティカルリージョンに入るときに生存している値にはリージョンの先頭で Upsilon ノードが割り当てられる点に特に注目してください。これは catch ブロックが関数外からの見えない制御フロー辺を持つものとみなされるためです。この結果として catch ブロックを支配する SSA 値は存在せず、外からの値は必ず φᶜ ノードを通して取得されるようになります。

メインの SSA データ構造

メインに使われるデータ構造 SSAIR は議論に値します。SSAIR は LLVM IR と Webkit の B3 IR からインスピレーションを受けています。このデータ構造のコアは文が平坦に並んだベクトルであり、それぞれの文はこのベクトルにおける位置に応じて SSA 値が割り当てられます (例えば添え字 1 の場所にある文には SSAValue(1) でアクセスできます)。加えて各 SSA 値が持つ型も管理されます。SSA 値には厳密に一度だけ値が割り当てられるので、この型は対応した添え字を持つ式の評価結果の型でもあります。

この表現は代入の明示的なエンコードが必要ないために効率の面で優れますが、ノードの順序が意味論的に重要であるためにノードの並び替えや挿入を独立に行えない欠点も当然あります。さらに use リストは保持されない (def から use を計算するには明示的な計算が必要です ──逆の def リストは添え字アクセスで自明に得られます) ので、LLVM にある RAUW (replace-all-uses-with) 操作は利用できません。

この問題を解決するために、次の処理が行われます:

適切なノードの挿入・自明なコピーの伝播・変更された SSA 値の改名を行って上記のデータ構造をコンパクトにする関数 compact! が提供されます。ただしこのスキームが賢いのは、compact! を後に続くパスの一部として遅延させて行える点です。多くの最適化パスは文のリスト全体を走査しながら解析や改変を行いますが、そのときはイテレータ IncrementalCompact が利用できます。このイテレータは反復ごとに必要になったコンパクト化を行ってからノードに対する新しい添え字とノード自身を返します。このイテレータを使えば、反復中に def-use チェーンの走査や IR に対する変更および削除が可能になります (ただし挿入は行えません)。

この取り決めの背後にあるアイデアは「いずれにせよ最適化パスは文リストが格納されるメモリに触れてメモリアクセスによるペナルティを食らう。これに比較すれば、追加で情報を個別に保持することによるオーバーヘッドは無視できる程度だろう (さらに IR を変更するときにそういったデータ構造を保持するオーバーヘッドも節約できる)」というものです。


日本語 Julia 書籍 (Amazon アソシエイト)
1 から始める Julia プログラミング
Julia プログラミングクックブック―言語仕様からデータ分析、機械学習、数値計算まで
スタンフォード ベクトル・行列からはじめる最適化数学