Vulkan 后端¶

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

注意

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

什么是 Vulkan？¶

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

Vulkan 已被 GPU 供应商广泛采用，市场上的大多数现代 GPU（包括桌面和移动设备）都支持 Vulkan。从 Android 7.0 开始，Android 也包含了 Vulkan。

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

Vulkan 计算库¶

ExecuTorch Vulkan Delegate 是一个独立运行时（称为 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 delegate 当前支持以下特性

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

除了增加算子覆盖范围外，以下特性目前正在开发中

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

端到端示例¶

为了进一步理解 Vulkan Delegate 的特性及使用方法，请考虑以下使用简单单算子模型的端到端示例。

将模型编译和 Lowering 到 Vulkan Delegate¶

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

设置并安装 ExecuTorch 后，可以使用以下脚本生成一个简单的模型并将其 Lowering 到 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 delegates 一样，可以使用 to_backend() API 将模型 Lowering 到 Vulkan Delegate。Vulkan Delegate 实现了 VulkanPartitioner 类，该类用于识别图中 Vulkan delegate 支持的节点（即算子），并将模型中兼容的部分分离出来，以便在 GPU 上执行。

这意味着即使模型包含一些不受支持的算子，也可以将其 Lowering 到 Vulkan delegate。这仅仅意味着图的某些部分将在 GPU 上执行。

注意

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

构建 Vulkan Delegate 库¶

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

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

首先，请确保已安装 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

构建和安装带 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