編譯快取

從 Android 10 開始,神經網路 API (NNAPI) 提供了支援編譯工件快取的功能,從而減少了應用程式啟動時的編譯時間。使用此快取功能,驅動程式不需要管理或清理快取的檔案。這是一項可選功能,可以使用 NN HAL 1.2 來實現。有關此函數的更多信息,請參閱ANeuralNetworksCompilation_setCaching

該驅動程式還可以實現獨立於 NNAPI 的編譯快取。無論是否使用 NNAPI NDK 和 HAL 快取功能,都可以實現這一點。 AOSP 提供了一個低階實用程式庫(快取引擎)。有關更多信息,請參閱實現緩存引擎

工作流程概述

本節介紹實作編譯快取功能的一般工作流程。

提供的快取資訊和緩存命中

  1. 該應用程式傳遞一個快取目錄和一個模型特有的校驗和。
  2. NNAPI 運行時會根據校驗和、執行首選項和分區結果查找快取檔案並找到檔案。
  3. NNAPI 開啟快取檔案並使用prepareModelFromCache將句柄傳遞給驅動程式。
  4. 驅動程式直接從快取檔案準備模型並返回準備好的模型。

快取資訊提供和快取未命中

  1. 該應用程式傳遞模型特有的校驗和和快取目錄。
  2. NNAPI 運行時會根據校驗和、執行首選項和分區結果查找快取文件,但找不到快取文件。
  3. NNAPI 根據校驗和、執行首選項和分區創建空緩存文件,打開緩存文件,並使用prepareModel_1_2將句柄和模型傳遞給驅動程式。
  4. 驅動程式編譯模型,將快取資訊寫入快取文件,並返回準備好的模型。

未提供快取信息

  1. 該應用程式調用編譯而不提供任何快取資訊。
  2. 該應用程式不會傳遞任何與快取相關的內容。
  3. NNAPI 運行時使用prepareModel_1_2將模型傳遞給驅動程式。
  4. 驅動程式編譯模型並返回準備好的模型。

快取資訊

提供給驅動程式的快取資訊由令牌和快取檔案句柄組成。

代幣

令牌是長度為Constant::BYTE_SIZE_OF_CACHE_TOKEN快取令牌,用於標識準備好的模型。使用prepareModel_1_2儲存快取檔案並使用prepareModelFromCache檢索準備好的模型時,會提供相同的標記。驅動程式的客戶端應該選擇衝突率低的令牌。驅動程式無法偵測到令牌衝突。衝突會導致執行失敗或成功執行但產生不正確的輸出值。

快取檔案句柄(兩種類型的快取檔案)

快取檔案有兩種類型:資料快取模型快取

  • 資料快取:用於快取常數數據,包括預處理和轉換的張量緩衝區。資料快取的修改不應導致比在執行時產生錯誤輸出值更糟糕的影響。
  • 模型快取:用於快取安全敏感數據,例如以裝置本機二進位格式編譯的可執行機器碼。模型快取的修改可能會影響驅動程式的執行行為,惡意用戶端可能會利用它來超出授予的權限執行。因此,驅動程式必須在從快取準備模型之前檢查模型快取是否已損壞。有關詳細信息,請參閱安全性

驅動程式必須決定如何在兩種類型的快取檔案之間分配快取信息,並使用getNumberOfCacheFilesNeeded報告每種類型需要多少個快取檔案。

NNAPI 運行時始終開啟具有讀寫權限的快取檔案句柄。

安全

在編譯快取中,模型快取可能包含安全敏感數據,例如裝置本機二進位格式的已編譯可執行機器碼。如果保護不當,模型快取的修改可能會影響驅動程式的執行行為。由於快取內容儲存在app目錄中,因此客戶端可以修改快取檔案。有錯誤的客戶端可能會意外損壞緩存,而惡意用戶端可能會故意利用它在裝置上執行未經驗證的程式碼。根據設備的特性,這可能是安全問題。因此,驅動程式必須能夠在從快取準備模型之前檢測潛在的模型快取損壞。

實現此目的的一種方法是驅動程式維護從令牌到模型快取的加密雜湊的映射。當將編譯保存到快取時,驅動程式可以儲存其模型快取的令牌和雜湊。當從快取檢索編譯時,驅動程式使用記錄的令牌和雜湊對檢查模型快取的新雜湊。此映射應在系統重新啟動後保持不變。驅動程式可以使用Android 金鑰庫服務framework/ml/nn/driver/cache中的公用程式庫或任何其他適當的機制來實作映射管理器。驅動程式更新後,應重新初始化此映射管理器,以防止從早期版本準備快取檔案。

為了防止檢查時間到使用時間(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_2prepareModelFromCache期間複製檔案句柄,並在稍後讀取/更新快取內容。
  • 在普通編譯呼叫之外實作檔案鎖定邏輯,以防止寫入與讀取或其他寫入同時發生。

實施緩存引擎

除了 NN HAL 1.2 編譯快取介面之外,您還可以在frameworks/ml/nn/driver/cache目錄中找到快取實用程式庫。 nnCache子目錄包含驅動程式的持久性儲存程式碼,用於在不使用 NNAPI 快取功能的情況下實作編譯快取。這種形式的編譯快取可以使用任何版本的 NN HAL 來實現。如果驅動程式選擇實現與 HAL 介面斷開連接的緩存,則驅動程式負責在不再需要緩存的工件時釋放它們。