利用 LiteRT 实现 NPU 加速

LiteRT 提供了一个统一的接口来使用神经处理单元 (NPU),而无需您浏览特定于供应商的编译器、运行时或库依赖项。使用 LiteRT 进行 NPU 加速可提升实时推理和大型模型推理的性能,并通过使用零复制硬件缓冲区来最大限度地减少内存复制。

开始使用

NPU 供应商

LiteRT 支持以下供应商的 NPU 加速:

Qualcomm AI Engine Direct

MediaTek NeuroPilot

Google Tensor

Google Tensor SDK 尚处于实验性访问阶段。点此注册。

AOT 和设备端编译

LiteRT NPU 同时支持 AOT 和设备端编译,以满足您的特定部署要求:

  • 离线 (AOT) 编译:此方法最适合目标 SoC 已知的大型复杂模型。提前编译可显著降低初始化费用,并在用户启动应用时降低内存用量。
  • 在线(设备端)编译:也称为 JIT 编译。此格式非常适合以平台无关的方式分发小型模型。该模型在初始化期间于用户设备上编译,无需额外的准备步骤,但首次运行成本较高。

以下指南将分三个步骤展示如何部署以进行 AOT 和设备端编译。

第 1 步:针对目标 NPU SoC 的 AOT 编译

您可以使用 LiteRT AOT(预先)编译器将 .tflite 模型编译为受支持的 SoC。您还可以在单个编译流程中同时定位多个 SoC 供应商和版本。如需了解详情,请参阅此 LiteRT AOT 编译笔记本。 虽然 AOT 编译是可选的,但强烈建议对大型模型使用 AOT 编译,以缩短设备端初始化时间。对于设备端编译,无需执行此步骤。

第 2 步:如果是在 Android 设备上,则通过 Google Play 进行部署

在 Android 上,使用 Google Play On-device AI (PODAI) 将模型和 NPU 运行时库与应用一起部署。

  • 对于设备端编译模型:将原始 .tflite 模型文件直接添加到应用的 assets/ 目录中。
  • 对于 AOT 编译模型:使用 LiteRT 将编译后的模型导出到单个 Google Play AI Pack 中。 然后,您将 AI 包上传到 Google Play,以便自动将正确的已编译模型交付到用户的设备。
  • 对于 NPU 运行时库,请使用 Play Feature Delivery 将正确的运行时库分发给用户的设备。

如需了解如何使用 Play AI PackPlay Feature Delivery 进行部署,请参阅以下部分。

使用 Play AI Pack 部署 AOT 模型

以下步骤将指导您使用 Play AI Pack 部署 AOT 编译的模型。

将 AI Pack 添加到项目中

将 AI Pack 复制到 Gradle 项目的根目录,以将 AI Pack 导入到 Gradle 项目中。例如:

my_app/
    ...
    ai_packs/
        my_model/...
        my_model_mtk/...

将每个 AI Pack 添加到 Gradle build 配置中:

// my_app/ai_packs/my_model/build.gradle.kts

plugins { id("com.android.ai-pack") }

aiPack {
  packName = "my_model"  // ai pack dir name
  dynamicDelivery { deliveryType = "on-demand" }
}

// Add another build.gradle.kts for my_model_mtk/ as well

将 AI 包添加到 Gradle 配置中

将生成的 AI Pack 中的 device_targeting_configuration.xml 复制到主应用模块的目录。然后更新 settings.gradle.kts

// my_app/setting.gradle.kts

...
// AI Packs
include(":ai_packs:my_model")
include(":ai_packs:my_model_mtk")

更新 build.gradle.kts

// my_app/build.gradle.kts

android {
 ...

 defaultConfig {
    ...
    // API level 31+ is required for NPU support.
    minSdk = 31
  }

  // AI Packs
  assetPacks.add(":ai_packs:my_model")
  assetPacks.add(":ai_packs:my_model_mtk")
}

配置 AI Pack 以实现按需分发

借助按需交付,您可以在运行时请求模型,如果模型仅在某些用户流程中需要,则此功能非常有用。模型将下载到应用的内部存储空间。在 build.gradle.kts 文件中配置 Android AI Pack 功能后,检查设备功能。 另请参阅 PODAI 中有关安装时交付快速跟进交付说明

val env = Environment.create(BuiltinNpuAcceleratorProvider(context))

val modelProvider = AiPackModelProvider(
    context, "my_model", "model/my_model.tflite") {
    if (NpuCompatibilityChecker.Qualcomm.isDeviceSupported())
      setOf(Accelerator.NPU) else setOf(Accelerator.CPU, Accelerator.GPU)
}
val mtkModelProvider = AiPackModelProvider(
    context, "my_model_mtk", "model/my_model_mtk.tflite") {
    if (NpuCompatibilityChecker.Mediatek.isDeviceSupported())
      setOf(Accelerator.NPU) else setOf()
}
val modelSelector = ModelSelector(modelProvider, mtkModelProvider)
val model = modelSelector.selectModel(env)

val compiledModel = CompiledModel.create(
    model.getPath(),
    CompiledModel.Options(model.getCompatibleAccelerators()),
    env,
)

通过 Play Feature Delivery 部署 NPU 运行时库

Play Feature Delivery 支持多种分发选项,可优化初始下载大小,包括安装时分发、按需分发、按条件分发和免安装分发。下面,我们将展示基本的安装时交付指南。

将 NPU 运行时库添加到项目

下载 litert_npu_runtime_libraries.zip(用于 AOT 编译)或 litert_npu_runtime_libraries_jit.zip(用于设备端编译),并将其解压缩到项目的根目录中:

my_app/
    ...
    litert_npu_runtime_libraries/
        mediatek_runtime/...
        qualcomm_runtime_v69/...
        qualcomm_runtime_v73/...
        qualcomm_runtime_v75/...
        qualcomm_runtime_v79/...
        qualcomm_runtime_v81/...
        fetch_qualcomm_library.sh

运行脚本以下载 NPU 支持库。例如,针对 Qualcomm NPU 运行以下命令:

$ ./litert_npu_runtime_libraries/fetch_qualcomm_library.sh

将 NPU 运行时库添加到 Gradle 配置中

将生成的 AI Pack 中的 device_targeting_configuration.xml 复制到主应用模块的目录。然后更新 settings.gradle.kts

// my_app/setting.gradle.kts

...
// NPU runtime libraries
include(":litert_npu_runtime_libraries:runtime_strings")

include(":litert_npu_runtime_libraries:mediatek_runtime")
include(":litert_npu_runtime_libraries:qualcomm_runtime_v69")
include(":litert_npu_runtime_libraries:qualcomm_runtime_v73")
include(":litert_npu_runtime_libraries:qualcomm_runtime_v75")
include(":litert_npu_runtime_libraries:qualcomm_runtime_v79")
include(":litert_npu_runtime_libraries:qualcomm_runtime_v81")

更新 build.gradle.kts

// my_app/build.gradle.kts

android {
 ...
 defaultConfig {
    ...
    // API level 31+ is required for NPU support.
    minSdk = 31

    // NPU only supports arm64-v8a
    ndk { abiFilters.add("arm64-v8a") }
    // Needed for Qualcomm NPU runtime libraries
    packaging { jniLibs { useLegacyPackaging = true } }
  }

  // Device targeting
  bundle {
      deviceTargetingConfig = file("device_targeting_configuration.xml")
      deviceGroup {
        enableSplit = true // split bundle by #group
        defaultGroup = "other" // group used for standalone APKs
      }
  }

  // NPU runtime libraries
  dynamicFeatures.add(":litert_npu_runtime_libraries:mediatek_runtime")
  dynamicFeatures.add(":litert_npu_runtime_libraries:qualcomm_runtime_v69")
  dynamicFeatures.add(":litert_npu_runtime_libraries:qualcomm_runtime_v73")
  dynamicFeatures.add(":litert_npu_runtime_libraries:qualcomm_runtime_v75")
  dynamicFeatures.add(":litert_npu_runtime_libraries:qualcomm_runtime_v79")
  dynamicFeatures.add(":litert_npu_runtime_libraries:qualcomm_runtime_v81")
}

dependencies {
  // Dependencies for strings used in the runtime library modules.
  implementation(project(":litert_npu_runtime_libraries:runtime_strings"))
  ...
}

第 3 步:使用 LiteRT 运行时在 NPU 上进行推理

LiteRT 可抽象化针对特定 SoC 版本进行开发的复杂性,让您只需几行代码即可在 NPU 上运行模型。它还提供强大的内置回退机制:您可以指定 CPU、GPU 或两者作为选项,如果 NPU 不可用,LiteRT 会自动使用它们。方便的是,AOT 编译还支持回退。它在 NPU 上提供部分委托,其中不受支持的子图在 CPU 或 GPU 上无缝运行(如指定)。

以 Kotlin 运行

请参阅以下演示版应用中的实现示例:

添加 Android 依赖项

您可以将最新的 LiteRT Maven 软件包添加到 build.gradle 依赖项中:

dependencies {
  ...
  implementation("com.google.ai.edge.litert:litert:+")
}

运行时集成

// 1. Load model and initialize runtime.
// If NPU is unavailable, inference will fallback to GPU.
val model =
    CompiledModel.create(
        context.assets,
        "model/mymodel.tflite",
        CompiledModel.Options(Accelerator.NPU, Accelerator.GPU)
    )

// 2. Pre-allocate input/output buffers
val inputBuffers = model.createInputBuffers()
val outputBuffers = model.createOutputBuffers()

// 3. Fill the first input
inputBuffers[0].writeFloat(...)

// 4. Invoke
model.run(inputBuffers, outputBuffers)

// 5. Read the output
val outputFloatArray = outputBuffers[0].readFloat()

以 C++ 跨平台运行

请参阅异步分割 C++ 应用中的示例实现。

Bazel 构建依赖项

C++ 用户必须使用 LiteRT NPU 加速功能构建应用的依赖项。用于打包核心应用逻辑(例如,cc_binarymain.cc)需要以下运行时组件:

  • LiteRT C API 共享库data 属性必须包含 LiteRT C API 共享库 (//litert/c:litert_runtime_c_api_shared_lib) 和 NPU 的特定于供应商的调度共享对象 (//litert/vendors/qualcomm/dispatch:dispatch_api_so)。
  • 特定于 NPU 的后端库:例如,适用于 Android 主机(如 libQnnHtp.solibQnnHtpPrepare.so)的 Qualcomm AI RT (QAIRT) 库和相应的 Hexagon DSP 库 (libQnnHtpV79Skel.so)。这可确保 LiteRT 运行时能够将计算任务分流到 NPU。
  • 属性依赖项deps 属性链接到必要的编译时依赖项,例如 LiteRT 的张量缓冲区 (//litert/cc:litert_tensor_buffer) 和 NPU 调度层 (//litert/vendors/qualcomm/dispatch:dispatch_api) 的 API。这样一来,您的应用代码就可以通过 LiteRT 与 NPU 进行交互。
  • 模型文件和其他资源:通过 data 属性包含。

此设置允许已编译的二进制文件动态加载并使用 NPU 来加速机器学习推理。

设置 NPU 环境

某些 NPU 后端需要运行时依赖项或库。使用编译后的模型 API 时,LiteRT 通过 Environment 对象来整理这些要求。使用以下代码查找合适的 NPU 库或驱动程序:

// Provide a dispatch library directory (following is a hypothetical path) for the NPU
std::vector<Environment::Option> environment_options = {
    {
      Environment::OptionTag::DispatchLibraryDir,
      "/usr/lib64/npu_dispatch/"
    }
};

LITERT_ASSIGN_OR_RETURN(auto env, Environment::Create(absl::MakeConstSpan(environment_options)));

运行时集成

以下代码段展示了整个流程在 C++ 中的基本实现:

// 1. Load the model that has NPU-compatible ops
LITERT_ASSIGN_OR_RETURN(auto model, Model::Load("mymodel_npu.tflite"));

// 2. Create a compiled model with NPU acceleration
//    See following section on how to set up NPU environment
LITERT_ASSIGN_OR_RETURN(auto compiled_model,
  CompiledModel::Create(env, model, kLiteRtHwAcceleratorNpu));

// 3. Allocate I/O buffers
LITERT_ASSIGN_OR_RETURN(auto input_buffers, compiled_model.CreateInputBuffers());
LITERT_ASSIGN_OR_RETURN(auto output_buffers, compiled_model.CreateOutputBuffers());

// 4. Fill model inputs (CPU array -> NPU buffers)
float input_data[] = { /* your input data */ };
input_buffers[0].Write<float>(absl::MakeConstSpan(input_data, /*size*/));

// 5. Run inference
compiled_model.Run(input_buffers, output_buffers);

// 6. Access model output
std::vector<float> data(output_data_size);
output_buffers[0].Read<float>(absl::MakeSpan(data));

利用 NPU 加速实现零复制

使用零复制可让 NPU 直接访问其自身内存中的数据,而无需 CPU 明确复制该数据。通过避免在 CPU 内存中复制数据,零复制可以显著缩短端到端延迟时间。

以下代码是使用 AHardwareBuffer 将数据直接传递给 NPU 的零复制 NPU 的示例实现。此实现避免了与 CPU 内存之间昂贵的往返,从而显著减少了推理开销。

// Suppose you have AHardwareBuffer* ahw_buffer

LITERT_ASSIGN_OR_RETURN(auto tensor_type, model.GetInputTensorType("input_tensor"));

LITERT_ASSIGN_OR_RETURN(auto npu_input_buffer, TensorBuffer::CreateFromAhwb(
    env,
    tensor_type,
    ahw_buffer,
    /* offset = */ 0
));

std::vector<TensorBuffer> input_buffers{npu_input_buffer};

LITERT_ASSIGN_OR_RETURN(auto output_buffers, compiled_model.CreateOutputBuffers());

// Execute the model
compiled_model.Run(input_buffers, output_buffers);

// Retrieve the output (possibly also an AHWB or other specialized buffer)
auto ahwb_output = output_buffers[0].GetAhwb();

串联多个 NPU 推理

对于复杂的流水线,您可以串联多个 NPU 推理。由于每个步骤都使用加速器友好型缓冲区,因此流水线主要保留在 NPU 管理的内存中:

// compiled_model1 outputs into an AHWB
compiled_model1.Run(input_buffers, intermediate_buffers);

// compiled_model2 consumes that same AHWB
compiled_model2.Run(intermediate_buffers, final_outputs);

NPU 设备端编译缓存

LiteRT 支持对 .tflite 模型进行 NPU 设备端(称为 JIT)编译。在无法提前编译模型的情况下,JIT 编译尤其有用。

不过,JIT 编译在将用户提供的模型按需转换为 NPU 字节码指令时,可能会产生一些延迟和内存开销。为了最大限度地降低性能影响,可以缓存 NPU 编译制品。

启用缓存后,LiteRT 仅会在需要时触发模型重新编译,例如:

  • 供应商的 NPU 编译器插件版本已更改;
  • Android build 指纹已更改;
  • 用户提供的模型发生了变化;
  • 编译选项已更改。

如需启用 NPU 编译缓存,请在环境选项中指定 CompilerCacheDir 环境标记。该值必须设置为应用的现有可写入路径。

   const std::array environment_options = {
        litert::Environment::Option{
            /*.tag=*/litert::Environment::OptionTag::CompilerPluginLibraryDir,
            /*.value=*/kCompilerPluginLibSearchPath,
        },
        litert::Environment::Option{
            litert::Environment::OptionTag::DispatchLibraryDir,
            kDispatchLibraryDir,
        },
        // 'kCompilerCacheDir' will be used to store NPU-compiled model
        // artifacts.
        litert::Environment::Option{
            litert::Environment::OptionTag::CompilerCacheDir,
            kCompilerCacheDir,
        },
    };

    // Create an environment.
    LITERT_ASSERT_OK_AND_ASSIGN(
        auto environment, litert::Environment::Create(environment_options));

    // Load a model.
    auto model_path = litert::testing::GetTestFilePath(kModelFileName);
    LITERT_ASSERT_OK_AND_ASSIGN(auto model,
                                litert::Model::CreateFromFile(model_path));

    // Create a compiled model, which only triggers NPU compilation if
    // required.
    LITERT_ASSERT_OK_AND_ASSIGN(
        auto compiled_model, litert::CompiledModel::Create(
                                 environment, model, kLiteRtHwAcceleratorNpu));

延迟时间和内存节省示例:

NPU 编译所需的时间和内存可能会因多种因素而异,例如底层 NPU 芯片、输入模型的复杂程度等。

下表比较了需要 NPU 编译时与由于缓存而可以跳过编译时的运行时初始化时间和内存消耗。在一款示例设备上,我们获得了以下信息:

TFLite 模型 使用 NPU 编译进行模型初始化 使用缓存的编译进行模型初始化 使用 NPU 编译初始化内存占用 使用缓存的编译初始化内存
torchvision_resnet152.tflite 7465.22 毫秒 198.34 毫秒 1525.24 MB 355.07 MB
torchvision_lraspp_mobilenet_v3_large.tflite 1592.54 毫秒 166.47 毫秒 254.90 MB 33.78 MB

在其他设备上,我们会获取以下信息:

TFLite 模型 使用 NPU 编译进行模型初始化 使用缓存的编译进行模型初始化 使用 NPU 编译初始化内存占用 使用缓存的编译初始化内存
torchvision_resnet152.tflite 2766.44 毫秒 379.86 毫秒 653.54 MB 501.21 MB
torchvision_lraspp_mobilenet_v3_large.tflite 784.14 毫秒 231.76 毫秒 113.14 MB 67.49 MB