Vulkan 后端¶

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

注意

Vulkan 代理目前正处于积极开发中，其组件可能会发生更改。

什么是 Vulkan？¶

Vulkan 是一个低级 GPU API 规范，作为 OpenGL 的继任者而开发。它旨在为开发者提供比以前的规范更精细的 GPU 控制，以减少开销并最大限度地发挥现代图形硬件的功能。

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

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

Vulkan 计算库¶

ExecuTorch Vulkan 代理是围绕一个名为 **Vulkan 计算库** 的独立运行时进行包装的。Vulkan 计算库的目的是通过 GLSL 计算着色器为 PyTorch 运算符提供 GPU 实现。

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

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 代理目前支持以下功能：

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

除了增加运算符覆盖范围外，以下功能目前也正在开发中：

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

端到端示例¶

为了进一步了解 Vulkan 代理的功能及其使用方法，请考虑以下使用简单单运算符模型的端到端示例。

编译并降低模型到 Vulkan 代理¶

假设 ExecuTorch 已设置并安装，可以使用以下脚本将降低的 MobileNet V2 模型生成为 vulkan_mobilenetv2.pte。

一旦 ExecuTorch 设置并安装完毕，就可以使用以下脚本生成一个简单的模型并将其降低到 Vulkan 代理。

# 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 代理一样，可以使用 to_backend() API 将模型降低到 Vulkan 代理。Vulkan 代理实现了 VulkanPartitioner 类，该类识别图中 Vulkan 代理支持的节点（即运算符），并将兼容的模型部分分离以在 GPU 上执行。

这意味着即使模型包含一些不受支持的运算符，也可以将其降低到 Vulkan 代理。这只会意味着图的一部分将在 GPU 上执行。

注意

可以检查支持的操作列表 Vulkan 分区器代码，以检查 Vulkan 代理当前实现了哪些操作。

构建 Vulkan 代理库¶

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

通过使用 CMake 构建时设置 -DEXECUTORCH_BUILD_VULKAN=ON 来构建 Vulkan 代理库。

首先，请确保您已安装 Android NDK；任何 NDK 版本在 r19c 之后都可以正常工作。请注意，本文档中的示例已通过 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

构建并安装 ExecuTorch 库（用于 Android）以及 Vulkan 代理

# 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 上运行。推理可能会很慢，因为大部分运算符都是通过便携式运算符执行的。

现在，部分代理的模型就可以（部分地）在您设备的 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