使用 LiteRT Next 的 GPU 加速

与 CPU 相比,图形处理单元 (GPU) 具有巨大的并行吞吐量,因此通常用于深度学习加速。LiteRT Next 允许用户在创建编译型模型 (CompiledModel) 时将硬件加速作为参数指定,从而简化了使用 GPU 加速的过程。LiteRT Next 还使用了 LiteRT 不提供的经过改进的新版 GPU 加速实现。

借助 LiteRT Next 的 GPU 加速功能,您可以创建适合 GPU 的输入和输出缓冲区,在 GPU 显存中实现零拷贝数据,并异步执行任务以最大限度地提高并行性。

如需查看支持 GPU 的 LiteRT Next 的实现示例,请参阅以下演示版应用:

添加 GPU 依赖项

请按照以下步骤为 Kotlin 或 C++ 应用添加 GPU 依赖项。

Kotlin

对于 Kotlin 用户,GPU 加速器是内置的,除了开始使用指南之外,无需执行其他步骤。

C++

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

  • LiteRT C API 共享库data 属性必须包含 LiteRT C API 共享库 (//litert/c:litert_runtime_c_api_shared_lib) 和 GPU 专用组件 (@litert_gpu//:jni/arm64-v8a/libLiteRtGpuAccelerator.so)。
  • 属性依赖项deps 属性通常包含 GLES 依赖项 gles_deps()linkopts 通常包含 gles_linkopts()。这两者与 GPU 加速密切相关,因为 LiteRT 通常在 Android 上使用 OpenGLES。
  • 模型文件和其他资源:通过 data 属性添加。

以下是 cc_binary 规则的示例:

cc_binary(
    name = "your_application",
    srcs = [
        "main.cc",
    ],
    data = [
        ...
        # litert c api shared library
        "//litert/c:litert_runtime_c_api_shared_lib",
        # GPU accelerator shared library
        "@litert_gpu//:jni/arm64-v8a/libLiteRtGpuAccelerator.so",
    ],
    linkopts = select({
        "@org_tensorflow//tensorflow:android": ["-landroid"],
        "//conditions:default": [],
    }) + gles_linkopts(), # gles link options
    deps = [
        ...
        "//litert/cc:litert_tensor_buffer", # litert cc library
        ...
    ] + gles_deps(), # gles dependencies
)

通过这种设置,编译后的二进制文件可以动态加载和使用 GPU 来加速机器学习推理。

开始使用

如需开始使用 GPU 加速器,请在创建编译模型 (CompiledModel) 时传递 GPU 参数。以下代码段展示了整个过程的基本实现:

C++

// 1. Load model
LITERT_ASSIGN_OR_RETURN(auto model, Model::CreateFromFile("mymodel.tflite"));

// 2. Create a compiled model targeting GPU
LITERT_ASSIGN_OR_RETURN(auto env, Environment::Create({}));
LITERT_ASSIGN_OR_RETURN(auto compiled_model, CompiledModel::Create(env, model, kLiteRtHwAcceleratorGpu));

// 3. Prepare input/output buffers
LITERT_ASSIGN_OR_RETURN(auto input_buffers, compiled_model.CreateInputBuffers());
LITERT_ASSIGN_OR_RETURN(auto output_buffers, compiled_model.CreateOutputBuffers());

// 4. Fill input data (if you have CPU-based data)
input_buffers[0].Write<float>(absl::MakeConstSpan(cpu_data, data_size));

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

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

Kotlin

// Load model and initialize runtime
val  model =
    CompiledModel.create(
        context.assets,
        "mymodel.tflite",
        CompiledModel.Options(Accelerator.GPU),
        env,
    )

// Preallocate input/output buffers
val inputBuffers = model.createInputBuffers()
val outputBuffers = model.createOutputBuffers()

// Fill the first input
inputBuffers[0].writeFloat(FloatArray(data_size) { data_value /* your data */ })

// Invoke
model.run(inputBuffers, outputBuffers)

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

如需了解详情,请参阅 C++ 使用入门Kotlin 使用入门指南。

LiteRT Next GPU 加速器

新的 GPU 加速器仅适用于 LiteRT Next,经过优化,比之前的版本更高效地处理 AI 工作负载,例如大型矩阵乘法和 LLM 的 KV 缓存。与 LiteRT 版本相比,LiteRT Next GPU 加速器具有以下主要改进:

  • 扩展的运算符覆盖范围:处理更大、更复杂的神经网络。
  • 更好的缓冲区互操作性:支持直接使用 GPU 缓冲区来处理相机帧、2D 纹理或大型 LLM 状态。
  • 异步执行支持:将 CPU 预处理与 GPU 推理重叠。

使用 GPU 加速的零拷贝

使用零拷贝功能可让 GPU 直接访问其自身内存中的数据,而无需 CPU 明确复制这些数据。由于无需在 CPU 内存中复制数据,零拷贝可以显著缩短端到端延迟时间。

以下代码是使用 OpenGL(用于渲染矢量图形的 API)实现零拷贝 GPU 的示例。该代码会将 OpenGL 缓冲区格式的图片直接传递给 LiteRT Next:

// Suppose you have an OpenGL buffer consisting of:
// target (GLenum), id (GLuint), size_bytes (size_t), and offset (size_t)
// Load model and compile for GPU
LITERT_ASSIGN_OR_RETURN(auto model, Model::CreateFromFile("mymodel.tflite"));
LITERT_ASSIGN_OR_RETURN(auto env, Environment::Create({}));
LITERT_ASSIGN_OR_RETURN(auto compiled_model,
    CompiledModel::Create(env, model, kLiteRtHwAcceleratorGpu));

// Create a TensorBuffer that wraps the OpenGL buffer.
LITERT_ASSIGN_OR_RETURN(auto tensor_type, model.GetInputTensorType("input_tensor_name"));
LITERT_ASSIGN_OR_RETURN(auto gl_input_buffer, TensorBuffer::CreateFromGlBuffer(env,
    tensor_type, opengl_buffer.target, opengl_buffer.id, opengl_buffer.size_bytes, opengl_buffer.offset));
std::vector<TensorBuffer> input_buffers{gl_input_buffer};
LITERT_ASSIGN_OR_RETURN(auto output_buffers, compiled_model.CreateOutputBuffers());

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

// If your output is also GPU-backed, you can fetch an OpenCL buffer or re-wrap it as an OpenGL buffer:
LITERT_ASSIGN_OR_RETURN(auto out_cl_buffer, output_buffers[0].GetOpenClBuffer());

异步执行

借助 LiteRT 的异步方法(例如 RunAsync()),您可以在安排 GPU 推理时继续使用 CPU 或 NPU 执行其他任务。在复杂的流水线中,GPU 通常与 CPU 或 NPU 一起异步使用。

以下代码段基于零拷贝 GPU 加速示例中提供的代码进行构建。该代码会异步使用 CPU 和 GPU,并将 LiteRT Event 附加到输入缓冲区。LiteRT Event 负责管理不同类型的同步基元,以下代码会创建类型为 LiteRtEventTypeEglSyncFence 的托管 LiteRT 事件对象。此 Event 对象可确保在 GPU 完成之前,我们不会从输入缓冲区读取数据。所有这些操作都无需涉及 CPU。

LITERT_ASSIGN_OR_RETURN(auto env, Environment::Create({}));
LITERT_ASSIGN_OR_RETURN(auto compiled_model,
    CompiledModel::Create(env, model, kLiteRtHwAcceleratorGpu));

// 1. Prepare input buffer (OpenGL buffer)
LITERT_ASSIGN_OR_RETURN(auto gl_input,
    TensorBuffer::CreateFromGlBuffer(env, tensor_type, opengl_tex));
std::vector<TensorBuffer> inputs{gl_input};
LITERT_ASSIGN_OR_RETURN(auto outputs, compiled_model.CreateOutputBuffers());

// 2. If the GL buffer is in use, create and set an event object to synchronize with the GPU.
LITERT_ASSIGN_OR_RETURN(auto input_event,
    Event::CreateManagedEvent(env, LiteRtEventTypeEglSyncFence));
inputs[0].SetEvent(std::move(input_event));

// 3. Kick off the GPU inference
compiled_model.RunAsync(inputs, outputs);

// 4. Meanwhile, do other CPU work...
// CPU Stays busy ..

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

支持的模型

LiteRT Next 支持以下模型的 GPU 加速。基准测试结果基于在 Samsung Galaxy S24 设备上运行的测试。

型号 LiteRT GPU 加速 LiteRT GPU(毫秒)
hf_mms_300m 完全委托 19.6
hf_mobilevit_small 完全委托 8.7
hf_mobilevit_small_e2e 完全委托 8.0
hf_wav2vec2_base_960h 完全委托 9.1
hf_wav2vec2_base_960h_dynamic 完全委托 9.8
isnet 完全委托 43.1
timm_efficientnet 完全委托 3.7
timm_nfnet 完全委托 9.7
timm_regnety_120 完全委托 12.1
torchaudio_deepspeech 完全委托 4.6
torchaudio_wav2letter 完全委托 4.8
torchvision_alexnet 完全委托 3.3
torchvision_deeplabv3_mobilenet_v3_large 完全委托 5.7
torchvision_deeplabv3_resnet101 完全委托 35.1
torchvision_deeplabv3_resnet50 完全委托 24.5
torchvision_densenet121 完全委托 13.9
torchvision_efficientnet_b0 完全委托 3.6
torchvision_efficientnet_b1 完全委托 4.7
torchvision_efficientnet_b2 完全委托 5.0
torchvision_efficientnet_b3 完全委托 6.1
torchvision_efficientnet_b4 完全委托 7.6
torchvision_efficientnet_b5 完全委托 8.6
torchvision_efficientnet_b6 完全委托 11.2
torchvision_efficientnet_b7 完全委托 14.7
torchvision_fcn_resnet50 完全委托 19.9
torchvision_googlenet 完全委托 3.9
torchvision_inception_v3 完全委托 8.6
torchvision_lraspp_mobilenet_v3_large 完全委托 3.3
torchvision_mnasnet0_5 完全委托 2.4
torchvision_mobilenet_v2 完全委托 2.8
torchvision_mobilenet_v3_large 完全委托 2.8
torchvision_mobilenet_v3_small 完全委托 2.3
torchvision_resnet152 完全委托 15.0
torchvision_resnet18 完全委托 4.3
torchvision_resnet50 完全委托 6.9
torchvision_squeezenet1_0 完全委托 2.9
torchvision_squeezenet1_1 完全委托 2.5
torchvision_vgg16 完全委托 13.4
torchvision_wide_resnet101_2 完全委托 25.0
torchvision_wide_resnet50_2 完全委托 13.4
u2net_full 完全委托 98.3
u2net_lite 完全委托 51.4
hf_distil_whisper_small_no_cache 部分委托 251.9
hf_distilbert 部分委托 13.7
hf_tinyroberta_squad2 部分委托 17.1
hf_tinyroberta_squad2_dynamic_batch 部分委托 52.1
snapml_StyleTransferNet 部分委托 40.9
timm_efficientformer_l1 部分委托 17.6
timm_efficientformerv2_s0 部分委托 16.1
timm_pvt_v2_b1 部分委托 73.5
timm_pvt_v2_b3 部分委托 246.7
timm_resnest14d 部分委托 88.9
torchaudio_conformer 部分委托 21.5
torchvision_convnext_tiny 部分委托 8.2
torchvision_maxvit_t 部分委托 194.0
torchvision_shufflenet_v2 部分委托 9.5
torchvision_swin_tiny 部分委托 164.4
torchvision_video_resnet2plus1d_18 部分委托 6832.0
torchvision_video_swin3d_tiny 部分委托 2617.8
yolox_tiny 部分委托 11.2