应用-使用SIMD加速Cython

目录

使用SIMD加速Cython

Cython允许你通过将Python代码翻译为C或C++来编写Python的编译扩展。通常你会用它来加速你的软件，尤其是在实现小型数据科学或科学计算算法时非常有用。

但是，当Cython仍然太慢时，你还能做些什么呢？通常你仍然可以进行一些速度优化。在之前的文章中，我们重点介绍了如何通过优化代码来利用指令级并行性等特性。

在本文中，我们将重点介绍另一种CPU特性——单指令多数据（SIMD），特别是在Cython中的使用。正如我们将看到的，在某些情况下，使用SIMD只需要对代码进行最小的修改。

什么是SIMD？

单指令多数据（SIMD）是CPU指令，可以在内存中相邻存储的一系列原始值（整数或浮点数）上执行相同的操作，使用单条指令。编译器需要显式生成这些特殊的指令，这使得利用这一功能变得更加困难。

例如，如果我们想将数组中的四个连续整数乘以同一个常数，我们可以使用一条SIMD CPU指令来完成，而不是使用4条普通的CPU指令。将CPU指令数量减少到原来的四分之一可以显著提高速度！

使用SIMD的三种方式

如果你想使用SIMD，有几种方法可以让编译器生成这些特殊的CPU指令：

1. 内联函数（Intrinsics）：你可以直接告诉编译器使用特定的CPU指令，使用称为“内联函数”的特殊编译器功能。你可以在Cython中这样做，但它不具备可移植性。例如，x86-64 CPU指令与ARM CPU指令不同。
2. 自动向量化（Auto-vectorization）：编译器可以自动生成SIMD指令，特别是如果你适当地结构化你的代码。由于这依赖于编译器的启发式方法，你的控制较少，不同版本的编译器可能会给出不同的结果。（不要将此与Python中的“向量化”混淆；在C等低级语言中，“向量化”是SIMD的另一种说法。）
3. SIMD库：C++有一些库提供了使用SIMD进行数学操作的数据结构。这比内联函数提供了更多的抽象，同时比自动向量化更可靠。Cython可能可以使用这些库，但我们不会在这里讨论它们。

在本文中，我们将重点介绍使用自动向量化：让编译器为你生成SIMD指令。

我们的起始代码

我们将从以下Cython扩展开始：

import numpy as np
cimport numpy as cnp

def average_arrays_1(
    cnp.float32_t[:] arr1,
    cnp.float32_t[:] arr2
):
    out = np.empty((len(arr1),), dtype=np.float32)
    cdef cnp.float32_t[:] out_view = out
    for i in range(len(arr1)):
        out_view[i] = (arr1[i] + arr2[i]) / 2
    return out

cnp.float32_t[:] 是Cython的内存视图API；我们在所有数组上使用它，以确保Cython更容易将代码翻译为C。这确保我们使用的是快速的底层操作，而不是通用的“索引随机Python对象”操作，这些操作不会比普通Python更快。

我们希望在两种情况下测量函数的速度：

• 两个数组的所有内存都是连续的线性内存。
• 两个数组的视图，我们只抓取每16个元素，使用NumPy的“步幅”功能。代码将平均的视图在内存中不相邻；我们的代码需要跳过并处理每16个元素。

例如，在以下代码片段中，我们抓取每2个元素。我们创建的对象是一个视图，意味着我们没有复制内存。请注意，如果我们修改视图，它会影响原始数组：

>>> import numpy as np
>>> arr = np.array([0, 1, 2, 3, 4, 5, 6])
>>> view = arr[::2]
>>> view
array([0, 2, 4, 6])
>>> view[1] = 100
>>> arr
array([  0,   1, 100,   3,   4,   5,   6])

这是我们的基准测试脚本：

from time import time
import numpy as np
import sys
import simd_example

# 使用第一个命令行参数来选择调用哪个函数：
average_arrays = getattr(simd_example, sys.argv[1])

def timeit(title, arr1, arr2):
    start = time()
    for i in range(1_000):
        out = average_arrays(arr1, arr2)
    elapsed = (time() - start) / 1_000
    print(f"Time per call, {title}: {elapsed * 1000:.2f} ms")

# 内存中线性布局的数组：
ARR1 = np.random.random((1_000_000,)).astype(np.float32)
ARR2 = np.random.random((1_000_000,)).astype(np.float32)
timeit("contiguous", ARR1, ARR2)

# 抓取每16个元素的数组：
ARR3 = np.random.random((16_000_000,)).astype(np.float32)
ARR4 = np.random.random((16_000_000,)).astype(np.float32)
timeit("strided", ARR3[::16], ARR4[::16])

我们没有使用SIMD！

在Intel CPU上，如果你使用的是浮点操作，你可以通过检查CPU计数器来查看是否使用了SIMD。在Linux上，这些计数器由perf程序暴露。对于整数操作，有更复杂的方法来检查，例如查看生成的汇编代码，但我们会采用简单的方法。我的CPU不支持512位SIMD，所以我们只查看128位和256位SIMD，如下所示：

$ perf stat -e instructions:u,fp_arith_inst_retired.128b_packed_single:u,fp_arith_inst_retired.256b_packed_single:u -- python benchmark.py average_arrays_1
Time per call, contiguous: 0.57 ms
Time per call, strided: 4.2 ms

 Performance counter stats for 'python benchmark.py average_arrays_1':

    31,773,910,107      cpu_core/instructions:u/
       459,588,037      cpu_atom/instructions:u/
                 0      cpu_core/fp_arith_inst_retired.128b_packed_single:u/
                 0      cpu_core/fp_arith_inst_retired.256b_packed_single:u/

“_single”后缀表示32位浮点数；对于64位浮点数，我们需要“_double”后缀。看起来没有进行任何“打包”（即SIMD）操作。让我们来解决这个问题。

第一步：在编译器中启用自动向量化

为什么没有使用SIMD？回想一下，要使用SIMD，编译器必须生成特殊的指令。具体来说：

1. 我们希望使用自动向量化，即编译器自动完成这项工作。
2. 这意味着编译器必须成功使用其启发式方法来生成SIMD指令。

事实证明，默认情况下，在Linux上，Python使用-O2编译器优化级别来编译扩展，这甚至不会尝试进行自动向量化。因此，我们需要切换到-O3优化级别以启用此功能和其他优化。

你可以在Cython的setup.py配置中执行此操作，但我们只需设置一个环境变量并重新编译：

$ export CFLAGS="-O3"
$ pip install .

不幸的是，如果我们重新运行基准测试，结果与之前相同。这表明编译器的启发式方法未能找出如何使用SIMD来优化代码。

第二步：禁用边界检查和其他隐式条件

回想一下，SIMD对一系列值执行相同的操作。例如，我们希望编译器决定使用SIMD指令将一组float32与另一组float32相加，使用一条CPU指令。

有时编译器无法使用SIMD指令。对于我们的目的，如果你有分支（如if语句和其他条件控制流），这使得使用单条指令执行相同操作变得更加困难。如果分支决定在第一个项目上退出循环，那么一条CPU指令如何处理剩余的项目？

我们的代码没有任何显式条件，所以这不是问题。然而，它确实有一些隐式条件：边界检查和索引包装。

1. 为了防止从（或写入！）无效的内存地址，Cython为每个数组查找添加了边界检查，以确保你没有超出范围。如果索引超出边界，它会引发异常——换句话说，这是一个分支。
2. 为了支持使用负数索引（如Python所做的那样），Cython还添加了分支。

在好的情况下，编译器可能会优化掉其中的一些分支，但在这种情况下，我们可能不得不禁用它们。这意味着我们必须格外小心，因为编译器不再帮助我们捕获错误了！

注意：关闭边界检查是一个糟糕的主意。它会导致崩溃、安全问题，甚至更糟：无声的失败，你在不知道的情况下得到错误的结果。

在使用Rust时，人们已经想出了一些方法，可以在保持边界检查的同时给编译器提示，以便它可以优化掉这些检查。我希望能研究一下在Cython中是否可能做到这一点。

这是我们的新版本：

import numpy as np
cimport numpy as cnp
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def average_arrays_2(
    cnp.float32_t[:] arr1,
    cnp.float32_t[:] arr2
):
    if len(arr1) != len(arr2):
        raise ValueError("Arrays must be the same size")
    out = np.empty((len(arr1),), dtype=np.float32)
    cdef cnp.float32_t[:] out_view = out
    for i in range(len(arr1)):
        out_view[i] = (arr1[i] + arr2[i]) / 2
    return out

新版本在线性情况下运行速度提高了约15%，这很好，但它仍然没有使用任何SIMD。

第三步：去除步幅支持

是时候检查Cython从.pyx文件生成的C代码了。我们可以使用cythonize工具对代码进行注释，这将生成一个我们可以打开的HTML文件。

$ cythonize --annotate src/simd_example/_extension.pyx
$ ls src/simd_example/*.html
src/simd_example/_extension.html

报告将Cython代码映射到C代码，如果我们查看为out_view[i] = (arr1[i] + arr2[i]) / 2生成的C代码，我们会看到：

*((__pyx_t_5numpy_float32_t *) ( /* dim=0 */ (__pyx_v_out_view.data + __pyx_t_13 * __pyx_v_out_view.strides[0]) )) = (((*(__pyx_t_5numpy_float32_t *) ( /* dim=0 */ (__pyx_v_arr1.data + __pyx_t_11 * __pyx_v_arr1.strides[0]) )) + (*(__pyx_t_5numpy_float32_t *) ( /* dim=0 */ (__pyx_v_arr2.data + __pyx_t_12 * __pyx_v_arr2.strides[0]) ))) / 2.0);

这比我们的Cython代码复杂得多……因为步幅。基本上，Cython必须处理我们传入的NumPy视图的情况，该视图不是读取所有内存，而是跳过每N个项目。它不能假设内存是连续的。

如果我们有非连续的内存，从编译器的自动向量化角度来看，使用SIMD可能会变得不那么有吸引力，因为不同的项目可能在内存中相距甚远。即使我们无法优化这种情况，连续内存的情况也应该更快。因此，作为第一步，让我们告诉Cython只支持连续内存，通过将float32_t[:]切换为float32_t[::1]来表示。有关此语法的更多详细信息，请参阅Cython文档。

@cython.boundscheck(False)
@cython.wraparound(False)
def average_arrays_3(cnp.float32_t[::1] arr1, cnp.float32_t[::1] arr2):
    if len(arr1) != len(arr2):
        raise ValueError("Arrays must be the same size")
    out = np.empty((len(arr1),), dtype=np.float32)
    cdef cnp.float32_t[::1] out_view = out
    for i in range(len(arr1)):
        out_view[i] = (arr1[i] + arr2[i]) / 2
    return out

生成的C代码不再引用步幅：

*((__pyx_t_5numpy_float32_t *) ( /* dim=0 */ ((char *) (((__pyx_t_5numpy_float32_t *) __pyx_v_out_view.data) + __pyx_t_13)) )) = (((*(__pyx_t_5numpy_float32_t *) ( /* dim=0 */ ((char *) (((__pyx_t_5numpy_float32_t *) __pyx_v_arr1.data) + __pyx_t_11)) ))) + (*(__pyx_t_5numpy_float32_t *) ( /* dim=0 */ ((char *) (((__pyx_t_5numpy_float32_t *) __pyx_v_arr2.data) + __pyx_t_12)) ))) / 2.0);

我们重新编译并运行基准测试：

Time per call, contiguous: 0.24 ms
Traceback (most recent call last):
  File "/home/itamarst/devel/simd-profiles/benchmark.py", line 28, in <module>
    timeit("strided", ARR3[::16], ARR4[::16])
  File "/home/itamarst/devel/simd-profiles/benchmark.py", line 14, in timeit
    out = average_arrays(arr1, arr2)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "src/simd_example/_extension.pyx", line 41, in simd_example._extension.average_arrays_3
    def average_arrays_3(cnp.float32_t[::1] arr1, cnp.float32_t[::1] arr2):
  File "<stringsource>", line 663, in View.MemoryView.memoryview_cwrapper
  File "<stringsource>", line 353, in View.MemoryView.memoryview.__cinit__
ValueError: ndarray is not C-contiguous

 Performance counter stats for 'python benchmark.py average_arrays_3':

     4,940,110,700      cpu_core/instructions:u/
        66,742,512      cpu_atom/instructions:u/
       500,000,000      cpu_core/fp_arith_inst_retired.128b_packed_single:u/
                 0      cpu_core/fp_arith_inst_retired.256b_packed_single:u/

好消息：

• 我们终于使用了SIMD！
• 在连续内存的情况下，我们现在快了两倍。

坏消息是我们无法处理非连续内存：步幅视图会崩溃。但我们可以解决这个问题。

第四步：重新添加对步幅内存的支持

我们解决这个问题的方法是让函数有两个版本：

• 一个用于连续内存，它将使用SIMD并运行得更快。
• 另一个将回退到更通用的版本，不使用SIMD。它会慢一些，但至少我们的代码可以运行。

为了减少代码重复，我们将使用Cython的一个称为“融合类型”（fused types）的功能，这是C++模板或Rust泛型的简化版本。基本上，Cython最终会为每种类型生成两个版本的函数。

# 这种类型在编译时将是以下两种底层类型之一：
ctypedef fused float_arr_t:
   cnp.float32_t[:]
   cnp.float32_t[::1]

@cython.boundscheck(False)
@cython.wraparound(False)
cdef average_impl(
    float_arr_t arr1,
    float_arr_t arr2
):
    out = np.empty((len(arr1),), dtype=np.float32)
    cdef cnp.float32_t[::1] out_view = out
    for i in range(len(arr1)):
        out_view[i] = (arr1[i] + arr2[i]) / 2
    return out

def average_arrays_4(arr1, arr2):
    if len(arr1) != len(arr2):
        raise ValueError("Arrays must be the same size")
    if arr1.data.contiguous and arr2.data.contiguous:
        return average_impl[cnp.float32_t[::1]](arr1, arr2)
    else:
        return average_impl[cnp.float32_t[:]](arr1, arr2)

如果我们编译并运行基准测试，我们得到：

Time per call, contiguous: 0.24 ms
Time per call, strided: 4.0 ms

 Performance counter stats for 'python benchmark.py average_arrays_4':

    13,953,180,922      cpu_core/instructions:u/
       412,866,119      cpu_atom/instructions:u/
       500,000,000      cpu_core/fp_arith_inst_retired.128b_packed_single:u/
                 0      cpu_core/fp_arith_inst_retired.256b_packed_single:u/

我们的函数在可以使用SIMD时使用SIMD，在不能使用时回退到较慢的路径。而且它的代码重复最少。

回顾我们的结果：加速及其实现方式

以下是我们不同版本的比较：

版本	连续运行时间（ms）	步幅运行时间（ms）
基线	0.57	4.2
-O3 编译器标志	0.57	4.2
禁用边界检查	0.46	4.1
连续内存的SIMD	0.24	4.0

步幅数组的运行时间差异是否显著尚不清楚，但对于连续数组，我们的代码现在快了两倍。而且这些更改并不是特别复杂。

总结一下，我们必须：

1. 添加-O3编译器标志。这适用于Linux和macOS编译器。在Windows上，默认值可能足够，或者你可能需要设置类似但不同的标志，我还没有检查。
2. 小心禁用边界和环绕检查。这里有一个好的测试套件很有帮助——我们确保测试了无效的输入！
3. 确保Cython可以生成连续数组启用的更简单的代码。

如果你正在维护Cython代码，请考虑进行这些更改！当然，这并不总是足以开始使用SIMD。这是一个简单的例子，自动向量化很可能会起作用，更复杂的代码可能会失败。但在许多情况下，它会使你的代码更快。

如果你无法让自动向量化工作，你可以：

• 启用gcc输出向量化问题：export CFLAGS="-O3 -fopt-info-vec-missed"，然后运行pip install --verbose .，以便查看编译器输出。
• 尝试切换到Clang编译器而不是默认的gcc（export CC=clang）。要获取有关编译器决策过程的提示，请将-Rpass-analysis=loop-vectorize添加到CFLAGS中。