Vulkan 后端¶

ExecuTorch Vulkan delegate 是基于跨平台 Vulkan GPU API 标准构建的 ExecuTorch 原生 GPU delegate。它主要设计用于利用 GPU 加速 Android 设备上的模型推理，但也可以用于支持 Vulkan 实现的任何平台：笔记本电脑、服务器和边缘设备。

注意

Vulkan delegate 目前正积极开发中，其组件可能会发生变化。

什么是 Vulkan？¶

Vulkan 是作为 OpenGL 的后继者开发的低级 GPU API 规范。它旨在与先前的规范相比，为开发者提供对 GPU 更明确的控制，以减少开销并最大限化现代图形硬件的能力。

Vulkan 已被 GPU 供应商广泛采用，市场上的大多数现代 GPU（包括桌面和移动）都支持 Vulkan。Android 7.0 及更高版本也包含了 Vulkan。

请注意，Vulkan 是一个 GPU API，而不是一个 GPU 数学库。也就是说，它提供了一种在 GPU 上执行计算和图形操作的方式，但没有内置高性能计算内核库。

Vulkan 计算库¶

ExecuTorch Vulkan Delegate 是一个围绕独立运行时（称为 Vulkan 计算库）的包装器。Vulkan 计算库旨在通过 GLSL 计算着色器为 PyTorch 运算符提供 GPU 实现。

Vulkan 计算库是 PyTorch Vulkan 后端的一个分支/迭代。PyTorch Vulkan 后端的核心组件被分支到 ExecuTorch 中，并适配了 AOT 图模式的模型推理风格（与 PyTorch 采用的 Eager 执行模型推理风格相对）。

Vulkan 计算库的组件包含在 executorch/backends/vulkan/runtime/ 目录中。核心组件如下所列和描述

runtime/
├── api/ .................... Wrapper API around Vulkan to manage Vulkan objects
└── graph/ .................. ComputeGraph class which implements graph mode inference
    └── ops/ ................ Base directory for operator implementations
        ├── glsl/ ........... GLSL compute shaders
        │   ├── *.glsl
        │   └── conv2d.glsl
        └── impl/ ........... C++ code to dispatch GPU compute shaders
            ├── *.cpp
            └── Conv2d.cpp

特性¶

Vulkan delegate 目前支持以下特性

内存规划
- 生命周期不重叠的中间张量将共享内存分配。这降低了模型推理的峰值内存使用量。
基于能力的分割:
- 图可以通过 partitioner 部分下沉到 Vulkan delegate，partitioner 会识别图中 Vulkan delegate 支持的节点（即运算符）并仅下沉支持的子图
支持上限动态形状:
- 张量可以在推理之间改变形状，只要其当前形状小于下沉期间指定的边界

除了增加运算符覆盖率外，以下特性目前正在开发中

量化支持
- 我们目前正在开发对 8 位动态量化的支持，并计划未来扩展到其他量化方案。
内存布局管理
- 内存布局是优化性能的重要因素。我们计划引入图 passes，在整个图中引入内存布局转换，以优化对内存布局敏感的运算符，例如卷积和矩阵乘法。
选择性构建
- 我们计划通过选择要构建哪些运算符/着色器来控制构建大小

端到端示例¶

为了进一步了解 Vulkan Delegate 的特性及如何使用它，请参考以下使用简单单运算符模型的端到端示例。

编译模型并将其下沉到 Vulkan Delegate¶

假设 ExecuTorch 已设置并安装，可以使用以下脚本生成一个已下沉的 MobileNet V2 模型，文件名为 vulkan_mobilenetv2.pte。

一旦 ExecuTorch 已设置并安装，可以使用以下脚本生成一个简单模型并将其下沉到 Vulkan delegate。

# Note: this script is the same as the script from the "Setting up ExecuTorch"
# page, with one minor addition to lower to the Vulkan backend.
import torch
from torch.export import export
from executorch.exir import to_edge

from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner

# Start with a PyTorch model that adds two input tensors (matrices)
class Add(torch.nn.Module):
  def __init__(self):
    super(Add, self).__init__()

  def forward(self, x: torch.Tensor, y: torch.Tensor):
      return x + y

# 1. torch.export: Defines the program with the ATen operator set.
aten_dialect = export(Add(), (torch.ones(1), torch.ones(1)))

# 2. to_edge: Make optimizations for Edge devices
edge_program = to_edge(aten_dialect)
# 2.1 Lower to the Vulkan backend
edge_program = edge_program.to_backend(VulkanPartitioner())

# 3. to_executorch: Convert the graph to an ExecuTorch program
executorch_program = edge_program.to_executorch()

# 4. Save the compiled .pte program
with open("vk_add.pte", "wb") as file:
    file.write(executorch_program.buffer)

与其他 ExecuTorch delegate 一样，可以使用 to_backend() API 将模型下沉到 Vulkan Delegate。Vulkan Delegate 实现了 VulkanPartitioner 类，该类识别图中 Vulkan delegate 支持的节点（即运算符），并将模型中兼容的部分分离出来以便在 GPU 上执行。

这意味着即使模型包含一些不支持的运算符，也可以将其下沉到 Vulkan delegate。这仅仅意味着图中只有一部分会在 GPU 上执行。

注意

可以检查支持的 ops 列表 Vulkan partitioner 代码，以查看 Vulkan delegate 中当前实现了哪些 ops。

构建 Vulkan Delegate 库¶

构建和测试 Vulkan Delegate 最简单的方法是针对 Android 进行构建并在本地 Android 设备上测试。Android 设备内置支持 Vulkan，并且 Android NDK 附带了 GLSL 编译器，这是编译 Vulkan 计算库的 GLSL 计算着色器所必需的。

通过在 CMake 构建时设置 -DEXECUTORCH_BUILD_VULKAN=ON 可以构建 Vulkan Delegate 库。

首先，确保你已安装 Android NDK；NDK r19c 之后的任何 NDK 版本都可以工作。请注意，本文档中的示例已使用 NDK r27b 进行验证。还应该安装 Android SDK，以便你可以访问 adb。

本页中的说明假设设置了以下环境变量。

export ANDROID_NDK=<path_to_ndk>
# Select the appropriate Android ABI for your device
export ANDROID_ABI=arm64-v8a
# All subsequent commands should be performed from ExecuTorch repo root
cd <path_to_executorch_root>
# Make sure adb works
adb --version

使用 Vulkan Delegate 构建和安装 ExecuTorch 库（适用于 Android）

# From executorch root directory
(rm -rf cmake-android-out && \
  pp cmake . -DCMAKE_INSTALL_PREFIX=cmake-android-out \
    -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
    -DANDROID_ABI=$ANDROID_ABI \
    -DEXECUTORCH_BUILD_VULKAN=ON \
    -DPYTHON_EXECUTABLE=python \
    -Bcmake-android-out && \
  cmake --build cmake-android-out -j16 --target install)

在设备上运行 Vulkan 模型¶

注意

由于目前运算符支持有限，只有二元算术运算符将在 GPU 上运行。预计推理速度会较慢，因为大部分运算符是通过 Portable 运算符执行的。

现在，部分 delegated 的模型可以在你的设备的 GPU 上（部分地）执行了！

# Build a model runner binary linked with the Vulkan delegate libs
cmake --build cmake-android-out --target vulkan_executor_runner -j32

# Push model to device
adb push vk_add.pte /data/local/tmp/vk_add.pte
# Push binary to device
adb push cmake-android-out/backends/vulkan/vulkan_executor_runner /data/local/tmp/runner_bin

# Run the model
adb shell /data/local/tmp/runner_bin --model_path /data/local/tmp/vk_add.pte