在 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 提供了部分环境变量用于性能调优。
名称 | 描述 |
---|---|
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 核争用。