コンパイルのキャッシュ

Android 10 以降、Neural Networks API（NNAPI）は、コンパイルのアーティファクトのキャッシュをサポートする機能を提供し、アプリ起動時にコンパイルにかかる時間を短縮します。このキャッシュ機能を使用すると、ドライバは、キャッシュされたファイルを管理またはクリーンアップする必要がありません。これはオプション機能で、NN HAL 1.2 で実装できます。この機能の詳細については、ANeuralNetworksCompilation_setCaching をご覧ください。

ドライバは、NNAPI とは独立したコンパイルのキャッシュの実装も行えます。その際、NNAPI NDK と HAL のキャッシュ機能が使用されているかどうかは関係ありません。AOSP は、低レベルのユーティリティライブラリ（キャッシュエンジン）を提供しています。詳細については、キャッシュエンジンの実装をご覧ください。

ワークフローの概要

このセクションでは、コンパイルのキャッシュ機能を実装した一般的なワークフローについて説明します。

提供されたキャッシュ情報とキャッシュヒット

アプリは、キャッシュディレクトリと、モデルに固有のチェックサムを渡します。
NNAPI ランタイムは、チェックサム、実行設定、パーティショニングの結果に基づいてキャッシュファイルを検索、検出します。
NNAPI はキャッシュファイルを開き、ハンドルを prepareModelFromCache でドライバに渡します。
ドライバはキャッシュファイルからモデルを直接準備し、完了したモデルを返します。

提供されたキャッシュ情報とキャッシュミス

アプリは、モデルに固有のチェックサムと、キャッシュディレクトリを渡します。
NNAPI ランタイムは、チェックサム、実行設定、パーティショニングの結果に基づいてキャッシュするファイルを検索しますが、キャッシュファイルは検出しません。
NNAPI は、チェックサム、実行設定、パーティショニングに基づいて空のキャッシュファイルを作成して開き、ハンドルとモデルを prepareModel_1_2 でドライバに渡します。
ドライバはモデルをコンパイルし、キャッシュ情報をキャッシュファイルに書き込み、準備が完了したモデルを返します。

提供されないキャッシュ情報

アプリはキャッシュ情報を提供せずにコンパイルを呼び出します。
アプリは、キャッシュに関連するものは何も渡しません。
NNAPI ランタイムはモデルを prepareModel_1_2 でドライバに渡します。
ドライバはモデルをコンパイルし、準備が完了したモデルを返します。

キャッシュ情報

ドライバに提供されるキャッシュ情報は、トークンとキャッシュファイルのハンドルで構成されます。

トークン

トークンは、長さが Constant::BYTE_SIZE_OF_CACHE_TOKEN のキャッシュトークンで、準備済みのモデルを識別します。prepareModel_1_2 でキャッシュファイルを保存して、準備済みのモデルを prepareModelFromCache で取得すると、同じトークンが提供されます。ドライバのクライアントは、衝突率の低いトークンを選択する必要があります。ドライバはトークンの衝突を検出できません。衝突すると、実行に失敗するか、実行されても正しい出力値が生成されません。

キャッシュファイルのハンドル（2 種類のキャッシュファイル）

キャッシュファイルには、データキャッシュとモデルキャッシュの 2 種類があります。

データキャッシュ: 前処理して変換したテンソルのバッファを含む定数データのキャッシュに使用します。データキャッシュを変更しても、実行時に不適切な出力値が生成されるほど悪い影響が出ることはないでしょう。
モデルキャッシュ: コンパイルされた実行可能マシンコードなど、機密データをデバイスのネイティブバイナリ形式でキャッシュするために使用します。モデルキャッシュを変更すると、ドライバの実行動作に影響が生じ、悪意のあるクライアントに許可された権限を越えて利用される可能性があります。したがって、キャッシュからモデルを準備する前に、ドライバはモデルキャッシュが破損していないかどうかを確認する必要があります。詳細については、セキュリティをご覧ください。

ドライバは、2 種類のキャッシュファイル間のキャッシュ情報の分配方法を決定し、getNumberOfCacheFilesNeeded で種類ごとに必要なキャッシュファイルの数を報告します。

NNAPI ランタイムは、読み取りと書き込みの両方の権限を持つキャッシュファイルのハンドルを常に開きます。

セキュリティ

コンパイルキャッシュでは、モデルキャッシュに、コンパイルされた実行可能マシンコードなどの機密データがデバイスのネイティブバイナリ形式で含まれる場合があります。適切に保護されていない場合、モデルキャッシュの変更がドライバの実行動作に影響する可能性があります。キャッシュの内容はアプリのディレクトリに保存されるため、クライアントがキャッシュファイルを変更できます。バグのあるクライアントが誤ってキャッシュを破損し、悪意のあるクライアントが意図的にこのコードを利用してデバイスで未確認のコードを実行する可能性があります。デバイスの特性によっては、セキュリティ上の問題になることがあります。したがって、キャッシュからモデルを準備する前に、ドライバは破損の可能性があるモデルキャッシュを検出できる必要があります。

これを行う方法の 1 つは、ドライバがトークンからモデルキャッシュの暗号ハッシュへのマップを維持することです。ドライバは、キャッシュにコンパイルを保存するときに、モデルキャッシュのトークンとハッシュを格納できます。また、キャッシュからコンパイルを取得するときに、記録されたトークンとハッシュのペアでモデルキャッシュの新しいハッシュを確認します。このマッピングは、永続的で、システムを再起動しても変わりません。ドライバは、Android キーストアサービス、framework/ml/nn/driver/cache のユーティリティライブラリその他の適切なメカニズムを使用して、マッピングマネージャーを実装できます。ドライバの更新時にはこのマッピングマネージャーを再初期化して、以前のバージョンからキャッシュファイルを準備できないようにしてください。

time-of-check to time-of-use（TOCTOU）攻撃を防止するために、ドライバはファイルに保存する前に記録されたハッシュを計算し、ファイルの内容を内部バッファにコピーした後に新しいハッシュを計算する必要があります。

以下のサンプルコードは、このロジックを実装する方法を示しています。

bool saveToCache(const sp<V1_2::IPreparedModel> preparedModel,
                 const hidl_vec<hidl_handle>& modelFds, const hidl_vec<hidl_handle>& dataFds,
                 const HidlToken& token) {
    // Serialize the prepared model to internal buffers.
    auto buffers = serialize(preparedModel);

    // This implementation detail is important: the cache hash must be computed from internal
    // buffers instead of cache files to prevent time-of-check to time-of-use (TOCTOU) attacks.
    auto hash = computeHash(buffers);

    // Store the {token, hash} pair to a mapping manager that is persistent across reboots.
    CacheManager::get()->store(token, hash);

    // Write the cache contents from internal buffers to cache files.
    return writeToFds(buffers, modelFds, dataFds);
}

sp<V1_2::IPreparedModel> prepareFromCache(const hidl_vec<hidl_handle>& modelFds,
                                          const hidl_vec<hidl_handle>& dataFds,
                                          const HidlToken& token) {
    // Copy the cache contents from cache files to internal buffers.
    auto buffers = readFromFds(modelFds, dataFds);

    // This implementation detail is important: the cache hash must be computed from internal
    // buffers instead of cache files to prevent time-of-check to time-of-use (TOCTOU) attacks.
    auto hash = computeHash(buffers);

    // Validate the {token, hash} pair by a mapping manager that is persistent across reboots.
    if (CacheManager::get()->validate(token, hash)) {
        // Retrieve the prepared model from internal buffers.
        return deserialize<V1_2::IPreparedModel>(buffers);
    } else {
        return nullptr;
    }
}

高度なユースケース

特定の高度なユースケースでは、コンパイルの呼び出し後にドライバがキャッシュのコンテンツにアクセス（読み取りまたは書き込み）することを要求します。たとえば、次のような場合があります。

実行時コンパイル: コンパイルは最初の実行まで遅延します。
多段階コンパイル: 高速コンパイルが最初に実行され、オプションの最適化コンパイルは使用頻度に応じて後で実行されます。

コンパイルの呼び出し後にキャッシュのコンテンツにアクセス（読み取りまたは書き込み）するために、ドライバは以下を行います。

prepareModel_1_2 または prepareModelFromCache の呼び出し中にファイルのハンドルを複製し、キャッシュコンテンツを後で読み込んで読み取り / 書き込みを行います。
通常のコンパイルの呼び出しの外側にファイルロックのロジックを実装し、同時書き込みや読み取りと書き込みの同時発生を防ぎます。

キャッシュエンジンの実装

NN HAL 1.2 コンパイルキャッシュのインターフェースに加えて、frameworks/ml/nn/driver/cache ディレクトリにはキャッシュユーティリティライブラリもあります。nnCache サブディレクトリは、ドライバの永続ストレージコードが含まれており、NNAPI キャッシュ機能を使用せずにコンパイルキャッシュを実装します。この形式のコンパイルキャッシュは、NN HAL の任意のバージョンで実装できます。HAL インターフェースから切り離されたキャッシュをドライバに実装することを選択した場合、ドライバは、キャッシュされたアーティファクトが不要になったときに解放する必要があります。