广播的 NumPy 算法 - 为什么一种方法的性能如此之高？

Broadcasted NumPy arithmetic - why is one method so much more performant?

This question is a follow up to my answer in .

设置如下：

x = np.arange(5000)  # an integer array
N = 4

现在，我将以两种不同的方式计算 Vandermonde matrix：

m1 = (x ** np.arange(N)[:, None]).T

而且，

m2 = x[:, None] ** np.arange(N)

完整性检查：

np.array_equal(m1, m2)
True

这些方法是相同的，但它们的性能不同：

%timeit m1 = (x ** np.arange(N)[:, None]).T
42.7 µs ± 271 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit m2 = x[:, None] ** np.arange(N)
150 µs ± 995 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

因此，第一种方法尽管需要在最后进行换位，但仍然比第二种方法快 3 倍。

唯一的区别是第一种情况下广播的是smaller数组，而第二种情况是larger ].

So, with a fairly decent understanding of how numpy works, I can guess that the answer would involve the cache. The first method is a lot more cache friendly than the second. However, I'd like an official word from someone with more experience than me.

时间上出现这种鲜明对比的原因是什么？

虽然恐怕我的结论不会比你的更基础 ("probably caching")，但我相信我可以通过一组更本地化的测试来帮助我们集中注意力。

考虑您的示例问题：

M,N = 5000,4
x1 = np.arange(M)
y1 = np.arange(N)[:,None]
x2 = np.arange(M)[:,None]
y2 = np.arange(N)
x1_bc,y1_bc = np.broadcast_arrays(x1,y1)
x2_bc,y2_bc = np.broadcast_arrays(x2,y2)
x1_cont,y1_cont,x2_cont,y2_cont = map(np.ascontiguousarray,
                                      [x1_bc,y1_bc,x2_bc,y2_bc])

如你所见，我定义了一堆数组来比较。 x1、y1 和 x2、y2 分别对应于您的原始测试用例。 ??_bc 对应于这些数组的显式广播版本。这些与原始数据共享数据，但它们具有明确的 0 步幅以获得适当的形状。最后，??_cont 是这些广播数组的连续版本，就好像用 np.tile.

构造的一样

所以 x1_bc、y1_bc、x1_cont 和 y1_cont 的形状都是 (4, 5000)，但是前两者的步幅为零，而后者两个是连续的数组。对于所有意图和目的，利用这些对应数组对中的任何一个的力量应该给我们相同的连续结果（正如 hpaulj 在评论中指出的那样，转置本身基本上是免费的，所以我将忽略最外层的转置以下）。

以下是与您的原始支票相对应的时间：

In [143]: %timeit x1 ** y1
     ...: %timeit x2 ** y2
     ...: 
52.2 µs ± 707 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
96 µs ± 858 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

以下是显式广播数组的时间安排：

In [144]: %timeit x1_bc ** y1_bc
     ...: %timeit x2_bc ** y2_bc
     ...: 
54.1 µs ± 906 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
99.1 µs ± 1.51 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each

同样的事情。这告诉我，差异不是由于从索引表达式到广播数组的转换而造成的。这在很大程度上是预料之中的，但检查一下也无妨。

最后，连续数组：

In [146]: %timeit x1_cont ** y1_cont
     ...: %timeit x2_cont ** y2_cont
     ...: 
38.9 µs ± 529 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
45.6 µs ± 390 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

很大一部分差异消失了！

那我为什么要检查这个？如果您在 python 中沿大尾随维度使用向量化操作，则可以使用 CPU 缓存的一般经验法则。更具体地说，对于行优先 ("C-order") 数组，尾随维度是连续的，而对于列优先 ("fortran-order") 数组，前导维度是连续的。对于足够大的尺寸，arr.sum(axis=-1) 应该比 arr.sum(axis=0) 快，对于行主 numpy 数组给出或采用一些细则。

这里发生的是两个维度之间存在巨大差异（分别为 4 和 5000），但是两个转置情况之间的 巨大 性能不对称只发生对于广播案例。我公认的挥手印象是广播使用 0 步来构建适当大小的视图。这些 0 步意味着在更快的情况下，长 x 数组的内存访问看起来像这样：

[mem0,mem1,mem2,...,mem4999, mem0,mem1,mem2,...,mem4999, ...] # and so on

其中 mem* 仅表示 float64 值 x 位于 RAM 中的某处。将此与我们使用形状 (5000,4):

的较慢情况进行比较 
[mem0,mem0,mem0,mem0, mem1,mem1,mem1,mem1, mem2,mem2,mem2,mem2, ...]

我天真的想法是，使用前者允许 CPU 一次缓存更大块的 x 的单个值，因此性能非常好。在后一种情况下，0 步使 CPU 在 x 的同一内存地址上跳转四次，每次跳转 5000 次。我发现有理由相信此设置不利于缓存，从而导致整体性能不佳。这也与连续情况未显示此性能影响的事实一致：CPU 必须与所有 5000*4 个唯一 float64 值一起使用，并且缓存可能不会受到阻碍通过这些奇怪的阅读。

我也想看看broadcast_arrays:

In [121]: X,Y = np.broadcast_arrays(np.arange(4)[:,None], np.arange(1000))
In [122]: timeit X+Y
10.1 µs ± 31.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [123]: X,Y = np.broadcast_arrays(np.arange(1000)[:,None], np.arange(4))
In [124]: timeit X+Y
26.1 µs ± 30.6 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [125]: X.shape, X.strides
Out[125]: ((1000, 4), (4, 0))
In [126]: Y.shape, Y.strides
Out[126]: ((1000, 4), (0, 4))

np.ascontiguousarray 将 0 跨步维度转换为完整维度

In [132]: Y1 = np.ascontiguousarray(Y)
In [134]: Y1.strides
Out[134]: (16, 4)
In [135]: X1 = np.ascontiguousarray(X)
In [136]: X1.shape
Out[136]: (1000, 4)

使用完整数组的操作速度更快：

In [137]: timeit X1+Y1
4.66 µs ± 161 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

因此，使用 0 跨步数组会产生某种时间损失，即使它没有首先显式扩展数组也是如此。成本与形状有关，也可能与扩展的维度有关。

我不相信缓存真的是这里最有影响力的因素。

我也不是训练有素的计算机科学家，所以我很可能是错的，但让我向您介绍几个观察结果。 为简单起见，我使用@hpaulj 的调用，即“+”显示与“**”基本相同的效果。

我的工作假设是外循环的开销，我认为它比连续的可向量化最内层循环要昂贵得多。

所以让我们首先尽量减少重复的数据量，这样缓存不太可能产生太大影响：

>>> from timeit import repeat
>>> import numpy as np
>>> 
>>> def mock_data(k, N, M):
...     x = list(np.random.randint(0, 10000, (k, N, M)))
...     y = list(np.random.randint(0, 10000, (k, M)))
...     z = list(np.random.randint(0, 10000, (k, N, 1)))
...     return x, y, z
...   
>>> k, N, M = 500, 5000, 4
>>>
>>> repeat('x.pop() + y.pop()', setup='x, y, z = mock_data(k, M, N)', globals=globals(), number=k)
[0.017986663966439664, 0.018148145987652242, 0.018077059998176992]
>>> repeat('x.pop() + y.pop()', setup='x, y, z = mock_data(k, N, M)', globals=globals(), number=k)
[0.026680009090341628, 0.026304758968763053, 0.02680662798229605]

这两种情况都有连续的数据，相同的添加次数，但具有 5000 次外部迭代的版本要慢得多。当我们恢复缓存时，尽管在试验中差异保持大致相同，但比率变得更加明显：

>>> repeat('x[0] + y[0]', setup='x, y, z = mock_data(k, M, N)', globals=globals(), number=k)
[0.011324503924697638, 0.011121788993477821, 0.01106808998156339]
>>> repeat('x[0] + y[0]', setup='x, y, z = mock_data(k, N, M)', globals=globals(), number=k)
[0.020170683041214943, 0.0202067659702152, 0.020624138065613806]

回到最初的 "outer sum" 场景，我们看到在非连续的长维度情况下我们变得更糟。由于我们不必读取比连续场景中更多的数据，因此不能用未缓存数据来解释。

>>> repeat('z.pop() + y.pop()', setup='x, y, z = mock_data(k, M, N)', globals=globals(), number=k)
[0.013918839977122843, 0.01390116906259209, 0.013737019035033882]
>>> repeat('z.pop() + y.pop()', setup='x, y, z = mock_data(k, N, M)', globals=globals(), number=k)
[0.0335254140663892, 0.03351909795310348, 0.0335453050211072]

进一步从跨试验缓存中获益：

>>> repeat('z[0] + y[0]', setup='x, y, z = mock_data(k, M, N)', globals=globals(), number=k)
[0.012061356916092336, 0.012182610924355686, 0.012071475037373602]
>>> repeat('z[0] + y[0]', setup='x, y, z = mock_data(k, N, M)', globals=globals(), number=k)
[0.03265167598146945, 0.03277428599540144, 0.03247103898320347]

从缓存器的角度来看，这充其量是不确定的。

所以让我们看看源代码。 从 tarball 构建当前的 NumPy 后，您会在树中的某处找到一个名为 'loops.c' 的文件中价值近 15000 行的计算机生成代码。这些循环是 ufunc 的最内层循环，与我们的情况最相关的部分似乎是 

#define BINARY_LOOP\
    char *ip1 = args[0], *ip2 = args[1], *op1 = args[2];\
    npy_intp is1 = steps[0], is2 = steps[1], os1 = steps[2];\
    npy_intp n = dimensions[0];\
    npy_intp i;\
    for(i = 0; i < n; i++, ip1 += is1, ip2 += is2, op1 += os1)

/*
 * loop with contiguous specialization
 * op should be the code working on `tin in1`, `tin in2` and
 * storing the result in `tout * out`
 * combine with NPY_GCC_OPT_3 to allow autovectorization
 * should only be used where its worthwhile to avoid code bloat
 */
#define BASE_BINARY_LOOP(tin, tout, op) \
    BINARY_LOOP { \
        const tin in1 = *(tin *)ip1; \
        const tin in2 = *(tin *)ip2; \
        tout * out = (tout *)op1; \
        op; \
    }

etc.

我们案例中的有效载荷似乎足够精简，尤其是如果我正确解释了关于连续专业化和自动矢量化的评论。现在，如果我们只进行 4 次迭代，开销与有效负载的比率开始看起来有点麻烦，而且还不止于此。

在文件 ufunc_object.c 中，我们找到以下片段

/*
 * If no trivial loop matched, an iterator is required to
 * resolve broadcasting, etc
 */

NPY_UF_DBG_PRINT("iterator loop\n");
if (iterator_loop(ufunc, op, dtypes, order,
                buffersize, arr_prep, arr_prep_args,
                innerloop, innerloopdata) < 0) {
    return -1;
}

return 0;

实际循环看起来像

    NPY_BEGIN_THREADS_NDITER(iter);

    /* Execute the loop */
    do {
        NPY_UF_DBG_PRINT1("iterator loop count %d\n", (int)*count_ptr);
        innerloop(dataptr, count_ptr, stride, innerloopdata);
    } while (iternext(iter));

    NPY_END_THREADS;

innerloop就是我们上面看到的内循环。 iternext 有多少开销？

为此，我们需要转到文件 nditer_templ.c.src，我们在其中找到 

/*NUMPY_API
 * Compute the specialized iteration function for an iterator
 *
 * If errmsg is non-NULL, it should point to a variable which will
 * receive the error message, and no Python exception will be set.
 * This is so that the function can be called from code not holding
 * the GIL.
 */
NPY_NO_EXPORT NpyIter_IterNextFunc *
NpyIter_GetIterNext(NpyIter *iter, char **errmsg)
{

etc.

这个函数 returns 一个函数指针，指向

预处理所做的事情之一 
/* Specialized iternext (@const_itflags@,@tag_ndim@,@tag_nop@) */
static int
npyiter_iternext_itflags@tag_itflags@_dims@tag_ndim@_iters@tag_nop@(
                                                      NpyIter *iter)
{

etc.

解析这超出了我的范围，但无论如何它是一个函数指针，必须在外循环的每次迭代中调用，据我所知函数指针不能内联，所以与 4 次迭代相比一个简单的循环体，这将是重要的。

我应该大概分析一下，但我的技能不够。