numpy - 哪些 numpy 指令是 SIMD 优化的？我怎么知道？

Question

很明显，对于简单的操作，例如A + B，np.sum(A, axis=0)这些是缓存优化的。

A对于复杂的操作，例如在矩阵,上应用 FFT 也很明显B，这些都不是缓存优化的。

但是，问题是针对中间操作的，例如np.where，np.apply_along_axis（这可能未优化）、np.einsum（这可能已优化）np.vstack等。我如何知道给定的 numpy 函数是否针对缓存命中进行了优化，并且它是比两个嵌套循环快for？

标签： numpyvectorization

Answer 1

首先，SO 搜索[numpy] simd会出现许多答案。

一，https://stackoverflow.com/a/45798012/901925，找到了一个src/umath/simd.inc.src文件。它将自己描述为“当前包含基于 amd64、x32 或非通用版本 (CFLAGS=-march=...) 构建的 sse2 函数”。这是低级代码，根据构建，可能会被合并到numpy二进制文件中。这不是您作为 Python 级别的用户将能够检测或控制的东西。

最近有很多点击的问题是 numpy 怎么这么快？. 但答案主要涉及c++比较代码及其内存使用。所以它真的没有解决numpy用法。

但是出于您的目的，真正的问题是操作是否使用编译numpy方法，或者是否使用 Python 级别的迭代和对象。

首先，您了解numpy数组是如何存储的，以及它与列表有何不同？在不知道这种差异的情况下，许多关于 numpy 速度的讨论将难以理解。作为一般规则，使用数组就好像它们是列表一样，使用迭代和列表推导会更慢。并且在numpy函数中使用列表会导致速度损失，因为列表必须首先转换为数组。

dtype 数组也object将它们的数据存储在对象引用中，因此它们的计算以列表理解的速度运行。快速 numpy 方法仅适用于数字 dtypes，可以用c本机类型编译的那些 - 浮点数、整数等。

至于您的示例表达式

A + B  

像这样的运算符被实现为ufunc，它充分利用了数组数据存储。由于它可以与多维数组一起工作，并且使用broadcasting的底层代码相当复杂，而且不是您或我可以轻松阅读的内容。在某些较低级别上，它可能会利用处理器兑现和特殊指令，但这更多是c代码宏和编译器选项的功能。

np.sum(A, axis=0)

sum实际上是 a np.add.reduce，所以上述评论适用。但是对于列表，本机 pythonsum毫不逊色。

np.where

np.nonzero是较简单的编译函数之一。它首先用于np.count_nonzero查找有多少非零元素。它使用它来分配它将返回的数组元组，然后再次遍历参数以填充索引。c它相当快，因为​​它以干净的代码循环遍历数组的数据缓冲区。

np.apply_along_axis

即使与列表推导相比，这也很慢。它必须为每个一维数组调用一次你的函数。这是对 python 函数的重复调用，它比实际的迭代方法花费的时间最多。像这样的函数不会编译你的函数，所以在某种程度上，它们只是 Python 级迭代的覆盖。python代码可供研究。

np.einsum

这是一个复杂的功能，根据输入以不同的方式工作。对于更简单的情况，它只使用np.matmul/@，这可能非常快，具体取决于BLAS您拥有的类似库。多年前，当我为它写了一个补丁时，einsum在.nditercython

np.vstack

这是一个封面np.concatenate。python代码很容易阅读。 concatenate被编译。但是这些函数应该正确使用，带有完整的数组列表。在循环中重复使用比 list 更糟糕append。