Swiftにおける所有権付きSSA形式

こんにちは!Merpay Advent Calendar 2020 の7日目は、メルペイ エキスパートチームインターン の kateinoigakukun がお送りします。

この記事ではSwiftコンパイラ内部で使われている、最適化パスの正確性を保証するためのOwnership SSA(Static Single Assignment、静的単一代入)という技術がSwiftのメモリ管理を支え、パフォーマンス向上にどのように寄与しているのかを紹介します。

注意:この記事で触れるOwnership(所有権)とは主にメモリにおけるコピーと参照の管理の手法の名称で、題材にするテーマとしてSwift言語自体に備わる所有権ではなく、中間言語における所有権モデルのお話をします。

Automatic Reference Countingの歴史

Ownership SSAの背景を語る上で重要になってくるのがAutomatic Reference Countingの歴史です。

Swiftのメモリ管理はARC(Automatic Reference Counting, 自動参照カウント)によって行われています。Objective-CではMRC(Manual Reference Counting, 手動参照カウント)でメモリ管理を行うこともできますが、原理的にもプログラマによるエラーが多く発生してしまうため、現在は基本的にARCが使われます。

ARCはMRCで頻繁に発生していたUse After Freeやメモリの開放し忘れの問題を、参照カウントのインクリメント/デクリメント操作を言語レベルでコンパイラによって自動生成することで解決しました。

しかし、MRCではプログラマが決めていたオブジェクトのライフタイムを、ARCではコンパイラが自動で保守的に決めるため、MRCに対してインクリメント/デクリメント回数が多くなる傾向があります。

たとえばMRCが可能な架空のSwiftを考えると、 cat変数とanimal変数のライフタイムはプログラマにとって明確なので、1回のrelease操作で済みます。

class Animal {
  func bark() {}
}
class Cat: Animal {
  override func bark() { print("meow") }
}

let cat = Cat()
let animal: Animal = cat as Animal
animal.bark()
release(cat)

一方でARCを用いるSwiftでは素直にコード生成すると以下のように、animal変数にcatをエイリアスするためretain(cat)が挿入され、animal変数のライフタイムが終わるタイミングでrelease(animal)が挿入されてしまいます。

class Animal {
  func bark() {}
}
class Cat: Animal {
  override func bark() { print("meow") }
}

let cat = Cat()
// retain(cat) // Inserted by compiler
let animal: Animal = cat as Animal

animal.bark()

// release(animal)
// release(cat)

このような無駄なretain/release操作は、実行時パフォーマンスに無視できないオーバーヘッドを生み出します。特にSwiftにおいてはオブジェクトの唯一性がArrayやDictionary等で使われているCopy on Write(メモリアロケーションにおける最適化戦略)の重要な役割を果たすため、パフォーマンスへのインパクトが大きくなります。たとえば、以下のようにArrayのストレージに対して無駄なretain命令が実行されることで、ストレージが他のオブジェクトと共有されることになり、変更操作のタイミングで配列全体のコピーが行われてしまいます。

var items: [Int] = [1, 2, 3]
// retain(items.storage)
let shared = items
let value = shared[2]
items[0] = 1 + value // コピー発生
// release(shared.storage)
// release(items.storage)

このようなオーバーヘッドを回避するため、コンパイラはインクリメント/デクリメント操作を可能な限り最適化を行います。

ARCの最適化

Swift固有の最適化は基本的にSIL(Swift Intermediate Language)という中間言語に対して適用されます。SILはSSA形式の中間言語で、LLVM IRに比べて高レベルな操作を記述できます。ARCのretain/release操作もSILの命令として表現されます。

たとえば上に示したようなARCのコードは以下のようにSILに変換されます。(説明のため実際に出力されるSILに手を加えています)

sil hidden @main : $@convention(thin) () -> () {
bb0:

    // let cat = Cat()
  %0 = metatype $@thick Cat.Type
  %1 = function_ref @"Cat.__allocating_init()" : $@convention(method) (@thick Cat.Type) -> @owned Cat
  %2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat

    // let animal: Animal = cat as Animal
  strong_retain %2 : $Cat
  %5 = upcast %2 : $Cat to $Animal

  // animal.bark()
  %7 = class_method %5 : $Animal, #Animal.bark : (Animal) -> () -> (), $@convention(method) (@guaranteed Animal) -> ()
  %8 = apply %7(%5) : $@convention(method) (@guaranteed Animal) -> ()

  strong_release %5 : $Animal
  strong_release %2 : $Cat

  %13 = tuple ()
  return %13 : $()
}

コンパイラはこのように保守的に生成されたコードを、できるだけMRCを用いたコードに近づけるための最適化を行います。たとえば、%5の実態は %2 と共有されていることがupcast命令から分かるため、%5に対するretain/release命令は%2に対するretain/release命令に書き換えることができます。

sil hidden @main : $@convention(thin) () -> () {
bb0:

    // let cat = Cat()
  %0 = metatype $@thick Cat.Type
  %1 = function_ref @"Cat.__allocating_init()" : $@convention(method) (@thick Cat.Type) -> @owned Cat
  %2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat

    // let animal: Animal = cat as Animal
  strong_retain %2 : $Cat
  %5 = upcast %2 : $Cat to $Animal

  // animal.bark()
  %7 = class_method %5 : $Animal, #Animal.bark : (Animal) -> () -> (), $@convention(method) (@guaranteed Animal) -> ()
  %8 = apply %7(%5) : $@convention(method) (@guaranteed Animal) -> ()

  strong_release %2 : $Cat // replaced
  strong_release %2 : $Cat

  %13 = tuple ()
  return %13 : $()
}

この命令列に対して、静的にretain/release操作をカウントすると、参照カウントが1以上であることが保証できる生存範囲の間に、1ペアの不要なretain/release操作があることを検出できます。このような解析を経て、コンパイラは元のコードを以下のような命令列に最適化できます。

sil hidden @main : $@convention(thin) () -> () {
bb0:

    // let cat = Cat()
  %0 = metatype $@thick Cat.Type
  %1 = function_ref @"Cat.__allocating_init()" : $@convention(method) (@thick Cat.Type) -> @owned Cat
  %2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat

    // let animal: Animal = cat as Animal
  %5 = upcast %2 : $Cat to $Animal

  // animal.bark()
  %7 = class_method %5 : $Animal, #Animal.bark : (Animal) -> () -> (), $@convention(method) (@guaranteed Animal) -> ()
  %8 = apply %7(%5) : $@convention(method) (@guaranteed Animal) -> ()

  strong_release %2 : $Cat

  %13 = tuple ()
  return %13 : $()
}

無事、MRCを用いた場合と同等の命令列に最適化できましたが、この最適化手法にはいくつかの問題点があります。

まず、upcast命令で生成される%5%2が同じオブジェクトを指す、という情報がSILの言語レベルで表現されておらず、命令固有の知識となっているため、最適化が複雑になりやすい点が挙げられます。

また、SILレベルではretain/releaseのペア関係が表現されておらず、生存範囲が命令列からの推論ベースで決められるため、命令順序の入れ替えによるアグレッシブな性能向上が難しくなっています。

さらに、retain/releaseのペア関係が言語レベルで静的に検証されていないため、MRCをプログラマが管理する場合と同様に、最適化プログラムのミスコンパイルによりメモリリークやUse After Freeが発生する可能性が高くなります。

このような問題を解決するためにAppleでSwiftコンパイラのパフォーマンス最適化やビルドシステムの改善進めているMichael Gottesmanさんが導入した仕組みがOwnership SSAです。

Ownership SSA

Ownership SSAは通常のSSAに所有権の概念を追加した拡張形式です。

コンパイラは各最適化パスごとにOwnership SSAを検証し、所有権モデルに違反するようなコードが生成されると、問題のある命令を診断してくれます。検証は基本的にコンパイラのデバッグ時にのみ行われるため、実際にコンパイラのエンドユーザーが使う際にはオーバーヘッドはほとんどありません。さらに、最適化の妥当性を検証するだけでなく、所有権情報を使った積極的な最適化も行われています。

形式的なOwnership SSAの定義はMichael Gottesmanさんのフォーラム投稿を参照してください。

SIL Ownership Model Proposal (Refreshed)

以下に有効なOwnership SSAの例を示します。所有権無しの場合と表現力を比較をすると、begin_borrow/end_borrowによって借用を表現していたり、retainの代わりにcopy_valueで意味上の値をコピーしその結果をupcastオペランドに使うことでコピー情報を解析しやすくしている、といった違いがあります。

class Animal {
  func bark()
}
class Cat: Animal {
  override func bark()
}

sil hidden [ossa] @main : $@convention(thin) () -> () {
bb0:
  // let cat = Cat()
  %0 = metatype $@thick Cat.Type
  %1 = function_ref @"Cat.__allocating_init()" : $@convention(method) (@thick Cat.Type) -> @owned Cat
  %2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat

  // let animal: Animal = cat as Animal
  %3 = begin_borrow %2 : $Cat
  %4 = copy_value %3 : $Cat
  %5 = upcast %4 : $Cat to $Animal
  end_borrow %3 : $Cat

  // animal.bark()
  %7 = begin_borrow %5 : $Animal
  %8 = class_method %7 : $Animal, #Animal.bark : (Animal) -> () -> (), $@convention(method) (@guaranteed Animal) -> () // user: %11
  %9 = apply %8(%7) : $@convention(method) (@guaranteed Animal) -> ()
  end_borrow %9 : $Animal

  destroy_value %5 : $Animal
  destroy_value %2 : $Cat

  %13 = tuple ()
  return %13 : $()
}

Ownership SSAのセマンティクス

Ownership SSAのセマンティクスを簡単に紹介します。

消費

まず、Ownership SSAにおける命令はオペランドの値を「消費」することがあります。消費された値はその後の命令列で使用することはできません。たとえば、destroy_value命令やupcast命令はオペランドの値を消費します。また、特にupcast命令は消費した値の所有権を結果に転送し、ムーブセマンティクスを表現します。ムーブセマンティクスは「コピーしない」ことを表現する強力な道具です。特にSwiftにおいては唯一性が静的に検証しやすく成ることで、前述のCopy on Writeを使いながらパフォーマンスを担保できるようになります。

このような命令のセマンティクスの上で、Ownership SSAには「全ての値は必ず一度だけ消費されなければならない」という制約があり、この制約を元にSSAを検証することでメモリリークやUse After Freeを防いでいます。

// let cat = Cat()
%0 = metatype $@thick Cat.Type
%1 = function_ref @"Cat.__allocating_init()" : $@convention(method) (@thick Cat.Type) -> @owned Cat
%2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat

// let animal: Animal = cat as Animal
%3 = begin_borrow %2 : $Cat
%4 = copy_value %3 : $Cat
%5 = upcast %4 : $Cat to $Animal
end_borrow %3 : $Cat

...

destroy_value %5 : $Animal
// destroy_value %2 : $Cat // <- ❗

もし、上のSILプログラムのように %2を破棄し忘れた場合、Swiftコンパイラの開発ツールであるsil-optコマンドを使うことで以下のように診断を出力してくれます。

$ sil-opt main.sil -sil-ownership-verifier-enable-testing -ownership-verifier-textual-error-dumper -o /dev/null
Error#: 0. Begin Error in Function: 'main'
Error! Found a leaked owned value that was never consumed.
Value:   %2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat // user: %3

Error#: 0. End Error in Function: 'main'

借用

値の借用はbegin_borrow/end_borrowによって表現され、そのスコープ内での値の生存を保証し、スコープの中で値が消費されないことを検証します。

上の例ではbarkメソッドを呼ぶ際にanimal変数を借用し、値を消費すること無くメソッド呼び出しを表現しています。

// let cat = Cat()
%0 = metatype $@thick Cat.Type
%1 = function_ref @"Cat.__allocating_init()" : $@convention(method) (@thick Cat.Type) -> @owned Cat
%2 = apply %1(%0) : $@convention(method) (@thick Cat.Type) -> @owned Cat

// let animal: Animal = cat as Animal
%3 = begin_borrow %2 : $Cat
%4 = copy_value %3 : $Cat
%5 = upcast %4 : $Cat to $Animal
end_borrow %3 : $Cat

// animal.bark()
%6 = begin_borrow %5 : $Animal
destroy_value %6 // <- ❗
%8 = class_method %7 : $Animal, #Animal.bark : (Animal) -> () -> (), $@convention(method) (@guaranteed Animal) -> () // user: %11
%9 = apply %8(%7) : $@convention(method) (@guaranteed Animal) -> ()
end_borrow %9 : $Animal

たとえば、上のように借用中の値%6をdestroy_valueで破棄した場合、所有権モデルに違反するため、次のような診断が出力されミスコンパイルに気がつけます。

$ sil-opt main.sil -sil-ownership-verifier-enable-testing -ownership-verifier-textual-error-dumper -o /dev/null
Error#: 0. Begin Error in Function: 'main'
Found outside of lifetime use?!
Value:   %5 = upcast %4 : $Cat to $Animal                // users: %8, %7
Consuming User:   destroy_value %5 : $Animal                      // id: %8
Non Consuming User:   end_borrow %7 : $Animal                         // id: %11
Block: bb0

Error#: 0. End Error in Function: 'main'

この他にもObjective-Cとの相互運用のためのルールやストレージポインタのライフタイム管理など興味深い機能がいくつかあるので、実装を追ってみると楽しいかもしれないですね。

まとめ

Ownership SSAが必要となっている経緯と所有権モデルの概要について簡単に説明しました。

Ownership SSAはSwiftを使ってプログラムを書く際にはあまり表に現れない機能ですが、コンパイルの正確性を保証する重要な機能です。私たちのプログラムが正確に最適化され動いている裏で活躍している技術がより理解されることで、Swift言語の信頼度が上がると嬉しいです。

明日のMerpay Advent Calendar 2020 執筆担当は、 Frontend Engineer のtanakaworldさんです。引き続きお楽しみください。