【Swift】async/awaitはどのようにして動いているのかLLVMのレベルで調べてみる

Swift 5.5によって導入されたConcurrencyのうちの、async/awaitがどのように動いているのか調べてみました。

きっかけ

私はもともとasync/awaitなどはJavascriptで使っていたので、Swiftに導入されると聞いたときも「ようやく便利になるね」ぐらいにしか思っていなかったのですが、iOSDC2021のこの発表を見たときにたまげました。

特に驚いたのが、await式がサスペンションポイント(Suspension point)となり、実行が一時保留されるが、その間もスレッドは別の処理に使用されることや、再開されたときは別のスレッドに回ることもあり得るし、同じスレッドの場合もありうる、という話でした。

しかも、並列する複数のタスク間でのコンテキストスイッチは発生せず、呼び出しのコストのみで実現できるので、並列処理を軽量に実行できるそうです。

「なんでこんなことが実現できるの！？」というのが今回興味を持ったきっかけでした。

私はもともと、Mugichaというプログラミング言語を作ったことがあるのですが、この言語はコンパイルしてLLVM-IRを出力することができます。

（👆気に入った方はスターしていただけると嬉しいです）

なので、LLVMは多少知っていました。このため、SwiftがLLVMレベルでどのように実現しているのか調べてみることにしてみました。

ではさっそく、この謎を読み解いていきましょう。

async/awaitとCoroutine

swiftにおける、async/awaitはCoroutine（コルーチン）と同じような仕組みを使用しているようです。

こちらの動画にはCoroutineの仕組みが解説されており、swiftのasync/awaitについてもあわせて説明がありました。

動画の内容をまとめると次のようになります。

動画まとめその１：Coroutineの実装で考慮すべき点

1. 制御権の移動をどのように実現するか

2通りの方法がある。

• コンテキストスイッチする
  スレッドを複数生成してコンテキストスイッチする。
  コンパイラの実装は簡単で、ランタイムで処理できるが、実行時のオーバーヘッドが大きい。

• 関数を分割する
  サスペンションポイントで関数をramp function(降下関数）と、resume fuction(再開関数）に分ける。
  分割方法は、再開関数を1つの関数にまとめるか、サスペンションポイントごとに複数に分割するかなどがある。

2. ローカル変数をどうやって保存するか

3通りの方法がある。

• ローカル関数をスタックに積んだままコンテキストスイッチする（stackful corotuines）
  関数の分割が不要。
  OSレベルでの実装が必要。

• 関数分割して、それぞれローカル関数を保持する(side allocation）
  関数分割が必要。
  ローカル関数の保持にヒープメモリのアロケートが必要。

• スタックの同居（cohabitation)
  関数分割が必要。
  コルーチンに入るときには普通にスタックフレームにローカル変数をpushするがyieldで制御を戻すときはスタックをPOPしない
  コンパイラの実装は複雑。

3. データをどうやって生成するか（yielding data)

2通りの方法がある。

• ramp/resume functionから値を返す
  すぐに値を利用するのに便利

• 固定したメモリに保存する
  async/awaitなど値の利用が遅れてやって来る場合に便利

動画まとめその２：Swiftで要求されたこと

1. async/awaitは特殊なケースのジェネレータであること
2. 呼び出し元と呼び出し先で頻繁に情報をやりとりする
3. yield valueと効率的なアクセスができることを優先する
4. Returned-Continuation Loweringを使用する

Returned-Continuation Loweringは、コルーチンcoroutine lowering（コルーチンのハンドリング手法）の一種です。（なんて和訳したら良いのかわかりませんでした）

サスペンドポイントが「生成された値(yield values)」のリストを取得し、このリストが継続関数と呼ばれる関数ポインタとともに呼び出し元に返されます。コルーチンは、この継続関数ポインターを呼び出すことで再開します。

WWDC21での解説

WWDC21でも、Swift Concurrencyの仕組みについて触れている部分があります。

この中で、async functionsのローカル変数はheap領域を用いていることが解説されています。

(WWDC21「Swift concurrency: Behind the scenes」より引用)

２つの動画から得られた仮説

2つの動画からわかることをまとめます。
まず、SwiftではReturned-Continuation Lowering を使用している、という説明がありました。このため、制御権の移動は、関数分割の手法を用いていることがわかります。また、この手法の場合呼び出し先の関数は、制御を戻すためのポインターを受け取っていると考えられます。

また、WWDC21の動画から、ローカル変数の保持は、side allocationの手法を使っていることがわかります。

LLVM-IRを出力して確認してみる

仮説を検証するために、次のコードをLLVM-IRに変換してみます。
mybarfunc関数が、await付きでmyfoofunc関数を呼び出しています。

なるべくシンプルにしているためエントリー部分もなく、実行しても何も起きません😉

simple_async_await.swift

@available(macOS 12.0.0, *)
func myfoofunc()  async {
    print("ok")
}

@available(macOS 12.0.0, *)
func mybarfunc()  async {
  await myfoofunc()
}

この呼出関係がどうなるか先程の仮説によると、mybarfuncはawaitのところで２つに分割されます。１つめがramp function, ２つめがresume functionとなります。処理の流れとしてはmybarfunc(ramp)からスタートしてまずmyfoofuncが呼ばれ、次に、mybarfunc(resume)が呼ばれます。

図にすると次のようになります。

LLVM-IRを出力してこの仮説を確認してみましょう。

$ swiftc -emit-ir simple_async_await.swift

このぐらいのコードでも300行ぐらいのLLVM-IRが出力されます。
まず、関数宣言のところだけ見てみましょう。

simple_async_await.ll

; myfoofuncの宣言
define hidden swifttailcc void @"$s18simple_async_await9myfoofuncyyYaF"(%swift.context* swiftasync %0) #0 {

; mybarfuncの宣言その１
define hidden swifttailcc void @"$s18simple_async_await9mybarfuncyyYaF"(%swift.context* swiftasync %0) #0 {

; mybarfuncの宣言その２
define internal swifttailcc void @"$s18simple_async_await9mybarfuncyyYaFTQ0_"(i8* swiftasync %0) #0 {

それぞれの関数名をdemangleすると次のようになります。

simple_async_await.ll

; myfoofuncの宣言
simple_async_await.myfoofunc() async -> ()

; mybarfuncの宣言その１
simple_async_await.mybarfunc() async -> ()

; mybarfuncの宣言その２
await resume partial function for simple_async_await.mybarfunc() async -> ()

mybarfuncは、２つの関数に分かれていて、２つめはresume partial functionがついていることがわかります。これは、関数の分割であり、１つ目がramp function、２つめがresume functionになっていることがわかります。

次に、これらの関数の宣言には swifttailcc という呼び出し規約が入っていることがわかります。

swifttailccは、次のように説明されています。

この呼び出し規則は、ほとんどの点で swiftcc と似ていますが、呼び出し先がスタックの引数領域をポップするので、tailcc のように必須のテールコールが可能になります。

tail call呼び出し規約を用いると、末尾位置の呼び出しは常に末尾呼び出しが最適化します。

末尾呼び出しの最適化とは何でしょうか。Wikipediaから引用します。

末尾呼出しのコードを、戻り先を保存しないジャンプに変換することによって、スタックの累積を無くし、効率の向上などを図る手法である。

末尾呼び出しの最適化を用いると、スタックを消費せずにジャンプすることができます。

なんでこんなことをしたいのかというと、よくあるのが再帰処理の最適化で、再帰呼び出しを、ジャンプに置き換えることでスタックの消費を防ぐことができます。

では、この仕組を使ってどのようにしてasync functionから制御を戻しているのでしょうか？
この疑問を確認するために、呼び出されるmyfoofuncのLLVM-IR表現をみてみましょう。

simple_async_await.ll

define hidden swifttailcc void @"$s18simple_async_await9myfoofuncyyYaF"(%swift.context* swiftasync %0) #0 {
entry:
  call void @coro.devirt.trigger(i8* null)
  %1 = alloca %swift.context*, align 8
  store %swift.context* %0, %swift.context** %1, align 8
  %2 = call swiftcc { %swift.bridge*, i8* } @"$ss27_allocateUninitializedArrayySayxG_BptBwlF"(i64 1, %swift.type* getelementptr inbounds (%swift.full_type, %swift.full_type* @"$sypN", i32 0, i32 1))
  %3 = extractvalue { %swift.bridge*, i8* } %2, 0
  %4 = extractvalue { %swift.bridge*, i8* } %2, 1
  %5 = bitcast i8* %4 to %Any*
  %6 = call swiftcc { i64, %swift.bridge* } @"$sSS21_builtinStringLiteral17utf8CodeUnitCount7isASCIISSBp_BwBi1_tcfC"(i8* getelementptr inbounds ([3 x i8], [3 x i8]* @0, i64 0, i64 0), i64 2, i1 true)
  %7 = extractvalue { i64, %swift.bridge* } %6, 0
  %8 = extractvalue { i64, %swift.bridge* } %6, 1
  %9 = getelementptr inbounds %Any, %Any* %5, i32 0, i32 1
  store %swift.type* @"$sSSN", %swift.type** %9, align 8
  %10 = getelementptr inbounds %Any, %Any* %5, i32 0, i32 0
  %11 = bitcast [24 x i8]* %10 to %TSS*
  %._guts = getelementptr inbounds %TSS, %TSS* %11, i32 0, i32 0
  %._guts._object = getelementptr inbounds %Ts11_StringGutsV, %Ts11_StringGutsV* %._guts, i32 0, i32 0
  %._guts._object._countAndFlagsBits = getelementptr inbounds %Ts13_StringObjectV, %Ts13_StringObjectV* %._guts._object, i32 0, i32 0
  %._guts._object._countAndFlagsBits._value = getelementptr inbounds %Ts6UInt64V, %Ts6UInt64V* %._guts._object._countAndFlagsBits, i32 0, i32 0
  store i64 %7, i64* %._guts._object._countAndFlagsBits._value, align 8
  %._guts._object._object = getelementptr inbounds %Ts13_StringObjectV, %Ts13_StringObjectV* %._guts._object, i32 0, i32 1
  store %swift.bridge* %8, %swift.bridge** %._guts._object._object, align 8
  %12 = call swiftcc %swift.bridge* @"$ss27_finalizeUninitializedArrayySayxGABnlF"(%swift.bridge* %3, %swift.type* getelementptr inbounds (%swift.full_type, %swift.full_type* @"$sypN", i32 0, i32 1))
  %13 = call swiftcc { i64, %swift.bridge* } @"$ss5print_9separator10terminatoryypd_S2StFfA0_"()
  %14 = extractvalue { i64, %swift.bridge* } %13, 0
  %15 = extractvalue { i64, %swift.bridge* } %13, 1
  %16 = call swiftcc { i64, %swift.bridge* } @"$ss5print_9separator10terminatoryypd_S2StFfA1_"()
  %17 = extractvalue { i64, %swift.bridge* } %16, 0
  %18 = extractvalue { i64, %swift.bridge* } %16, 1
  call swiftcc void @"$ss5print_9separator10terminatoryypd_S2StF"(%swift.bridge* %12, i64 %14, %swift.bridge* %15, i64 %17, %swift.bridge* %18)
  call void @swift_bridgeObjectRelease(%swift.bridge* %18) #1
  call void @swift_bridgeObjectRelease(%swift.bridge* %15) #1
  call void @swift_bridgeObjectRelease(%swift.bridge* %12) #1
  %19 = load %swift.context*, %swift.context** %1, align 8
  %20 = bitcast %swift.context* %19 to <{ %swift.context*, void (%swift.context*)*, i32 }>*
  %21 = getelementptr inbounds <{ %swift.context*, void (%swift.context*)*, i32 }>, <{ %swift.context*, void (%swift.context*)*, i32 }>* %20, i32 0, i32 1
  %22 = load void (%swift.context*)*, void (%swift.context*)** %21, align 8
  musttail call swifttailcc void %22(%swift.context* swiftasync %19) #1
  ret void
}

最後のところで、 musttail call という予約語が使われています。これは、末尾呼び出しの最適化を必須にした状態で関数を呼び出していることを示しています。

何を呼び出しているのでしょうか？ 呼び出している%22 へ代入される変数をさかのぼっていくと、途中キャストしたり色々していますが、

%22⇒%21(%20の2つめの要素を使用する)⇒%20⇒%19⇒%1⇒%0

というふうにたどれます。%0はmyfoofunc関数の引数です。つまり引数として渡された領域から参照できる関数を呼び出しています。

では、myfoofunc関数はどのような引数で呼ばれるのでしょうか？

呼び出しているmybarfunc関数を抜粋します。

simple_async_await.ll

define hidden swifttailcc void @"$s18simple_async_await9mybarfuncyyYaF"(%swift.context* swiftasync %0) #0 {
entry:
  call void @coro.devirt.trigger(i8* null)
  %1 = bitcast %swift.context* %0 to i8*
  %async.ctx.frameptr = getelementptr inbounds i8, i8* %1, i32 24
  %FramePtr = bitcast i8* %async.ctx.frameptr to %"$s18simple_async_await9mybarfuncyyYaF.Frame"*
  %2 = getelementptr inbounds %"$s18simple_async_await9mybarfuncyyYaF.Frame", %"$s18simple_async_await9mybarfuncyyYaF.Frame"* %FramePtr, i32 0, i32 0
  store %swift.context* %0, %swift.context** %2, align 8
  %3 = load i32, i32* getelementptr inbounds (%swift.async_func_pointer, %swift.async_func_pointer* @"$s18simple_async_await9myfoofuncyyYaFTu", i32 0, i32 1), align 8
  %4 = zext i32 %3 to i64
  %5 = call swiftcc i8* @swift_task_alloc(i64 %4) #1
  %.spill.addr = getelementptr inbounds %"$s18simple_async_await9mybarfuncyyYaF.Frame", %"$s18simple_async_await9mybarfuncyyYaF.Frame"* %FramePtr, i32 0, i32 1
  store i8* %5, i8** %.spill.addr, align 8
  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %5)
  %6 = bitcast i8* %5 to <{ %swift.context*, void (%swift.context*)*, i32 }>*
  %7 = load %swift.context*, %swift.context** %2, align 8
  %8 = getelementptr inbounds <{ %swift.context*, void (%swift.context*)*, i32 }>, <{ %swift.context*, void (%swift.context*)*, i32 }>* %6, i32 0, i32 0
  store %swift.context* %7, %swift.context** %8, align 8
  %9 = getelementptr inbounds <{ %swift.context*, void (%swift.context*)*, i32 }>, <{ %swift.context*, void (%swift.context*)*, i32 }>* %6, i32 0, i32 1
  store void (%swift.context*)* bitcast (void (i8*)* @"$s18simple_async_await9mybarfuncyyYaFTQ0_" to void (%swift.context*)*), void (%swift.context*)** %9, align 8
  %10 = bitcast i8* %5 to %swift.context*
  musttail call swifttailcc void @"$s18simple_async_await9myfoofuncyyYaF"(%swift.context* swiftasync %10) #1
  ret void
}

最後の方にある musttail call　でmyfoofunc関数が呼ばれていることがわかります。引数は、%10ですが、たどっていくと、

%10⇒%5⇒swift_task_allocで確保した領域

となっていることがわかります。swift_task_allocは調べてみてもわからなかったのですが、WWDC21で説明されていたContinuation（継続、タスクを管理するためのオブジェクト）を取得しているのではないかと思われます。

そして、この領域には、resume functionであるmybarfunc(resume)への参照が格納されています。このあたりです。

simple_async_await.ll

store void (%swift.context*)* bitcast (void (i8*)* @"$s18simple_async_await9mybarfuncyyYaFTQ0_" to void (%swift.context*)*), void (%swift.context*)** %9, align 8

%9は、次のようにたどれます。

%9(%6の2つめの要素を使用する)⇒%6⇒%5

%5は先程、swift_task_allocで確保した領域であり、myfoofuncの引数になっていましたね。

myfoofuncは、この引数を受け取って musttail call で呼び出していたので、mybarfunc(resume)が呼ばれることが確認できました。これは上に書いた図の通りの呼び出し関係です。

結論

ということで、async/awaitの実装には、Returned-Continuation Loweringを使用し、関数分割を用いていることがわかりました。また、呼び出し元への制御の移動は、Continuationオブジェクトを使用していることもわかりました。

このように見てみると、awaitをつけて呼び出した場合、制御が戻ってきたときに別のスレッドになる可能性があることは自然なことのように思えます。呼び出し先で別のスレッドに制御が移動して、tail callでジャンプして戻ってくれば、awaitで呼び出した次の行は別のスレッドになります。コード上は次の行ですが、制御上は異なるコンテキストになるわけです。

なお、ローカル変数の保存やデータ生成については調査していません。また、スレッドが生成された場合の制御なども調べたかったのですが、制御権の移動を調べるだけで、LLVM-IRとのにらめっこは限界です😂

最後に

今回の内容についての質問や感想は、TwitterでDMなどいただけると嬉しいです。
調べていてだいぶ面白かったので、興味を持ってくれた方はきっと話が合うでしょう🌈

これ以外では、主にiOSを使った3D処理やAR、ML、音声処理などの作品やサンプル、技術情報を発信しています。
作品ができたらTwitterで発信していきますのでフォローをお願いします🙏

Twitterは作品や記事のリンクを貼っています。
https://twitter.com/jugemjugemjugem

Qiitaは、iOS開発、とくにARや機械学習、グラフィックス処理、音声処理について発信しています。
https://qiita.com/TokyoYoshida

Noteでは、連載記事を書いています。
https://note.com/tokyoyoshida

Zennは機械学習が多めです。
https://zenn.dev/tokyoyoshida