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