编译缓存

从 Android 10 开始,Neural Networks 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 运行时打开缓存文件句柄时始终具有读写权限。

安全性

在编译缓存中,模型缓存可能包含对安全敏感的数据,例如采用设备原生二进制文件格式且已编译的可执行机器代码。如果未得到适当保护,修改模型缓存可能会影响驱动程序的执行行为。由于缓存内容存储在应用目录中,因此客户端可以修改缓存文件。有漏洞的客户端可能会意外损坏缓存,而恶意客户端可能会故意利用这一点在设备上执行未经验证的代码。这可能引发安全问题,具体取决于设备特性。因此,驱动程序必须能够检测潜在的模型缓存损坏情况,然后再从缓存准备模型。

执行此操作的方法之一是让驱动程序维护从令牌到模型缓存的加密哈希的映射。将编译保存到缓存时,驱动程序可以存储令牌及其模型缓存的哈希。在从缓存检索编译时,驱动程序会使用已记录的令牌和模型缓存的哈希对检查新哈希。此映射应在系统多次重新启动后保持不变。驱动程序可以使用 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 接口断开连接的缓存,那么,驱动程序负责在不再需要缓存工件时将其释放。