R8 が Android での Kotlin コルーチンの速度を 2 倍にした方法

AGP 9.2.0 以降、R8 はほとんどの Atomic*FieldUpdater 呼び出しを、一般的なオペレーションで 2 ～ 4 倍のパフォーマンスを発揮する Unsafe バリアントに最適化します。これは、kotlinx.coroutines のアトミックを実装する kotlinx.atomicfu ライブラリに特に大きな影響を与え、コルーチンの起動とキャンセルを最大 2 倍高速化します。この特典を得るには、AGP を 9.2.0 以降に更新してください。

Android アプリの大部分で Kotlin がメインの言語として採用されているため、kotlinx.coroutines は非同期プログラミングの事実上の標準となっています。このライブラリは、Kotlin ネイティブの同時実行フローを管理するための、適切に設計された構造化された方法を提供します。Jetpack Compose も例外ではなく、ポインタ イベント、アニメーション、その他のインタラクションの管理にコルーチンを採用しています。このドキュメントの執筆時点では、Compose のほとんどの同時実行 API は内部で suspend 関数を呼び出し、更新を処理するためにコルーチンを起動またはキャンセルしています。

Compose チームがパフォーマンスの調査を開始したところ、コルーチンがコンポジション外で発生する多くのオペレーションのボトルネックになっていることがわかりました。たとえば、Modifier.clickable の作成と更新に費やされた時間の 80% は、InteractionSource の更新を処理する内部コルーチンの起動とキャンセルに費やされました。これらの観察結果に基づいて、初期のパフォーマンス作業の多くは、デフォルト パスからコルーチンを削除し、初期化を必要なときまで遅らせることに重点が置かれました。

コルーチンのコスト

Android で関数の内部動作を分析する最も簡単な方法は、Android ランタイム（ART）メソッド トレースをキャプチャすることです。ART メソッド トレースは、アプリの実行フローを記録するツールです。どのメソッドがどの順序で呼び出され、各メソッドにどのくらいの時間が費やされたかを正確に把握できるため、デベロッパーはパフォーマンスのボトルネックを特定できます。空の LaunchedEffect { } 呼び出しの場合、次のようになります。

上記のメソッド トレースは、次の 3 つの部分に分けることができます。

• 新しいコルーチンの初期化
• コルーチンの開始
• コルーチンの完了（すぐに終了するため）

LaunchedEffect のキャンセルは、CancellationException も作成する点を除いて、通常の完了と似ています。

上記のプロファイルから、すぐに疑わしいとわかるのは、java.util.concurrent.AtomicReferenceFieldUpdater（j… ラベルの付いた紫または緑のボックス）への頻繁な呼び出しです。各呼び出しは比較的高速ですが、頻度が問題です。複数の呼び出しに分散された無視できないオーバーヘッドは、顕著な回帰につながる可能性があります。通話を拡大すると、時間のほとんどがリフレクション チェックに費やされていることがわかります。

コルーチンは、親子関係にロックフリーのツリー構造を実装し、構造化された同時実行を可能にします。kotlinx.atomicfu ライブラリは、よく知られている JVM プリミティブである AtomicReferenceFieldUpdater を使用して、ロックフリーの原子操作を実装しています。アップデータは、クラス参照とフィールド名を使用して実行時にアトミック オペレーションを実行します。また、フィールドが存在し、アクセス可能であることを確認するために、いくつかのリフレクション セーフティ チェックを実行する必要があります。コルーチンの各オペレーション（開始、一時停止、キャンセル、完了）は少なくとも 1 つのアトミック オペレーションを呼び出すため、処理が遅いとコルーチンはうまく機能しません。

AtomicReferenceFieldUpdater の調査

しかし、先走らないようにしましょう。実際には、AtomicReferenceFieldUpdater は JVM で 10 年以上前から最適化されており、メソッド トレースでは、VM レベルの最適化（ジャストインタイム（JIT）コンパイルまたは事前（AOT）コンパイル）によって完全に削除されるオーバーヘッドがキャプチャされる可能性があります。パフォーマンスを確認するために、kotlinx.atomicfu と java.util.concurrent.atomic のアトミック参照の違いを測定するベンチマークをいくつか作成しましょう。

@RunWith(AndroidJUnit4::class)
class AtomicReferenceBenchmark {
    @get:Rule
    val benchmarkRule = BenchmarkRule()
    
    private val atomicReference = java.util.concurrent.atomic.AtomicReference(false)
    private val atomicRef = kotlinx.atomicfu.atomic<Boolean>(false)
    
    @Test
    fun atomicReference_compareAndSet() {
        benchmarkRule.measureRepeated { 
            atomicReference.compareAndSet(true, false)
            atomicReference.compareAndSet(false, true)
        }
    }

     @Test
    fun atomicRef_compareAndSet() {
        benchmarkRule.measureRepeated {
            atomicRef.compareAndSet(true, false)
            atomicRef.compareAndSet(false, true)
        }
    }

    /* measuring other methods from the method traces above */
}

Google Pixel 5 でこのベンチマークを実行すると（ウォームアップ中に AtomicReferenceFieldUpdater#compareAndSet が JIT コンパイルされることを確認）、Google Pixel 5（API 33）で次の結果が得られます。

 50.7 ns  atomicReference_compareAndSet
135   ns  atomicRef_compareAndSet

測定結果から、kotlinx.atomicfu バージョンが約 2.7 倍遅いことが明らかになりました。これにより、ART が隠れた最適化を行わず、リフレクション アクセス チェックが実行時に実際のオーバーヘッドを追加することが確認されます。

元のメソッド トレースを振り返ると、AtomicReferenceFieldUpdater が実行する意味のある作業は、基盤となるアトミック オペレーションを実際に実行する Unsafe.getObjectVolatile への内部呼び出しのみです。ほとんどの場合、アップデータ イニシャライザは静的であり、周囲のクラスの構造に基づいて常に正しいことが証明できます。そのため、AtomicReferenceFieldUpdater の使用状況のほとんどを静的に分析し、コンパイル時に内部の Unsafe バリアントに置き換えることができます。Android ビルド ツールチェーンには、まさにそれを実行できる独自の最適化コンパイラがあります。

R8 による最適化

Atomic*FieldUpdater クラスは、微妙な動的でリフレクション ベースの使用をサポートしますが、静的に明らかなパターンでよく使用されます。これにより、ベースライン パフォーマンスが遅いことと、最適化の必要性の両方を説明できます。R8 はフルプログラム最適化コンパイラであり、より単純なパターンを見抜いて、リフレクションの安全チェックのオーバーヘッドを削減するのに適しています。R8 は Java コンパイラまたは Kotlin コンパイラの後に JVM バイトコードを受け取りますが、読みやすさを考慮して、これらの例は Java 構文で示されています。このため、AtomicReferenceFieldUpdater には型引数がありません。

class Example {
    volatile String data = "";
    static final AtomicReferenceFieldUpdater updater =
        AtomicReferenceFieldUpdater.newUpdater(Example.class, String.class, "data");

    void example() {
        // ...
        updater.compareAndSet(this, "", "new");
        // ...
    }
}

基本例では、ホルダー、型、フィールド名に単純な定数引数を使用して volatile フィールドにアクセスする static final アップデータを作成します。使用されるリフレクションは完全に透明です。このアップデータが有効なフィールドを参照しており、アップデータ作成サイトがそのフィールドへの有効なアクセス権を持っていることが明確にわかります。

本質的に、Atomic*FieldUpdater はフィールド オフセットと Unsafe の呼び出しをラップするものです。最適化の最良のシナリオは、更新フィールドをオフセット フィールドに置き換え、更新呼び出しを Unsafe の呼び出しに置き換えることです。

Atomic*FieldUpdater の最適化

最適化は、インストルメンテーション、置換、クリーンアップの 3 つの部分で実装されます。

インストゥルメンテーション

最初の手順は、Unsafe 呼び出しによる直接アクセスを容易にするため、オフセット フィールドを更新プログラム フィールドとともに導入することです。

static final long updater$offset =
    SyntheticUnsafe.UNSAFE.objectFieldOffset(Example.class.getDeclaredField("data"))

このフィールドにはリフレクションを介してアクセスし、Unsafe を使用してクラスのフィールド オフセットを抽出します。このコードは、リフレクション検証を無視した場合の Atomic*FieldUpdater の内部処理を表しています。代わりに、アップデータのホルダー型と揮発性フィールドのフィールド型がコンパイラで静的に追跡されます。

元のフィールドとその初期化はそのまま残されます。最適化プロセスでは、楽観的に使用を促進して最適化し、後でクリーンアップします。これは実装の簡単なアプローチですが、アップデータ フィールドの部分的な最適化も可能です。一部の用途はそのまま残し、他の用途は最適化します。

交換

コンパイラのこの時点では、適切な同時実行結合ポイントの後に、計測されたアップデータ フィールドのリストがあります。つまり、いくつかの条件に基づいて各呼び出しサイトを個別に最適化できます。次の呼び出し例を考えてみましょう。

updater.compareAndSet(holder, expectedValue, newValue);

Atomic*FieldUpdater に必要な条件は次のとおりです。

• updater はインストルメント化されたフィールドから取得されますか？つまり、静的解析でオブジェクトの値を追跡して、計測された更新プログラムのフィールド読み取りに戻すことはできますか？
• holder は、元々定義されたホルダー型と同じクラスか、そのサブクラスですか？
• newValue は、元々定義されたフィールド型と同じクラスですか、またはサブクラスですか？

すべての条件が満たされると、呼び出しはリフレクション チェックなしで Unsafe への呼び出しに置き換えられます。

SyntheticUnsafe.UNSAFE.compareAndSwapObject(holder, Example.updater$offset, expectedValue, newValue)

この新しい呼び出しは高速でシンプルですが、updater と holder の null 値の処理に関して元の呼び出しとは異なります。静的に除外されない限り、両方に null チェックが挿入されます。

クリーンアップ

この時点で、保持クラスには元の更新フィールドと新しいオフセット フィールドがあり、どちらかを使用する可能性のある呼び出しサイトがあります。呼び出しサイトが最適化されていない場合はオフセット フィールドを削除し、呼び出しサイトがすべて最適化されている場合は更新プログラム フィールドを削除する必要があります。どちらの場合も、初期化呼び出しも削除する必要があります。未使用のフィールドの削除とデッドコードの削除はコンパイラですでに行われていますが、ここで初期化コードを削除するには、いくつかの追加の処理が必要です。

newUpdater と getDeclaredField の両方の呼び出しは、例外をスローする可能性があるため、副作用が生じる可能性があります（また、API バージョンによって実装が異なるため、実装も不明です）。つまり、一般的な最適化では安全に削除できません。そのため、このクリーンアップでは、例外がないことが静的にわかっている計測フィールドを明示的に考慮する必要がありました。

最終的に、上記のシンプルなアップデータ例は、最適化後に次のようになります。

結果

これらの最適化により、kotlinx.atomicfu と AtomicInt/Long/ReferenceFieldUpdater のほとんどの明示的な使用が、R8 を適用した AtomicReference のパフォーマンスと一致するようになりました。実際、一部のベンチマークではさらに高速です。kotlinx.atomicfu には、atomic インスタンスをフィールドにインライン化できるコンパイラ プラグインがあり、アトミックに更新されたフィールドの作成に必要な割り当てを減らすことができます。

この取り組みの主な恩恵を受けたのは Jetpack Compose でした。Compose ランタイムには、コルーチンのパフォーマンスを非常に厳密に追跡して、パフォーマンスの回帰を早期に検出するマイクロベンチマークが多数あります。ベンチマークが新しいバージョンの R8 に更新されたとき、LaunchedEffect でコルーチンを起動してキャンセルする際に 2 倍の改善が見られました。

それ以外にも、ART チームはこれらの最適化を VM レベルでネイティブに実装しています。アプリが API 36 を対象としていて、最新バージョンの Android で実行されている場合、デバイスですでに同様の方法でコルーチンが最適化されている可能性があります。上記のコルーチン ベンチマークでは、最近のバージョンの ART で JIT を更新した後、パフォーマンスが約 15% 向上しました。

AGP 9.2.0 にアップグレードするか、R8 9.2.0 を直接使用すると、アプリにこの最適化がデフォルトで適用されます。詳しくは、D8 dexer と R8 shrinker をご覧ください。