ExecuTorch 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 开始包含在 Android 中。

请注意，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 r25 验证过。还应该安装 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 上运行。由于大多数运算符通过可移植运算符执行，因此推理速度预计会很慢。

现在，部分委托模型可以在您设备的 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