哪些 numpy 指令经过 SIMD 优化？我怎么知道？

Question

对于简单的操作来说这是显而易见的，例如

A + B

，

np.sum(A, axis=0)

，这些都是缓存优化的。

对于复杂的操作也是显而易见的，例如对矩阵

A

、

B

应用 FFT，这些不是缓存优化的。

但是，问题是针对中间的操作，例如

np.where

、

np.apply_along_axis

（这可能未优化）、

np.einsum

（这可能已优化）、

np.vstack

等。我如何知道给定的 numpy 函数是否针对缓存命中进行了优化，并且它是比两个嵌套

for

循环更快？

Answer 1

首先在

[numpy] simd

上进行SO搜索会找到很多答案。

一个，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

函数中使用列表会导致速度损失，因为列表必须首先转换为数组。

此外

object

dtype 数组将其数据存储为对象引用，因此它们的计算以列表理解速度运行。 快速 numpy 方法仅适用于数字数据类型，可以使用

c

本机类型（浮点数、整数等）进行编译。

至于你的示例表达式

A + B  

像这样的运算符被实现为

ufunc

，它充分利用了数组数据存储。 由于它可以处理多维数组，并使用

broadcasting

，底层代码非常复杂，不是你或我可以轻松阅读的东西。 在某些较低级别，它可能会利用处理器缓存和特殊指令，但这更多是

c

代码宏和编译器选项的功能。

np.sum(A, axis=0)

sum

实际上是一个

np.add.reduce

，所以上面的评论适用。 但对于列表来说，原生 python

sum

也毫不逊色。

np.where

np.nonzero

是更简单的编译函数之一。 它首先使用

np.count_nonzero

来查找有多少个非零元素。 它使用它来分配将返回的数组元组，然后再次循环参数以填充索引。 它相当快，因为它以干净的

c

代码循环遍历数组的数据缓冲区。

np.apply_along_axis

即使与列表推导式相比，这也很慢。 它必须为每个一维数组调用一次函数。对 python 函数的重复调用花费了最多的时间，比实际的迭代方法花费的时间更多。 像这样的函数不会编译您的函数，因此无论如何它们只是 Python 级别迭代的覆盖。 python代码可供研究。

np.einsum

这是一个复杂的函数，根据输入以不同的方式工作。 对于更简单的情况，它只使用

np.matmul/@

，这可能非常快，具体取决于您拥有的

BLAS

之类的库。 几年前，当我为它编写补丁时，

einsum

在

nditer

中使用了

cython

。

np.vstack

这是

np.concatenate

的封面。 python 代码很容易阅读。

concatenate

已编译。但这些函数应该正确使用，并使用整个数组列表。 在循环中重复使用比列表更糟糕

append

。

Answer 2

问题的标题询问的是 SIMD 优化的 numpy 函数。 Numpy 2.0.0 预计将于 2024 年夏季发布，并将包含一个新函数来准确回答这个问题：

numpy.lib.introspect.opt_func_info()

对于指定的函数和数据类型，该函数在运行时返回“当前”支持的 SIMD 扩展。

在文中，您询问缓存优化和缓存命中，我不知道有任何文档指定给定函数的任何缓存优化。