ExecuTorch Vulkan 委托¶

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

注意

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

什么是 Vulkan？¶

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

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

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

Vulkan 计算库¶

ExecuTorch Vulkan 委托是围绕一个名为 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 委托目前支持以下功能

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

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

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

端到端示例¶

为了进一步了解 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 版本高于 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

要构建和安装带有 Vulkan 委托的 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