WeGO PyTorch 中的 ZenDNN - 3.5 简体中文

Vitis AI 用户指南 (UG1414)

Document ID
UG1414
Release Date
2023-09-28
Version
3.5 简体中文

在 WeGO PyTorch 中启用 ZenDNN

ZenDNN 初始为禁用状态,但您可选择通过 WeGO-Torch 的编译 API 将其启用:
wego_mod = wego_torch.compile(mod, wego_torch.CompileOptions( 
...
optimize_options = wego_torch.OptimizeOptions(zendnn_enable = True)) 
) 

启用 ZenDNN 后,已编译的 WeGO 计算图中的 CPU 运算符(DPU 不支持的运算符)将被替换为 ZenDNN 运算符,并且会使用 ZenDNN 内核执行这些运算符以供加速。

环境变量

ZenDNN 提供了部分环境变量用于性能调优。
表 1. 环境变量
名称 描述
OMP_DYNAMIC 要启用 ZenDNN 时,请将其显式设置为 FALSE。
ZENDNN_GEMM_ALGO 默认值为 3。您可设置 [0, 1, 2, 3, 4] 来调整不同 GEMM ALGO 路径。
OMP_NUM_THREADS

默认值是用户系统的物理核数。您需要根据推断线程数进行调整以达成更好的性能。欲知详情,请参阅调谐准则。

调谐准则

ZenDNN 使用 OpenMP 作为底层库。OMP_NUM_THREADS 环境变量用于控制运算符内部并行度,即 ZenDNN 内核中的多核并行度。对于 OpenMP,不同的应用线程或运算符间的线程可使用不同的 OpenMP 线程池来处理运算符内的任务。因此,在多线程应用程序中可能会使用许多 OpenMP 线程,这将耗用大量 CPU 内核资源并降低整体性能。因此,建议的 OMP_NUM_THREADS 调谐值是根据目标 CPU 平台中的核心数以及应用中使用的线程数设置的,以避免超额预订。例如,如果您在单个应用中启动 16 个线程,并且您的平台上有 64 个 CPU 核,那么您可设置 OMP_NUM_THREADS <= 4 以避免 CPU 核争用。