后端最佳实践¶
x86 CPU¶
现代 x86 CPU 上的编译工作负载通常通过单指令多数据 (SIMD) 指令集进行优化。SIMD 是一种典型的高性能计算并行处理技术,例如深度学习模型训练和推理。应用 SIMD 后,每个计算单元在任何给定时间段内使用分配到的不同数据执行相同的指令。最常用的支持 SIMD 的 x86 指令集架构 (ISA) 包括 AVX、AVX2、AVX-512 和 AMX。
您可以通过使用 collect_env 脚本 来检查您的机器支持的 ISA。由于该脚本提供了 PyTorch 的完整环境信息,我们可以使用 grep
来提取包含 ISA 信息的行
python collect_env.py | grep "a[(v|m)]x"
通常,如果支持 AVX-512,应该会看到以“avx512”开头的指令(例如 avx512f
、avx512bw
、avx512_vnni
)。如果支持 AMX,应该会看到以“amx”开头的指令(例如 amx_tile
、amx_bf16
、amx_int8
)。
具体来说,对于启用了 AMX 指令集的服务器,可以通过利用 AMX 来进一步提升工作负载性能。