Mercari Engineering Blog

We're the software engineers behind Mercari. Check out our blog to see the tech that powers our marketplace.

Swiftにおける末尾再帰とCompilerによる最適化を探る

初めに

この記事はMercari Advent Calendar 2019の16日目の記事です。

こんにちは、iOSエンジニアの@kagemikuです。MercariのiOSアプリ開発をしています。

突然ですが、皆さんは普段、再帰関数を書いていますか? リスト探索や木探索を行うアルゴリズムを、きれいな再帰関数にバチッと落とし込めた瞬間はなんとも言えない快感がありますよね!キラやば〜っです。 iOSアプリ開発においても、何かしらのデータ処理を実装する際に再帰関数を書くことが度々あるかと思います。 Swiftの登場以降、iOS開発の際に使用する言語はほぼSwiftになったと言っていいでしょう。 そんなSwiftにおいても、当たり前ですが再帰関数を書くことは可能です。 この記事では、そんなSwiftにおける再帰関数と、Compilerによる末尾再帰除去という最適化について探っていきます。

再帰関数、および末尾再帰関数を定義し、その挙動の違いを確認していきます。 その後、Swift Compilerの各Componentからの生成物を比べていき、最適化なし・ありでどのように生成物が異なっているのかを確認していきます。 最終的に、初めに定義した再帰関数を末尾再帰の形にした上で、Compilerによる末尾再帰除去最適化がなされた場合、実行速度が7~8倍高速になっていることを見ていきます。

なお、この記事を通して使用するSwiftのバージョンなどは次の通りです。

$ swift --version
Apple Swift version 5.1.2 (swiftlang-1100.0.278 clang-1100.0.33.9)
Target: x86_64-apple-darwin19.0.0

Swiftにおける再帰関数と末尾再帰

再帰関数

まずは、簡単な再帰呼び出しを行う関数を定義します。ここでは1から、引数として与えられたnまでの総和を計算するsumという関数を定義しています。ついでに、10を引数としてsumを呼び出した結果も表示するようにしています。

func sum(_ n: Int) -> Int {
    if n == 0 {
        return 0
    }

    return n + sum(n - 1)
}

print(sum(10))

このコードをmain.swiftとして保存した上で、Compile、実行すると次のような結果が得られます。

$ swiftc -Onone main.swift && ./main
55

では、引数として与える数を1000000に変更し再度Compile、実行してみます。

$ swiftc -Onone main.swift && ./main
zsh: segmentation fault  ./main

親の顔より見たセグフォ(こちらの結果は環境に依存するかと思います)。 StackにStack frameが乗り切らず、Stack Overflowを引き起こしてしまいました。当然と言えば当然ですね。

Compilerがいい感じにしてくれないかなと期待し、最適化を有効にした上でもう一度実行してみても、セグフォするはずです。

$ swiftc -O main.swift && ./main
zsh: segmentation fault  ./main

困りました。しかし、今日はどうしてもsum関数を再帰関数として定義したい気分です。なんとかしたいのです。 というわけで、この再帰関数を末尾再帰の形で書き直してみます。(理由は後述)

末尾再帰

末尾再帰(関数)とは、「再帰関数において、再帰呼び出しが計算の最後のステップになっている」再帰関数のことです。 ちなみに末尾呼び出しというものもあり、これは「関数において関数呼び出しが計算の最後のステップになっている」関数のことです。 つまり、末尾再帰は末尾呼び出しの再帰関数版ということになります。

早速、末尾再帰の形で書き直した実装を見てみましょう。

func sum(_ n: Int, _ res: Int) -> Int {
    if n == 0 {
        return res
    }

    return sum(n - 1, res + n)
}

 print(sum(1000000, 0))

最初の例とは異なり、最後のステップであるreturnではsum関数の呼び出し、すなわち再帰呼び出しだけになっています。 早速実行してみます。

$ swiftc -Onone main.swift && ./main
zsh: segmentation fault  ./main

変わらずセグフォしてしまいました。Compilerによる最適化を有効にした上でもう一度実行してみます。

swiftc -O main.swift && ./main
500000500000

成功しました🎉 最初の例との違いは次の通りです。

  • 末尾再帰にした
  • Compilerの最適化を有効にした。

この2つの条件が揃ったとき、Compilerによる末尾再帰除去という最適化がなされたため、Stack Overflowが回避されました。キラやば〜っ。

末尾再帰除去最適化とは、末尾再帰の状態になっている再帰関数において、再帰呼び出しをループに置き換える最適化手法です。 ループに置き換えることで、(再帰)関数呼び出しの際にStack frameをStack上へ積むことなく新たな処理を開始できます。

  • 末尾再帰にした
  • Compilerの最適化を有効にした。

この2つの条件が揃ったときに、Compiler内部では何が起きたのでしょうか。これから探っていきます。

Compilerによる最適化

CompilerのPipelineについて

SwiftのCompilerは次のようなPipelineに沿って、最終的なバイナリを吐き出していきます。(Swift コンパイラのアーキテクチャより引用しています)

f:id:kagemiku:20191211193941j:plain
Swift CompilerのPipeline

それぞれのComponentと処理内容、吐き出す生成物は(雑に)次の通りです。lib/と接頭辞のあるものはSwift Compilerのフロントエンドです。バックエンドにはLLVMを採用しています。

Component 処理内容 生成物
lib/Parse 字句解析を行いつつAST構築(型情報や意味解析には関知しない) AST
lib/Sema 意味解析など type-checked AST
lib/SILGen SIL(Swift Intermediate Language)生成 raw SIL
lib/SILOptimizer raw SILのデータフロー解析など(必要に応じて最適化) canonical SIL
lib/IRGen LLVM IRの生成 LLVM IR
LLVM バイナリ生成(必要に応じて最適化) バイナリ(.oオブジェクトファイル)

SwiftのCompilerはOSSとして公開されているので、興味のある方はぜひチェックしてください。

github.com

また、Swift Compilerのアーキテクチャについては、こちらの記事がとても分かりやすいです。

qiita.com

末尾再帰にしたとして、最適化を無効にしたときと有効にしたときで実行結果が変わってくるのであれば、Compilerの生成物も当然異なっているはずです。 それでは、最適化を無効にしたときと有効にしたときで、上記のどのタイミング(Component)でどのような生成物の違いがあるのか、見ていきます。

各Componentにおける差分

次の内容をsum.swiftとして保存し、生成物の差分を実際に見ていきます。

func sum(_ n: Int, _ res: Int) -> Int {
    if n == 0 {
        return res
    }

    return sum(n - 1, res + n)
}

雑に予想できることなのですが、lib/Parseからlib/SILGenまでは生成物に違いはありません。

  • lib/Parse
$ swiftc -dump-parse -Onone sum.swift > onone.txt
$ swiftc -dump-parse -O sum.swift > o.txt
$ diff onone.txt o.txt
$
  • lib/Sema
$ swiftc -dump-ast -Onone sum.swift > onone.txt
$ swiftc -dump-ast -O sum.swift > o.txt
$ diff onone.txt o.txt
$
  • lib/SILGen
$ swiftc -emit-silgen -Onone sum.swift > onone.txt
$ swiftc -emit-silgen -O sum.swift > o.txt
$ diff onone.txt o.txt
$

lib/SILOptimizerでは差分が確認できるのですが、今回の記事に関係のある差分ではありません。

  • lib/SILOptimizer
$ swiftc -emit-sil -Onone sum.swift > onone.txt
$ swiftc -emit-sil -O sum.swift > o.txt
$ diff onone.txt o.txt
19,20c19,20
< // %0                                             // users: %5, %12, %20, %2
< // %1                                             // users: %19, %10, %3
---
> // %0                                             // users: %5, %2
> // %1                                             // users: %16, %8, %3
.
.
.
.
# 省略

lib/IRGenでは、ようやく気になる差分が確認できました。最適化無効、有効それぞれにおいて、sum関数に関する部分だけを掲載します(なお、一部のdebugコードは省略しています)。

  • lib/IRGen (最適化なし)
$ swiftc -emit-sil -Onone sum.swift
.
.
define hidden swiftcc i64 @"$s3sumAAyS2i_SitF"(i64, i64) #0 {
entry:
  .
  .
  %4 = icmp eq i64 %0, 0
  br i1 %4, label %5, label %6

; <label>:5:                                      ; preds = %entry
  br label %16

; <label>:6:                                      ; preds = %entry
  %7 = call { i64, i1 } @llvm.ssub.with.overflow.i64(i64 %0, i64 1)
  %8 = extractvalue { i64, i1 } %7, 0
  %9 = extractvalue { i64, i1 } %7, 1
  br i1 %9, label %18, label %10

; <label>:10:                                     ; preds = %6
  %11 = call { i64, i1 } @llvm.sadd.with.overflow.i64(i64 %1, i64 %0)
  %12 = extractvalue { i64, i1 } %11, 0
  %13 = extractvalue { i64, i1 } %11, 1
  br i1 %13, label %19, label %14

; <label>:14:                                     ; preds = %10
  %15 = call swiftcc i64 @"$s3sumAAyS2i_SitF"(i64 %8, i64 %12)
  br label %16

; <label>:16:                                     ; preds = %5, %14
  %17 = phi i64 [ %15, %14 ], [ %1, %5 ]
  ret i64 %17

; <label>:18:                                     ; preds = %6
  call void @llvm.trap()
  unreachable

; <label>:19:                                     ; preds = %10
  call void @llvm.trap()
  unreachable
}
.
.
  • lib/IRGen (最適化あり)
$ swiftc -emit-sil -O sum.swift
.
.
; Function Attrs: nounwind
define hidden swiftcc i64 @"$s3sumAAyS2i_SitF"(i64, i64) local_unnamed_addr #1 {
entry:
  %2 = icmp eq i64 %0, 0
  br i1 %2, label %tailrecurse._crit_edge, label %.lr.ph

.lr.ph:                                           ; preds = %entry, %tailrecurse
  %.tr15 = phi i64 [ %9, %tailrecurse ], [ %1, %entry ]
  %.tr4 = phi i64 [ %4, %tailrecurse ], [ %0, %entry ]
  %3 = tail call { i64, i1 } @llvm.ssub.with.overflow.i64(i64 %.tr4, i64 1)
  %4 = extractvalue { i64, i1 } %3, 0
  %5 = extractvalue { i64, i1 } %3, 1
  br i1 %5, label %11, label %6

; <label>:6:                                      ; preds = %.lr.ph
  %7 = tail call { i64, i1 } @llvm.sadd.with.overflow.i64(i64 %.tr15, i64 %.tr4)
  %8 = extractvalue { i64, i1 } %7, 1
  br i1 %8, label %12, label %tailrecurse

tailrecurse:                                      ; preds = %6
  %9 = extractvalue { i64, i1 } %7, 0
  %10 = icmp eq i64 %4, 0
  br i1 %10, label %tailrecurse._crit_edge, label %.lr.ph

tailrecurse._crit_edge:                           ; preds = %tailrecurse, %entry
  %.tr1.lcssa = phi i64 [ %1, %entry ], [ %9, %tailrecurse ]
  ret i64 %.tr1.lcssa

; <label>:11:                                     ; preds = %.lr.ph
  tail call void asm sideeffect "", "n"(i32 0) #4
  tail call void @llvm.trap()
  unreachable

; <label>:12:                                     ; preds = %6
  tail call void asm sideeffect "", "n"(i32 1) #4
  tail call void @llvm.trap()
  unreachable
}
.
.

それぞれにおいて、sum関数の内容を再度実行しようとしている部分をさらに抜き出します。 最適化なしの方では、単にcall命令でsum関数を呼び出していました。

; <label>:14:                                     ; preds = %10
  %15 = call swiftcc i64 @"$s3sumAAyS2i_SitF"(i64 %8, i64 %12)
  br label %16

他方、最適化なしの方では、br命令による単なるラベルジャンプになっています(%4、すなわちn0ならretを行うtailrecurse._crit_edgeブロックへジャンプ、そうでないなら再び.lr.phブロックへジャンプ)。

tailrecurse:                                      ; preds = %6
  %9 = extractvalue { i64, i1 } %7, 0
  %10 = icmp eq i64 %4, 0
  br i1 %10, label %tailrecurse._crit_edge, label %.lr.ph

ラベルジャンプ、すなわちループ処理に置き換わっていますね。

ここまできたらLLVMにより吐き出されるAssemblyも確認しておきましょう(同じく、sum関数に関する部分だけを掲載します)。

  • LLVM (最適化なし)
$ swiftc -emit-assembly -Onone sum.swift
.
.
_$s3sumAAyS2i_SitF:
    .cfi_startproc
    pushq   %rbp
    .cfi_def_cfa_offset 16
    .cfi_offset %rbp, -16
    movq    %rsp, %rbp
    .cfi_def_cfa_register %rbp
    subq    $96, %rsp
    xorl    %eax, %eax
    leaq    -8(%rbp), %rcx
    movq    %rdi, -24(%rbp)
    movq    %rcx, %rdi
    movq    %rsi, -32(%rbp)
    movl    %eax, %esi
    movl    $8, %ecx
    movq    %rcx, %rdx
    movl    %eax, -36(%rbp)
    movq    %rcx, -48(%rbp)
    callq   _memset
    leaq    -16(%rbp), %rcx
    movq    %rcx, %rdi
    movl    -36(%rbp), %esi
    movq    -48(%rbp), %rdx
    callq   _memset
    movq    -24(%rbp), %rcx
    movq    %rcx, -8(%rbp)
    movq    -32(%rbp), %rdx
    movq    %rdx, -16(%rbp)
    cmpq    $0, %rcx
    jne LBB1_2
    movq    -32(%rbp), %rax
    movq    %rax, -56(%rbp)
    jmp LBB1_5
LBB1_2:
    movq    -24(%rbp), %rax
    decq    %rax
    seto    %cl
    movq    %rax, -64(%rbp)
    movb    %cl, -65(%rbp)
    jo  LBB1_6
    movq    -32(%rbp), %rax
    movq    -24(%rbp), %rcx
    addq    %rcx, %rax
    seto    %dl
    movq    %rax, -80(%rbp)
    movb    %dl, -81(%rbp)
    jo  LBB1_7
    movq    -64(%rbp), %rdi
    movq    -80(%rbp), %rsi
    callq   _$s3sumAAyS2i_SitF
    movq    %rax, -56(%rbp)
LBB1_5:
    movq    -56(%rbp), %rax
    addq    $96, %rsp
    popq    %rbp
    retq
LBB1_6:
    ud2
LBB1_7:
    ud2
    .cfi_endproc
.
.
  • LLVM (最適化あり)
$ swiftc -emit-assembly -O sum.swift
.
.
_$s3sumAAyS2i_SitF:
    pushq   %rbp
    movq    %rsp, %rbp
    movq    %rsi, %rax
    testq   %rdi, %rdi
    je  LBB1_5
    movq    %rdi, %rcx
    .p2align    4, 0x90
LBB1_2:
    decq    %rcx
    jo  LBB1_6
    addq    %rdi, %rax
    jo  LBB1_7
    movq    %rcx, %rdi
    testq   %rcx, %rcx
    jne LBB1_2
LBB1_5:
    popq    %rbp
    retq
LBB1_6:
    ## InlineAsm Start
    ## InlineAsm End
    ud2
LBB1_7:
    ## InlineAsm Start
    ## InlineAsm End
    ud2
.
.

再び、それぞれにおいて、sum関数の内容を再度実行しようとしている部分をさらに抜き出します。 最適化なしの方では、単にcall命令でsum関数を呼び出していました。 その際n - 1res + nはレジスタに格納してsum関数へ渡していますが、呼び出し時点での各レジスタの状態などは一部Stackに残っています。 これでは、関数呼び出しをすればするほど、Stackの枯渇に近づいていきますね。

    movq    -64(%rbp), %rdi
    movq    -80(%rbp), %rsi
    callq   _$s3sumAAyS2i_SitF

他方、最適化ありの方では、きちんとループに置き換わっていますね!(jne命令の部分です) キラやば〜っです。 Stackも消費しておらず、レジスタの更新だけで済んでおり、省エネです。末尾再帰除去はエコ。

LBB1_2:
    decq    %rcx
    jo  LBB1_6
    addq    %rdi, %rax
    jo  LBB1_7
    movq    %rcx, %rdi
    testq   %rcx, %rcx
    jne LBB1_2

以上のことから、末尾再帰除去最適化は、少なくともlib/IRGen以降に行われていることがわかります。 より厳密に述べると、lib/IRGen自体がLLVM IRの吐き出しロジックを持っているのではなく、LLVMのAPI(IRBuilder.CreateCallなど)を用いてLLVM IRにおけるinstructionに変換した後、LLVMのdumpメソッドによりLLVM IRのテキストフォーマットを出力しています。 その過程で、lib/IRGenはLLVMに対し最適化の依頼をし、LLVMが末尾再帰除去最適化(など)を行っています。 すなわち末尾再帰除去最適化そのものは、Swift Compilerのフロントエンドではなく、バックエンドであるLLVM上で行われています。

Swift CompilerのPipeline図で言うと、赤で囲った部分です。

f:id:kagemiku:20191216115822j:plain
Swift CompilerのPipeline(再掲)

LLVMにおける最適化

ついでなのでLLVMモジュールのコードも眺めてみましょう(AppleはオリジナルのLLVM repositoryをforkしています)。

github.com

その中にTailRecursionElimination.cppというファイルがあります。このファイルの先頭コメント部分を見てみると次のように書いてあります。

This file transforms calls of the current function (self recursion) followed by a return instruction with a branch to the entry of the function, creating a loop.

どうやらこのファイルで末尾再帰除去最適化の処理が行われているようです。今回は実装の詳細は追いませんが、興味のある方はぜひ覗いてみてください。

なおScalaなど、言語によっては特定のアノテーション(@tailrecのような)を付けてあげると、それが末尾再帰の形になっていない場合にCompile Errorを出してくれる機能もあります。 以前に同様の機能をSwift4に入れようとProposalが出されていたようなのですが、結局Swift4には入りませんでした。該当のIssueは現在Closeされてしまっています。

github.com

github.com

パフォーマンス測定

最適化によってループへの置き換えがなされているということは、関数呼び出しに係るオーバヘッドもなくなり、いくらか高速になっているはずです。パフォーマンス測定もしてしまいましょう。

次のような内容のファイルを作成します。

  • sum.swift
func sum(_ n: Int, _ res: Int) -> Int {
    if n == 0 {
        return res
    }

    return sum(n - 1, res + n)
}
  • main.swift
SumTests.defaultTestSuite.run()
  • benchmark.swift
import XCTest

class SumTests: XCTestCase {

    func testSum100() {
        measure {
            _ = sum(100, 0)
        }
    }

    func testSum100000() {
        measure {
            _ = sum(100000, 0)
        }
    }

}

これらを最適化なし、ありの両方でCompile、実行し、パフォーマンスを見てみます。

$ swiftc -Onone -F /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/Library/Frameworks -Xlinker -rpath -Xlinker /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/Library/Frameworks -lswiftCore benchmark.swift sum.swift main.swift -o benchmark
$ ./benchmark
Test Suite 'SumTests' started at 2019-12-09 00:52:28.258
Test Case '-[benchmark.SumTests testSum100]' started.
<unknown>:0: Test Case '-[benchmark.SumTests testSum100]' measured [Time, seconds] average: 0.000, relative standard deviation: 119.699%, values: [0.000043, 0.000009, 0.000006, 0.000006, 0.000005, 0.000005, 0.000005, 0.000005, 0.000005, 0.000005], performanceMetricID:com.apple.XCTPerformanceMetric_WallClockTime, baselineName: "", baselineAverage: , maxPercentRegression: 10.000%, maxPercentRelativeStandardDeviation: 10.000%, maxRegression: 0.100, maxStandardDeviation: 0.100
Test Case '-[benchmark.SumTests testSum100]' passed (0.347 seconds).
Test Case '-[benchmark.SumTests testSum10000]' started.
<unknown>:0: Test Case '-[benchmark.SumTests testSum10000]' measured [Time, seconds] average: 0.000, relative standard deviation: 94.556%, values: [0.000482, 0.000093, 0.000086, 0.000085, 0.000086, 0.000085, 0.000085, 0.000085, 0.000085, 0.000085], performanceMetricID:com.apple.XCTPerformanceMetric_WallClockTime, baselineName: "", baselineAverage: , maxPercentRegression: 10.000%, maxPercentRelativeStandardDeviation: 10.000%, maxRegression: 0.100, maxStandardDeviation: 0.100
Test Case '-[benchmark.SumTests testSum10000]' passed (0.254 seconds).
Test Suite 'SumTests' passed at 2019-12-09 00:52:28.861.
     Executed 2 tests, with 0 failures (0 unexpected) in 0.601 (0.603) seconds
$ swiftc -O -F /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/Library/Frameworks -Xlinker -rpath -Xlinker /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/Library/Frameworks -lswiftCore benchmark.swift sum.swift main.swift -o benchmark
$ ./benchmark
Test Suite 'SumTests' started at 2019-12-09 00:52:39.138
Test Case '-[benchmark.SumTests testSum100]' started.
<unknown>:0: Test Case '-[benchmark.SumTests testSum100]' measured [Time, seconds] average: 0.000, relative standard deviation: 127.815%, values: [0.000040, 0.000008, 0.000006, 0.000005, 0.000004, 0.000004, 0.000004, 0.000004, 0.000004, 0.000004], performanceMetricID:com.apple.XCTPerformanceMetric_WallClockTime, baselineName: "", baselineAverage: , maxPercentRegression: 10.000%, maxPercentRelativeStandardDeviation: 10.000%, maxRegression: 0.100, maxStandardDeviation: 0.100
Test Case '-[benchmark.SumTests testSum100]' passed (0.348 seconds).
Test Case '-[benchmark.SumTests testSum10000]' started.
<unknown>:0: Test Case '-[benchmark.SumTests testSum10000]' measured [Time, seconds] average: 0.000, relative standard deviation: 86.648%, values: [0.000057, 0.000024, 0.000015, 0.000012, 0.000009, 0.000009, 0.000009, 0.000009, 0.000009, 0.000009], performanceMetricID:com.apple.XCTPerformanceMetric_WallClockTime, baselineName: "", baselineAverage: , maxPercentRegression: 10.000%, maxPercentRelativeStandardDeviation: 10.000%, maxRegression: 0.100, maxStandardDeviation: 0.100
Test Case '-[benchmark.SumTests testSum10000]' passed (0.254 seconds).
Test Suite 'SumTests' passed at 2019-12-09 00:52:39.742.
     Executed 2 tests, with 0 failures (0 unexpected) in 0.602 (0.603) seconds

ものすごく見にくいですね。気合で読んでください。一応、ざっと平均値をまとめます。

最適化 測定結果平均(testSum100) 測定結果平均(testSum100000)
なし (-Onone) 0.0000094 0.0001257
あり (-O) 0.0000083 0.0000162

sum(100000, 0)の場合、最適化ありの方は7 ~ 8倍ほど早くなっていますね。キラやば〜っです。

終わりに

今回の記事は「そういえばSwiftでは末尾再帰除去は行われるんだろうか」という率直な疑問をきっかけにして生まれました。 Assemblyまで吐かせてしまえば間違いのない、確かな答えが得られるだろうという発想の元、愚直に丁寧にCompilerの生成物を調査してました。 その過程はとても楽しく、(おそらく)これかもSwiftを書いていくiOSエンジニアとして、ほんの少しばかりSwiftの気持ちが分かるようになったかなと思います。 かなり長くなってしまいましたがここまで読んで頂きありがとうございました!

明日17日目の執筆担当は@lightnet328さんです。引き続きMercari Advent Calendar 2019をお楽しみください!