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

Question

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

对于复杂的操作，例如在矩阵A、B上应用快速傅立叶变换，这些不是高速缓存优化的，这也是显而易见的。

然而，问题是中间的操作，例如np.where，np.apply_along_axis (这可能没有优化)，np.einsum (这可能是优化的)，np.vstack，等等。我如何知道给定的numpy函数是否针对缓存命中进行了优化，以及它是否比两个嵌套的for-loops更快？

Answer 1

首先，在[numpy] simd上搜索一下，就会找到一些答案。

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

点击率较高的最近问题是How is numpy so fast?。但答案主要涉及c++比较代码及其内存使用。所以它并没有解决numpy的使用问题。

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

首先，您是否了解numpy数组是如何存储的，以及它与列表有何不同？如果不知道这一区别，许多麻木的速度讨论将很难理解。一般来说，使用数组就好像它们是列表一样，带有迭代和列表理解的数组将会更慢。在numpy函数中使用列表会导致速度损失，因为必须首先将列表转换为数组。

此外，object数据类型数组存储对象引用的数据，因此它们的计算以列表理解的速度进行。快速的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在cython中使用了nditer。

np.vstack

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