NEP 38 — 使用SIMD优化指令提升性能#

作者:

Sayed Adel, Matti Picus, Ralf Gommers

状态:

最终

类型:

标准

创建时间:

2019-11-25

解决时间:

NumPy讨论

摘要#

虽然编译器在利用特定硬件指令优化代码方面越来越好,但它们有时并不能产生最优结果。此外,我们希望能够将二进制优化的C扩展模块从一台机器复制到具有相同基础架构(x86、ARM或PowerPC)但具有不同功能的另一台机器,而无需重新编译。

我们在ufunc机制中有一个机制,可以通过CPU功能名称索引来构建替代循环。在导入时(在InitOperators中),会从候选中选择与运行时CPU信息匹配的循环函数。此NEP提出了一种机制,可以为更多功能和架构扩展此功能。提出的步骤是

  • 建立一套定义良好、与架构无关、通用的内建函数(intrinsics),用于捕获跨架构可用的功能。

  • 使用C宏集捕获这些通用的内建函数,并使用宏从基线到架构上可用的功能集的最大集来构建功能集代码路径。将这些作为有限数量的编译好的替代代码路径提供。

  • 在运行时,发现可用的CPU功能,并据此从可能的代码路径中选择。

动机与范围#

传统上,NumPy依赖编译器为目标架构生成特定的优化代码。然而,如今很少有用户在本地为自己的机器编译NumPy。大多数用户使用二进制包,这些包必须为最低兼容的CPU架构提供运行时支持。因此,NumPy无法利用其CPU处理器更高级的功能,因为它们可能并非在所有用户的系统上都可用。

传统上,CPU功能通过内建函数暴露,这些内建函数是编译器特定的指令,直接映射到汇编指令。最近,有关于添加更多内建函数有效性的讨论(例如,用于浮点数的AVX优化的gh-11113)。过去,NumPy中曾为各种ufunc中的快速avx512例程添加了特定于架构的代码,使用上述机制为架构选择最佳循环。然而,这些代码不是通用的,也无法推广到其他架构。

最近,OpenCV开始使用硬件抽象层(HAL)中的通用内建函数,为常用的共享单指令多数据(SIMD)结构提供了良好的抽象。此NEP为NumPy提出了类似的机制。使用该机制有三个阶段:

  • 代码中提供了抽象内建函数的底层结构。ufunc机制将使用这些抽象内建函数的集合进行扩展,以便一个ufunc将表示为一组循环,从最小到最大可能可用的内建函数集。

  • 在编译时,编译器宏和CPU检测用于将抽象内建函数转换为具体的内建函数调用。平台上不可用的任何内建函数,无论是由于CPU不支持(因此无法测试)还是由于抽象内建函数在平台上没有对应的具体内建函数,都不会报错,而是相应的循环将不会被生成并添加到可能性集合中。

  • 在运行时,CPU检测代码将进一步限制可用循环集,并将为ufunc选择最佳循环。

当前的NEP仅提议将运行时功能检测和最佳循环选择机制用于ufunc。未来的NEP可能会为提出的解决方案提出其他用途。

ufunc机制已经能够为特定可用的CPU功能在运行时选择最佳循环,目前用于avx2fmaavx512f循环(在生成的__umath_generated.c文件中);通用内建函数将扩展生成的代码以包含更多循环变体。

用法与影响#

最终用户将能够获取其平台和编译器可用的内建函数列表。可选地,用户可能能够通过环境变量指定运行时可用的循环中的哪一个将被使用,也许是为了对不同循环的影响进行基准测试。对普通最终用户应该没有直接影响,所有循环的结果应该相同,误差在1-3个ULP以内。另一方面,拥有更强大机器的用户应该会注意到显著的性能提升。

二进制发行版 - PyPI上的wheels和conda包#

通过此过程发布的二进制文件会更大,因为它们包含了架构的所有可能循环。一些打包者可能希望限制循环数量以减小二进制文件大小,我们希望他们仍然支持广泛的架构族。请注意,Intel MKL产品已经存在此问题,其中二进制包包含针对各种CPU替代项的广泛的共享对象(DLL)集。

源码构建#

参见下面的“详细描述”。理论上,一个了解目标机器细节的打包者进行的源码构建,可以通过命令行参数选择仅编译目标所需的循环,从而生成更小的二进制文件。

如何运行基准测试以评估性能优势#

添加更多使用内建函数的代码会增加维护难度。因此,只有在带来显著性能优势时,才应添加此类代码。评估这种性能优势可能并不容易。为了帮助这一点,此NEP的实现将增加一种方式,通过环境变量(名称待定)在*运行时*选择可用的指令集。这种能力对于CI代码验证至关重要。

诊断#

Python中将提供一个新字典__cpu_features__。键是可用的功能,值是该功能是否可用的布尔值。内部将使用各种新的私有C函数来查询可用功能。这些可能会通过特定的c扩展模块暴露以供测试。

添加新的CPU架构特定优化的工作流程#

NumPy将始终为任何可能进行SIMD向量化的代码提供基线C实现。如果贡献者想为某个架构添加SIMD支持(通常是他们最感兴趣的架构),那么此评论是关于如何操作的教程的开端:numpy/numpy#13516

截至目前,NumPy为许多ufunc提供了avx512favx2fma SIMD循环。这些可能是首先移植到通用内建函数的候选者。期望新的实现可能会导致基准测试出现回归,但不会增加二进制文件的大小。如果回归不显著,我们可能会选择为该平台保留X86特定的代码,并为其他平台使用通用内建函数代码。

任何使用内建函数实现ufunc的新PR都将期望使用通用内建函数。如果可以证明使用通用内建函数过于笨拙或性能不够,也可以接受特定于平台的代码。在极少数情况下,可以接受仅限于单个平台的PR,但必须在倾向于使用通用内建函数的解决方案的框架内进行审查。

接受新循环的主观标准是

  • 正确性:新代码在算法的边缘点也不能将精度降低超过1-3个ULP。

  • 代码膨胀:包括源代码大小,尤其是编译后的wheel的二进制大小。

  • 可维护性:代码的可读性如何

  • 性能:基准测试必须显示出显著的性能提升

添加新的内建函数#

如果贡献者想使用尚未作为通用内建函数支持的特定于平台的SIMD指令,那么

  1. 它应该作为所有平台的通用内建函数添加

  2. 如果它在其他平台上没有等效指令(例如AVX中的_mm512_mask_i32gather_ps),则不应添加通用内建函数,而应改为编写特定于平台的ufunc或简短的辅助函数。如果使用此类辅助函数,则必须用功能宏将其包装,并提供一个合理的非内建回退以供默认使用。

我们预计(2)将是例外。贡献者和维护者应考虑该仅限于单个平台的内建函数是否值得,相比于使用基于通用内建函数的最佳实现。

其他项目的重用#

如果通用内建函数能够被SciPy或Astropy等也构建ufunc的其他库使用,那将是很好的,但这并不是此NEP首次实现的明确目标。

向后兼容性#

不应对向后兼容性产生任何影响。

详细描述#

特定于CPU的指令被映射到通用内建函数,这些内建函数对于所有x86 SIMD变体、ARM SIMD变体等都是相似的。例如,NumPy的通用内建函数npyv_load_u32映射到

  • ARMNeon的vld1q_u32

  • 基于x86的AVX2的_mm256_loadu_si256

  • 基于x86的AVX-512的_mm512_loadu_si512

任何编写SIMD循环的人都将使用npyv_load_u32宏而不是特定于架构的内建函数。代码还为编译和运行时提供了保护宏,以便可以选择正确的循环。

runtests.pysetup.py提供了两个新的构建选项:--cpu-baseline--cpu-dispatch--cpu-baseline定义了编译所需的绝对最低功能。例如,在x86_64上,它默认为SSE3。如果编译器支持,将启用最低功能。--cpu-dispatch设置了可检测并用作分派要求的集合的其他内建函数。例如,在x86_64上,它默认为[SSSE3, SSE41, POPCNT, SSE42, AVX, F16C, XOP, FMA4, FMA3, AVX2, AVX512F, AVX512CD, AVX512_KNL, AVX512_KNM, AVX512_SKX, AVX512_CLX, AVX512_CNL, AVX512_ICL]。所有这些功能都被映射到一个c级布尔数组npy__cpu_have,以及一个c级便捷函数npy_cpu_have(int feature_id),该函数查询此数组,结果在运行时存储在__cpu_features__中。

导入ufunc时,会将被编译循环所需的功能与已检测到的功能进行匹配。匹配最佳的循环将被标记为由ufunc调用。

实现#

当前PR

第一个PR提供了编译时和运行时代码底层结构。第二个PR展示了为循环使用此底层结构的示例。NEP获批后,需要进行更多工作来编写使用NEP提供的机制的循环。

替代方案#

gh-13516中提出的一个替代方案是手动实现每个CPU架构的循环,而不尝试抽象SIMD内建函数中的常见模式(例如,拥有loops.avx512.c.srcloops.avx2.c.srcloops.sse.c.srcloops.vsx.c.srcloops.neon.c.src等)。这更类似于PIXMAX的做法。然而,这里存在大量的重复,手动代码重复需要一个专门负责实现和维护该平台循环代码的负责人。

讨论#

大部分讨论发生在接受此NEP的PRgh-15228上。邮件列表上的讨论提到了VOLK,它被添加到了相关工作部分。可维护性问题也在邮件列表和gh-15228上提出,并解决了以下问题:

  • 如果贡献者想利用特定的SIMD指令,是否需要为所有其他架构也添加该指令的软件实现?(参见工作流程的新内建函数部分)。

  • 谁有责任验证所有架构的代码和基准测试?如果用通用ufunc替换特定于架构的代码,对一个架构有益但损害了另一个架构的性能,会发生什么?(在工作流程的权衡部分已回答)。

参考文献和脚注#