在 numpy 中广播对 3D 数组的操作的时间缩放

Question

我正在尝试在两个 3D 数组上广播“>”的简单操作。一个具有尺寸 (m, 1, n)，另一个具有尺寸 (1, m, n)。如果我更改第三维 (n) 的值，我会天真地期望计算速度会缩放为 n。

但是，当我尝试明确测量这一点时，我发现当 n 从 1 增加到 2 时，计算时间增加了大约 10 倍，之后缩放是线性的。

为什么从 n=1 到 n=2 时计算时间会急剧增加？我假设它是 numpy 中的内存管理工件，但我正在寻找更多细节。

代码附在下面，结果图。

import numpy as np
import time
import matplotlib.pyplot as plt

def compute_time(n):

    x, y = (np.random.uniform(size=(1, 1000, n)), 
            np.random.uniform(size=(1000, 1, n)))

    t = time.time()
    x > y 
    return time.time() - t

a = [
        [
            n, np.asarray([compute_time(n) 
            for _ in range(100)]).mean()
        ]
        for n in range(1, 30, 1)
    ]

a = np.asarray(a)
plt.plot(a[:, 0], a[:, 1])
plt.xlabel('n')
plt.ylabel('time(ms)')
plt.show()

广播操作的时间图

Answer 1

我无法证明这一点，但我很确定这是由于一项仅在 n==1 时可用的简单优化。

目前，numpy ufunc 实现基于计算机生成的代码，用于映射到一个简单的 C 循环的最内层循环。封闭循环需要使用完全成熟的迭代器对象，该对象取决于有效负载，即最内层循环的大小和原子操作的成本可能是很大的开销。

现在，在 n==1 处，问题本质上是 2D（numpy 足够聪明，可以检测到），最里面的循环大小为 1000，因此迭代器对象有 1000 步。从 n==2 向上，最内层循环的大小为 n，我们有 1,000,000 步的迭代器对象，这说明了您正在观察的跳跃。

正如我所说，我无法证明它，但我可以让它看起来合理：如果我们将变量维度移到前面，那么最内层循环的大小恒定为 1000，而外层循环在 1000 个迭代步骤中线性增长.确实，这让跳跃消失了。

代码：

import numpy as np
import time
import matplotlib.pyplot as plt

def compute_time(n, axis=2):
    xs, ys = [1, 10], [10, 1]
    xs.insert(axis, n)
    ys.insert(axis, n)
    x, y = (np.random.uniform(size=xs),
            np.random.uniform(size=ys))

    t = time.perf_counter()
    x > y
    return time.perf_counter() - t

a = [
        [
            n,
            np.asarray([compute_time(n) for _ in range(100)]).mean(),
            np.asarray([compute_time(n, 0) for _ in range(100)]).mean()
        ]
        for n in range(0, 10, 1)
     ]

a = np.asarray(a)
plt.plot(a[:, 0], a[:, 1:])
plt.xlabel('n')
plt.ylabel('time(ms)')
plt.show()

相关：https://stackoverflow.com/a/48257213/7207392

Answer 2

@Paul 的理论非常正确。在这个答案中，我使用perf 和调试器深入研究以支持这个理论。

首先，让我们看看运行时间都花在了哪些地方（具体代码请参见下面的 run.py 清单）。

对于n=1，我们看到以下内容：

Event count (approx.): 3388750000
Overhead  Command  Shared Object                               Symbol                                                               
  34,04%  python   umath.cpython-36m-x86_64-linux-gnu.so       [.] DOUBLE_less
  32,71%  python   multiarray.cpython-36m-x86_64-linux-gnu.so  [.] _aligned_strided_to_contig_size8_srcstride0
  28,16%  python   libc-2.23.so                                [.] __memmove_ssse3_back
   1,46%  python   multiarray.cpython-36m-x86_64-linux-gnu.so  [.] PyArray_TransferNDimToStrided

与n=2相比：

Event count (approx.): 28954250000                                                              
Overhead  Command  Shared Object                               Symbol                                                               
  40,85%  python   libc-2.23.so                                [.] __memmove_ssse3_back
  40,16%  python   multiarray.cpython-36m-x86_64-linux-gnu.so  [.] PyArray_TransferNDimToStrided
   8,61%  python   umath.cpython-36m-x86_64-linux-gnu.so       [.] DOUBLE_less
   8,41%  python   multiarray.cpython-36m-x86_64-linux-gnu.so  [.] _contig_to_contig

对于 n=2，计数的事件数增加了 8.5 倍，但只有两倍的数据，因此我们需要解释 4 的减速因素。

另一个重要观察：n=2 和（不太明显）n=1（_aligned_strided_to_contig_size8_srcstride0 都是关于复制数据）的运行时间主要是内存操作，它们超过了比较成本 - @987654337 @。

很明显，PyArray_TransferNDimtoStrided 被两种尺寸都调用了，那么为什么它的运行时间份额会有如此大的差异呢？

PyArray_TransferNDimtoStrided显示的self-time不是复制所需的时间，而是开销：调整指针，使得在最后一个维度可以通过stransfer一次性复制：

 PyArray_TransferNDimToStrided(npy_intp ndim,
 ....
 /* A loop for dimensions 0 and 1 */
 for (i = 0; i < shape1; ++i) {
    if (shape0 >= count) {
        stransfer(dst, dst_stride, src, src_stride0,
                    count, src_itemsize, data);
        return 0;
    }
    else {
        stransfer(dst, dst_stride, src, src_stride0,
                    shape0, src_itemsize, data);
    }
    count -= shape0;
    src += src_stride1;
    dst += shape0*dst_stride;
}
...

这些 stransfer-functions 是 _aligned_strided_to_contig_size8_srcstride0（请参阅下面列表中生成的代码）和 _contig_to_contig：

• _contig_to_contig 用于 n=2 的情况下并传递 2-doubles（最后一个维度有 2 个值），调整指针的开销非常高！
• _aligned_strided_to_contig_size8_srcstride0 用于 n=1 并每次调用传输 1000 个双精度数（正如 @Paul 指出的，我们很快就会看到，numpy 足够聪明地丢弃 1 个元素长的维度），调整指针的开销可以忽略。

顺便说一句，为了使用现代 CPU 的矢量化，使用这些函数而不是简单的 for 循环：在编译时已知步幅，编译器能够对代码进行矢量化（编译器通常无法为步幅做这些）仅在运行时知道），因此 numpy 分析访问模式并分派给不同的预编译函数。

还有一个问题：numpy 真的会丢弃最后一个维度，如果它的大小为 1，正如我们的观察所表明的那样？

使用调试器很容易验证：

• 一个通过迭代器访问数据的ufunc，它是在iterator_loop通过NpyIter_AdvancedNew创建的
• in NpyIter_AdvancedNew，维度在npyiter_coalesce_axes 时被分析（并重新解释）

至于将n=2 与n=1 进行比较时“丢失”的速度因子4：它没有特殊含义，只是我机器上的一个随机值：从10 更改矩阵的维度^3 到 10^4 会将优势进一步转移（开销更少）甚至更远到n=1-case，这导致我的机器上的速度损失因子为 12。

---

运行.py

import sys
import numpy as np

n=int(sys.argv[1])

x, y = (np.random.uniform(size=(1, 1000, n)), 
        np.random.uniform(size=(1000, 1, n)))

for _ in range(10000):
    y<x

然后：

perf record python run.py 1
perf report
....
perf record python run.py 2
perf report

---

_aligned_strided_to_contig_size8_srcstride0的生成源：

/*
 * specialized copy and swap for source stride 0,
 * interestingly unrolling here is like above is only marginally profitable for
 * small types and detrimental for >= 8byte moves on x86
 * but it profits from vectorization enabled with -O3
 */
#if (0 == 0) && 1
static NPY_GCC_OPT_3 void
_aligned_strided_to_contig_size8_srcstride0(char *dst,
                        npy_intp dst_stride,
                        char *src, npy_intp NPY_UNUSED(src_stride),
                        npy_intp N, npy_intp NPY_UNUSED(src_itemsize),
                        NpyAuxData *NPY_UNUSED(data))
{
#if 8 != 16
#  if !(8 == 1 && 1)
    npy_uint64 temp;
#  endif
#else
    npy_uint64 temp0, temp1;
#endif
    if (N == 0) {
        return;
    }
#if 1 && 8 != 16
    /* sanity check */
    assert(npy_is_aligned(dst, _ALIGN(npy_uint64)));
    assert(npy_is_aligned(src, _ALIGN(npy_uint64)));
#endif
#if 8 == 1 && 1
    memset(dst, *src, N);
#else

#  if 8 != 16
    temp = _NPY_NOP8(*((npy_uint64 *)src));
#  else
#    if 0 == 0
        temp0 = (*((npy_uint64 *)src));
        temp1 = (*((npy_uint64 *)src + 1));
#    elif 0 == 1
        temp0 = _NPY_SWAP8(*((npy_uint64 *)src + 1));
        temp1 = _NPY_SWAP8(*((npy_uint64 *)src));
#    elif 0 == 2
        temp0 = _NPY_SWAP8(*((npy_uint64 *)src));
        temp1 = _NPY_SWAP8(*((npy_uint64 *)src + 1));
#    endif
#  endif

    while (N > 0) {
#  if 8 != 16
        *((npy_uint64 *)dst) = temp;
#  else
        *((npy_uint64 *)dst) = temp0;
        *((npy_uint64 *)dst + 1) = temp1;
#  endif
#  if 1
        dst += 8;
#  else
        dst += dst_stride;
#  endif
        --N;
    }
#endif/* @elsize == 1 && 1 -- else */
}
#endif/* (0 == 0) && 1 */