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

作者:

Sayed Adel，Matti Picus，Ralf Gommers

状态:

最终

类型:

标准

创建:

2019-11-25

决议:

https://mail.python.org/archives/list/[email protected]/thread/PVWJ74UVBRZ5ZWF6MDU7EUSJXVNILAQB/#PVWJ74UVBRZ5ZWF6MDU7EUSJXVNILAQB

摘要#

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

我们在 ufunc 机制中有一种机制可以构建替代循环，这些循环由 CPU 功能名称索引。在导入时（在 InitOperators 中），与运行时 CPU 信息匹配的循环函数被选中。本 NEP 提出了一种机制，可以在此基础上为更多功能和架构构建。提出的步骤是

  • 建立一套定义明确、与架构无关的通用内在函数，这些函数捕获跨架构可用功能。

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

  • 在运行时，发现可用的 CPU 功能，并从可能的代码路径中选择最佳路径。

动机和范围#

传统上，NumPy 依赖于编译器为目标架构生成最佳代码。但是，如今很少有用户为其机器在本地编译 NumPy。大多数用户使用二进制包，这些包必须为最低公分母 CPU 架构提供运行时支持。因此，NumPy 无法利用其 CPU 处理器的更高级功能，因为这些功能可能并非在所有用户的系统上都可用。

传统上，CPU 功能是通过内在函数公开的，内在函数是直接映射到汇编指令的特定于编译器的指令。最近，人们讨论了添加更多内在函数的有效性（例如，gh-11113 用于浮点数的 AVX 优化）。过去，针对快速 avx512 例程 向 NumPy 添加了特定于架构的代码，使用上面描述的机制为架构选择最佳循环。但是，代码不是通用的，并且不能推广到其他架构。

最近，OpenCV 开始在硬件抽象层 (HAL) 中使用通用内在函数，这为常见的共享单指令多数据 (SIMD) 结构提供了很好的抽象。本 NEP 为 NumPy 提出了一种类似的机制。使用该机制有三个阶段

  • 代码中提供了抽象内在函数的基础设施。将使用这些抽象内在函数集扩展 ufunc 机制，以便单个 ufunc 将表示为一系列循环，从最小到最大可能可用的内在函数集。

  • 在编译时，编译器宏和 CPU 检测用于将抽象内在函数转换为具体的内在函数调用。平台上任何不可用的内在函数（因为 CPU 不支持它们（因此无法测试）或者因为抽象内在函数在平台上没有并行的具体内在函数）都不会出错，而是相应的循环不会被生成并添加到可能性集中。

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

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

ufunc 机制已经能够在运行时为特定可用的 CPU 功能选择最佳循环，目前用于avx2、fma 和avx512f 循环（在生成的 __umath_generated.c 文件中）；通用内在函数将扩展生成的代码以包含更多循环变体。

用法和影响#

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

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

通过此过程发布的二进制文件将更大，因为它们包含该架构的所有可能的循环。一些打包程序可能更喜欢限制循环数量以限制二进制文件的大小，我们希望他们仍然支持各种架构系列。请注意，此问题已存在于 Intel MKL 产品中，其中二进制包包含大量针对各种 CPU 替代方案的替代共享对象 (DLL)。

源代码构建#

请参阅下面的“详细说明”。打包程序了解目标机器详细信息的源代码构建可以通过命令行参数选择仅编译目标所需的循环来理论上生成较小的二进制文件。

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

添加使用内在函数的更多代码将使代码更难以维护。因此，只有当它产生显着的性能优势时，才应添加此类代码。评估此性能优势可能并非易事。为了帮助解决这个问题，此 NEP 的实现将添加一种方法，可以通过环境变量在*运行时*选择可以使用哪些指令集。（名称待定）。此功能对于 CI 代码验证至关重要。

诊断#

Python 将提供一个新的字典 __cpu_features__。键表示可用的特性，值为布尔值，表示该特性是否可用。内部将使用各种新的私有 C 函数来查询可用的特性。这些函数可能会通过特定的 c 扩展模块公开以进行测试。

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

NumPy 将始终为任何可能适合 SIMD 矢量化的代码提供基线 C 实现。如果贡献者希望为某些架构（通常是他们最感兴趣的架构）添加 SIMD 支持，则此注释是关于如何操作的教程的开始：numpy/numpy#13516

目前，NumPy 具有许多针对许多 ufunc 的 avx512f、avx2 和 fma SIMD 循环。这些可能是首先移植到通用内联函数的候选者。预期新的实现可能会导致基准测试结果下降，但不会增加二进制文件的大小。如果下降幅度不小，我们可能会选择保留该平台的特定于 X86 的代码，并对其他平台使用通用内联函数代码。

任何使用内联函数实现 ufunc 的新的 PR 都应使用通用内联函数。如果可以证明使用通用内联函数过于繁琐或性能不足，则也可以接受特定于平台的代码。在极少数情况下，可以接受仅针对单个平台的 PR，但这必须在优先使用通用内联函数解决方案的框架内进行检查。

接受新循环的主观标准是

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

  • 代码膨胀：编译后的 wheel 的源代码大小，尤其是二进制文件大小。

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

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

添加新的内联函数#

如果贡献者希望使用尚未作为通用内联函数支持的特定于平台的 SIMD 指令，则

  1. 应将其添加为所有平台的通用内联函数

  2. 如果它在其他平台上没有等效指令（例如，_mm512_mask_i32gather_ps 在 AVX512 中），则不应添加通用内联函数，而应改为编写特定于平台的 ufunc 或简短的辅助函数。如果使用此类辅助函数，则必须使用特性宏对其进行包装，并提供一个合理的非内联函数回退作为默认值。

我们预计 (2) 将是例外情况。贡献者和维护者应该考虑，与使用基于最佳可用通用内联函数的实现相比，该单平台内联函数是否值得。

其他项目的重用#

如果通用内联函数可用于其他库（如 SciPy 或 Astropy），这些库也构建 ufunc，那就太好了，但这并非此 NEP 首次实现的明确目标。

向后兼容性#

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

详细说明#

特定于 CPU 的内容映射到通用内联函数，这些内联函数对于所有 x86 SIMD 变体、ARM SIMD 变体等都相似。例如，NumPy 通用内联函数 npyv_load_u32 映射到

  • vld1q_u32 用于基于 ARM 的 NEON

  • _mm256_loadu_si256 用于基于 x86 的 AVX2

  • _mm512_loadu_si512 用于基于 x86 的 AVX-512

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

runtests.py 和 setup.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.src、loops.avx2.c.src、loops.sse.c.src、loops.vsx.c.src、loops.neon.c.src 等）。这更类似于 PIXMAX 的做法。但是这里有很多重复，手动代码重复需要一位专门负责实现和维护该平台循环代码的负责人。

讨论#

大部分讨论发生在接受此 NEP 的 PR gh-15228 上。邮件列表上的讨论提到了 VOLK，该内容已添加到相关工作部分。可维护性问题也在邮件列表和 gh-15228 中提出并解决如下：

  • 如果贡献者想利用特定的 SIMD 指令，他们是否也需要为所有其他架构添加该指令的软件实现？（参见工作流程的 new-intrinsics 部分）。

  • 验证所有架构的代码和基准测试的责任在于谁？如果用通用 ufunc 代替特定于架构的代码有助于一个架构，但在另一个架构上损害性能，会发生什么？（在工作流程的 tradeoffs 部分中回答）。

参考文献和脚注#