带有 AVX SIMD 的 Cython：代码运行一次正确，但如果需要再次运行则挂起

Posted 2023-02-16

【中文标题】带有 AVX SIMD 的 Cython：代码运行一次正确，但如果需要再次运行则挂起

【英文标题】：Cython with AVX SIMD: Code runs once correctly but hang if needed to run again

【发布时间】：2020-04-01 16:52:30

【问题描述】：

我对 SIMD 内在函数还是很陌生，并且通过 Cython 进行了第二次尝试。经过这里的人的一些帮助（非常感谢），我有这个：

from libc.stdlib cimport malloc, free, calloc
cdef extern from "immintrin.h":  # in this example, we use AVX2
    ctypedef float  __m256
    __m256 _mm256_loadu_ps  (float *__P) nogil  
    __m256 _mm256_add_ps    (__m256 __A, __m256 __B) nogil
    __m256 _mm256_mul_ps    (__m256 __A, __m256 __B) nogil
    __m256 _mm256_fmadd_ps  (__m256 __A, __m256 __B, __m256 __C) nogil
    void   _mm256_storeu_ps (float *__P, __m256 __A) nogil  

cdef float [::1] Example_v2 (float *A, float *B, float *C, int n) :
    ### this example for A and B having more than 8 elements, possibly non-divisible by 8    
    cdef:
        __m256 mA, mB, mC
        float *out = <float*> malloc( n * sizeof(float)) 
        float [::1] out_mem = <float [:n]> out
        int i, j, m = <int>( (n-1) / 8) + 1

    with nogil:
        for i in range(m):
            j = 8*i
            mA = _mm256_loadu_ps( &A[j] )
            mB = _mm256_loadu_ps( &B[j] )
            mC = _mm256_loadu_ps( &C[j] )
            _mm256_storeu_ps( &out[j] , _mm256_fmadd_ps( mA, mB, mC ) )

    return out_mem

def run2(float [::1] A, float [::1] B, float [::1] C):
    return Example_v2( &A[0] , &B[0] , &C[0], A.shape[0] )

如果 n 可以被 8 整除。numpy A*B+C 和 cython 是相同的，这是唯一的好消息。

如果 n 不能被 8 整除。它只能运行一次，给出正确答案，然后我的 Spyder 会弹出下面的消息。第二次运行，不做任何更改将导致停顿/挂起或与第一次运行相同的重复。我不知道还有什么方法可以解决它。

另一个相关问题：有没有更好的方法来编写“with nogil”下的部分？ N = 400 的 Numpy 仍然更快，在循环（100,000）中运行通过以上。叹息……

  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\_pswindows.py", line 679, in wrapper
    return fun(self, *args, **kwargs)
  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\_pswindows.py", line 933, in create_time
    user, system, created = cext.proc_times(self.pid)
ProcessLookupError: [Errno 3] No such process (originated from GetExitCodeProcess != STILL_ACTIVE)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\__init__.py", line 373, in _init
    self.create_time()
  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\__init__.py", line 723, in create_time
    self._create_time = self._proc.create_time()
  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\_pswindows.py", line 681, in wrapper
    raise convert_oserror(err, pid=self.pid, name=self._name)
psutil.NoSuchProcess: psutil.NoSuchProcess process no longer exists (pid=9296)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\beng_\Anaconda3\lib\site-packages\qtconsole\manager.py", line 27, in poll
    super(QtKernelRestarter, self).poll()
  File "C:\Users\beng_\Anaconda3\lib\site-packages\jupyter_client\restarter.py", line 113, in poll
    self.kernel_manager.restart_kernel(now=True, newports=newports)
  File "C:\Users\beng_\Anaconda3\lib\site-packages\jupyter_client\manager.py", line 411, in restart_kernel
    self.shutdown_kernel(now=now, restart=True)
  File "C:\Users\beng_\Anaconda3\lib\site-packages\jupyter_client\manager.py", line 371, in shutdown_kernel
    self._kill_kernel()
  File "C:\Users\beng_\Anaconda3\lib\site-packages\spyder\plugins\ipythonconsole\utils\manager.py", line 78, in _kill_kernel
    self.kill_proc_tree(self.kernel.pid)
  File "C:\Users\beng_\Anaconda3\lib\site-packages\spyder\plugins\ipythonconsole\utils\manager.py", line 44, in kill_proc_tree
    parent = psutil.Process(pid)
  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\__init__.py", line 346, in __init__
    self._init(pid)
  File "C:\Users\beng_\Anaconda3\lib\site-packages\psutil\__init__.py", line 386, in _init
    raise NoSuchProcess(pid, None, msg)
psutil.NoSuchProcess: psutil.NoSuchProcess no process found with pid 9296

【问题讨论】：

NumPy 的速度有多快，您使用了什么编译器和哪些选项？ （在什么 CPU 上？）使用aligned_alloc 获得至少 32 字节对齐的内存是否有帮助？如果您受限于内存带宽，可能不会。

这能回答你的问题吗？ How dangerous is it to access an array out of bounds?

当 n 不是 8 的倍数时，你的 m 太大了。

@ead：当 n 不是 8 的倍数时，比如 n=26，所以我可以做 3 组 8 组来解决前 24 组，然后我还有两个元素要做但是这不是 8 的倍数，因此这是问题所在，因为我不知道如何处理它。什么是最佳实践？我是否将其零填充为 8 的倍数？退出循环并返回到非 SIMD 来处理附加元素？还是有更好的 SIMD 风格流程来处理这种情况？我认为 loadu 和 storeu 可以处理此类问题，但似乎没有。

@PeterCordes ：对于 n=400 并循环通过 100,000 ......这是 0.089 秒（numpy）与 0.25 秒。 i9-9900K（目前，一旦流程正确并且我可以真正加快速度，它将在 linux 集群上进行测试）。无论如何，现在关心的是让它在自己的桌面上工作。 msvc 的编译器选项是 /arch:AVX2 和 /O2。我会阅读aligned_alloc 并让你知道。谢谢。

【参考方案1】：

您只能安全地使用 SIMD 内部代码直到数据的前 m 位，其中 m 是 s 的最大倍数小于数据中的位数，其中 s 是 SIMD 数据类型的大小（例如 s = __m256 类型的 256 位）。因此，在对浮点数组（sizeof(float) = 4 字节或 32 位）执行 SIMD 操作并使用 __m256 的情况下，您可以安全地在第一个 l//8（向下舍入）上使用 SIMD，其中 l 是数组中的浮点数。

选项 1：应在不使用 SIMD 的情况下处理剩余数据。 选项 2：如果可以重复指令而不影响最终结果，您可以让最终的 SIMD 指令针对有效数据的最后 s 位重复一次。 选项 3：您可以将数据复制到具有更大 l' 大小的新浮点数组中，其中 l' 是 8 的倍数，无关的最后一个 @ 987654325@ 浮点结果被忽略。

您在使用 SIMD 内部函数时遇到的问题不一定是 Cython 的限制。我想如果你直接在 C 中使用 SIMD 内在函数而不解决上面的内存安全问题，你也会得到未定义的行为。

您还应该了解内存对齐、指针别名以及目标系统是否有这些指令的影响。

选项 4：通过提供正确的编译器选项（例如，为 msvc 提供“/arch:AVX2”，为 gcc 提供“-mavx -ftree-vectorize”），您可以在Cython 或 C，无需编写任何 SIMD 内部代码。您可能必须确保编译器知道指针没有别名。一种方法是首先将数据复制到本地声明的浮点数组。另一种方法是小心使用 C 中的 restrict 限定符来保证指针不会被别名。 Cython 不（还？）识别 restrict 限定符。

使用 Cython，您可以在 .pyx 或 .py 文件顶部声明带有标题注释的编译器选项。例如：

# distutils: extra_compile_args = -mavx -ftree-vectorize

【讨论】：

选项 4：如果您的数组大于一个向量，您有时可以安排在最后一个元素的末尾使用 ends 的未对齐向量来执行最后一个部分块.对于奇数大小的数组，这将与主循环中的最后一个向量部分重叠。这适用于像c[i] = a[i]+b[i] 这样的目标不与任何输入重叠的东西，因此您只需在重叠元素中进行相同的计算。或者对于像 a[i] |= 1 这样的事情，重做它是好的，或者如果你可以安排在 循环之前加载最终向量。

Re: 是否允许别名：__m256 之类的 Intel 内部类型 允许别名任何其他内容，即使是 float 以外的内容：Is `reinterpret_cast`ing between hardware SIMD vector pointer and the corresponding type an undefined behavior?。还与处理最终（部分）向量有关：Vectorizing with unaligned buffers: using VMASKMOVPS: generating a mask from a misalignment count? Or not using that insn at allvmaskmovps 也是标量清理的替代方法，如果您找不到更好的方法。

但是，当您可以保证编译器不混叠时，auto-vectorization 更有可能，或者至少更有效。现代 gcc/clang 将在自动矢量化时经常生成代码来检查重叠，否则会退回到标量循环。这里的总体优点，+1。此外，除了-mavx 之外，您通常还想使用-mtune=znver2 或-mtune=haswell。 (Why doesn't gcc resolve _mm256_loadu_pd as single vmovupd?)。或者只是使用 -march=native 为您自己的机器构建并启用它支持的所有内容，包括 FMA。

@PeterCordes 别名会影响代码的运行方式并给出错误的结果。但我不希望这个答案是关于矢量化的完整教程，所以我只说编码器应该意识到这些影响。我会将您的选项 4 添加到答案中。