NEP 10 — 优化迭代器/UFunc 性能#

作者:

Mark Wiebe <mwwiebe@gmail.com>

内容类型:

text/x-rst

创建时间:

2010年11月25日

状态:

最终

目录#

摘要#

此 NEP 提议用一个单一的新迭代器替换 NumPy 迭代器和多迭代器,该迭代器设计得更灵活,并允许更缓存友好的数据访问。新迭代器还包含了大部分核心 ufunc 功能,使得在不完全符合 ufunc 模型的情况下,也能轻松获得 ufunc 的好处。主要好处包括:

  • 自动重排以查找缓存友好的访问模式

  • 标准和可自定义的广播

  • 自动类型/字节序/对齐转换

  • 可选的缓冲以最小化转换内存使用

  • 可选的输出数组,未提供时自动分配

  • 自动输出或公共类型选择

此迭代器设计的大部分内容已实现并取得了令人鼓舞的结果。构造开销略大(a.flat: 0.5 us, nditer(a): 1.4 us, broadcast(a,b): 1.4 us, nditer([a,b]): 2.2 us),但是,如示例所示,在纯 Python 代码中结合迭代器已经可以改进内置 NumPy 机制的性能。一个示例重写了 np.add,对于某些 Fortran 连续数组获得了四倍的改进,另一个示例将图像合成代码从 1.4 秒提高到 180 毫秒。

该实现试图考虑到 NumPy 2.0 重构中做出的设计决策,以便将来能相对简单地将其集成到 libndarray 中。

动机#

NumPy 默认从 UFuncs 返回 C 连续数组。当处理结构不同的数据时,这可能导致极差的内存访问模式。一个简单的计时示例说明了这一点,与 Fortran 连续数组相加时性能下降超过八倍。所有计时均使用 NumPy 2.0dev (2010年11月22日) 在 Athlon 64 X2 4200+ 上,64位操作系统上完成。

In [1]: import numpy as np
In [2]: a = np.arange(1000000,dtype=np.float32).reshape(10,10,10,10,10,10)
In [3]: b, c, d = a.copy(), a.copy(), a.copy()

In [4]: timeit a+b+c+d
10 loops, best of 3: 28.5 ms per loop

In [5]: timeit a.T+b.T+c.T+d.T
1 loops, best of 3: 237 ms per loop

In [6]: timeit a.T.ravel('A')+b.T.ravel('A')+c.T.ravel('A')+d.T.ravel('A')
10 loops, best of 3: 29.6 ms per loop

在这种情况下,可以通过切换到内存视图,进行相加,然后再重塑回原样来轻松恢复性能。为了进一步探讨问题并查看它并非总是那么容易解决,让我们考虑 NumPy 中处理图像缓冲区的简单代码。

图像合成示例#

为了一个更真实的例子,让我们考虑一个图像缓冲区。图像通常以 Fortran 连续的顺序存储,颜色通道可以被视为一个结构化的“RGB”类型或一个长度为三的额外维度。由此产生的内存布局既不是 C 连续的,也不是 Fortran 连续的,但由于 ndarray 的灵活性,在 NumPy 中可以直接处理。这似乎是理想的,因为它使内存布局与典型的 C 或 C++ 图像代码兼容,同时在 Python 中提供自然访问。获取像素 (x,y) 的颜色就是“image[x,y]”。

此布局在 NumPy 中的性能非常差。以下代码创建了两个黑色图像,并对它们进行“over”合成操作。

In [9]: image1 = np.zeros((1080,1920,3), dtype=np.float32).swapaxes(0,1)
In [10]: alpha1 = np.zeros((1080,1920,1), dtype=np.float32).swapaxes(0,1)
In [11]: image2 = np.zeros((1080,1920,3), dtype=np.float32).swapaxes(0,1)
In [12]: alpha2 = np.zeros((1080,1920,1), dtype=np.float32).swapaxes(0,1)
In [13]: def composite_over(im1, al1, im2, al2):
   ....:     return (im1 + (1-al1)*im2, al1 + (1-al1)*al2)

In [14]: timeit composite_over(image1,alpha1,image2,alpha2)
1 loops, best of 3: 3.51 s per loop

如果我们放弃方便的布局,而使用 C 连续的默认值,性能大约会提高七倍。

In [16]: image1 = np.zeros((1080,1920,3), dtype=np.float32)
In [17]: alpha1 = np.zeros((1080,1920,1), dtype=np.float32)
In [18]: image2 = np.zeros((1080,1920,3), dtype=np.float32)
In [19]: alpha2 = np.zeros((1080,1920,1), dtype=np.float32)

In [20]: timeit composite_over(image1,alpha1,image2,alpha2)
1 loops, best of 3: 581 ms per loop

但这并非全部,因为事实证明,广播 alpha 通道也会带来性能损失。如果我们使用一个值为 3 的 alpha 通道而不是一个值,我们会得到

In [21]: image1 = np.zeros((1080,1920,3), dtype=np.float32)
In [22]: alpha1 = np.zeros((1080,1920,3), dtype=np.float32)
In [23]: image2 = np.zeros((1080,1920,3), dtype=np.float32)
In [24]: alpha2 = np.zeros((1080,1920,3), dtype=np.float32)

In [25]: timeit composite_over(image1,alpha1,image2,alpha2)
1 loops, best of 3: 313 ms per loop

作为最后的比较,让我们看看当使用一维数组以确保只有一个循环执行计算时,它的性能如何。

In [26]: image1 = np.zeros((1080*1920*3), dtype=np.float32)
In [27]: alpha1 = np.zeros((1080*1920*3), dtype=np.float32)
In [28]: image2 = np.zeros((1080*1920*3), dtype=np.float32)
In [29]: alpha2 = np.zeros((1080*1920*3), dtype=np.float32)

In [30]: timeit composite_over(image1,alpha1,image2,alpha2)
1 loops, best of 3: 312 ms per loop

为了获得参考性能数字,我用 C 语言(小心地使用与 NumPy 相同的编译选项)直接实现了这个简单的操作。如果我模仿 Python 代码的内存分配和布局,性能大约是 0.3 秒,与 NumPy 的性能非常接近。将操作合并到一个通道中将时间减少到大约 0.15 秒。

这个示例的一个小变体是使用一个具有四个通道(1920,1080,4)的单个内存块,而不是单独的图像和 alpha。这在图像处理应用程序中更常见,以下是 C 连续布局下的情况。

In [31]: image1 = np.zeros((1080,1920,4), dtype=np.float32)
In [32]: image2 = np.zeros((1080,1920,4), dtype=np.float32)
In [33]: def composite_over(im1, im2):
   ....:     ret = (1-im1[:,:,-1])[:,:,np.newaxis]*im2
   ....:     ret += im1
   ....:     return ret

In [34]: timeit composite_over(image1,image2)
1 loops, best of 3: 481 ms per loop

要查看本文档底部附近提议的内***迭代器实现可以带来的改进,请参阅提议 API 之后继续的示例。

改进缓存一致性#

为了获得 UFunc 调用的最佳性能,内存读取模式应尽可能规律。现代 CPU 会尝试预测内存读/写模式并提前填充缓存。最可预测的模式是所有输入和输出都以相同的顺序顺序处理。

我建议,默认情况下,UFunc 输出的内存布局应尽可能接近输入的布局。每当存在歧义或不匹配时,它将默认使用 C 连续布局。

为了理解如何实现这一点,我们首先考虑所有输入的步幅,在对形状进行广播归一化后。通过确定一组步幅是否兼容和/或含糊不清,我们可以确定一个输出内存布局,该布局可以最大化一致性。

在广播中,首先将输入形状转换为广播形状,方法是在前面添加单一维度,然后创建广播步幅,其中任何单一维度的步幅都设置为零。

步幅也可以是负数,在某些情况下,这可以归一化以适应以下讨论。如果特定轴的所有步幅都是负数或零,则可以在适当调整基数据指针后,将该维度的步幅取反。

下面是一个示例,说明三个具有 C 连续布局的输入如何导致广播步幅。为了简化,示例使用了 1 的项大小。

输入形状

(5,3,7)

(5,3,1)

(1,7)

广播形状

(5,3,7)

(5,3,1)

(1,1,7)

广播步幅

(21,7,1)

(3,1,0)

(0,0,1)

兼容步幅 — 一组步幅是兼容的,如果存在轴的排列,使得步幅对于集合中的每个步幅都是递减的,不包括零值条目。

上面的示例满足具有恒等排列的定义。在动机图像示例中,如果我们将颜色和 alpha 信息分开或不分开,步幅会略有不同。证明此处兼容性的排列是转置(0,1)。

输入/广播形状

图像 (1920, 1080, 3)

Alpha (1920, 1080, 1)

广播步幅(分开)

(3,5760,1)

(1,1920,0)

广播步幅(合并)

(4,7680,1)

(4,7680,0)

含糊不清的步幅 — 一组兼容的步幅是含糊不清的,如果存在多个轴的排列,使得步幅对于集合中的每个步幅都是递减的,不包括零值条目。

这通常发生在每个轴在一组步幅中都有一个 0 步幅时。最简单的例子是二维的,如下所示。

广播形状

(1,3)

(5,1)

广播步幅

(0,1)

(1,0)

然而,即使没有单个输入强制整个布局,也可能存在不含糊不清的兼容步幅,例如此示例。

广播形状

(1,3,4)

(5,3,1)

广播步幅

(0,4,1)

(3,1,0)

面对含糊不清的情况,我们可以选择完全抛弃步幅兼容的事实,或者尝试通过添加额外的约束来解决含糊不清。我认为合适的选择是解决它,方法是选择最接近 C 连续的内存布局,同时仍与输入步幅兼容。

输出布局选择算法#

我们希望生成的输出 ndarray 内存布局如下:

一致/不含糊不清的步幅

单一一致的布局

一致/含糊不清的步幅

最接近 C 连续的一致布局

不一致的步幅

C 连续

这是计算输出布局排列的算法的伪代码。

perm = range(ndim) # Identity, i.e. C-contiguous
# Insertion sort, ignoring 0-strides
# Note that the sort must be stable, and 0-strides may
# be reordered if necessary, but should be moved as little
# as possible.
for i0 = 1 to ndim-1:
    # ipos is where perm[i0] will get inserted
    ipos = i0
    j0 = perm[i0]
    for i1 = i0-1 to 0:
        j1 = perm[i1]
        ambig, shouldswap = True, False
        # Check whether any strides are ordered wrong
        for strides in broadcast_strides:
            if strides[j0] != 0 and strides[j1] != 0:
                if strides[j0] > strides[j1]:
                    # Only set swap if it's still ambiguous.
                    if ambig:
                        shouldswap = True
                else:
                    # Set swap even if it's not ambiguous,
                    # because not swapping is the choice
                    # for conflicts as well.
                    shouldswap = False
                ambig = False
        # If there was an unambiguous comparison, either shift ipos
        # to i1 or stop looking for the comparison
        if not ambig:
            if shouldswap:
                ipos = i1
            else:
                break
    # Insert perm[i0] into the right place
    if ipos != i0:
       for i1 = i0-1 to ipos:
         perm[i1+1] = perm[i1]
       perm[ipos] = j0
# perm is now the closest consistent ordering to C-contiguous
return perm

合并维度#

在许多情况下,内存布局允许使用一维循环来代替跟踪迭代器内的多个坐标。现有代码已经利用了这一点,当数据是 C 连续的时,但由于我们正在重新排序轴,我们可以更普遍地应用此优化。

一旦迭代步幅被排序为单调递减,任何可以合并的维度都将并排。如果对于所有操作数,增加 stride[i+1] * shape[i+1] 次与增加 stride[i] 相同,或者 strides[i+1]*shape[i+1] == strides[i],则维度 i 和 i+1 可以合并成一个单一维度。

这是合并的伪代码。

# Figure out which pairs of dimensions can be coalesced
can_coalesce = [False]*ndim
for strides, shape in zip(broadcast_strides, broadcast_shape):
    for i = 0 to ndim-2:
        if strides[i+1]*shape[i+1] == strides[i]:
            can_coalesce[i] = True
# Coalesce the types
new_ndim = ndim - count_nonzero(can_coalesce)
for strides, shape in zip(broadcast_strides, broadcast_shape):
    j = 0
    for i = 0 to ndim-1:
        # Note that can_coalesce[ndim-1] is always False, so
        # there is no out-of-bounds access here.
        if can_coalesce[i]:
            shape[i+1] = shape[i]*shape[i+1]
        else:
            strides[j] = strides[i]
            shape[j] = shape[i]
            j += 1

内部循环特化#

特化完全由内部循环函数处理,因此此优化与其它优化无关。一些特化已经完成,例如用于 reduce 操作。这个想法在 https://projects.scipy.org.cn/numpy/wiki/ProjectIdeas 中被提及,“使用内建指令(SSE 指令)来加速 NumPy 的底层循环。”

以下是两个参数函数的可能性,涵盖了加法/减法/乘法/除法的关键情况。

  • 第一个或第二个参数是单个值(即 0 步幅值)并且不别名输出。 arr = arr + 1; arr = 1 + arr

    • 可以只加载一次常量,而不是每次都从内存中重新加载

  • 步幅匹配数据类型的大小。例如 C 或 Fortran 连续数据

    • 可以使用简单的循环,无需使用步幅

  • 步幅匹配数据类型的大小,并且它们都对齐 16 字节(或与 16 字节对齐的偏移量相同)

    • 可以使用 SSE 一次处理多个值

  • 第一个输入和输出是相同的单个值(即约简操作)。

    • 这在现有代码中已针对许多 UFunc 进行了特化

上述情况通常不是互斥的,例如常量参数可以与 SSE 结合使用,当步幅匹配数据类型大小时,并且约简也可以使用 SSE 进行优化。

实现细节#

除了内部循环特化之外,讨论的优化显著影响了 ufunc_object.c 以及用于广播的 PyArrayIterObject/PyArrayMultiIterObject。总的来说,当需要时,应该可以模拟当前行为,但我认为默认行为应该是产生和操作能获得最佳性能的内存布局。

为了支持新的缓存友好行为,我们为任何 order= 参数引入了一个新的选项 ‘K’(表示“保持”)。

提议的 ‘order=’ 标志如下:

‘C’

C 连续布局

‘F’

Fortran 连续布局

‘A’

“C” (任何连续) — 如果输入具有 Fortran 连续布局,则为“F”,否则为“C”。

‘K’

“保持布局” — 相当于“C”然后进行某些轴的排列,尽可能接近输入的布局。

或者作为枚举

/* For specifying array memory layout or iteration order */
typedef enum {
        /* Fortran order if inputs are all Fortran, C otherwise */
        NPY_ANYORDER=-1,
        /* C order */
        NPY_CORDER=0,
        /* Fortran order */
        NPY_FORTRANORDER=1,
        /* An order as close to the inputs as possible */
        NPY_KEEPORDER=2
} NPY_ORDER;

也许一个好的策略是先实现这里讨论的功能,而不改变默认值。一旦它们实现并经过充分测试,默认值就可以在所有合适的地方从 order='C' 更改为 order='K'。UFuncs 还应该提供一个 order= 参数来控制其输出的布局。

迭代器可以进行自动转换,并且我已经创建了一系列逐步更宽松的转换规则。也许对于 2.0,NumPy 可以采用这个枚举作为其处理转换的首选方式。

/* For specifying allowed casting in operations which support it */
typedef enum {
        /* Only allow identical types */
        NPY_NO_CASTING=0,
        /* Allow identical and byte swapped types */
        NPY_EQUIV_CASTING=1,
        /* Only allow safe casts */
        NPY_SAFE_CASTING=2,
        /* Allow safe casts and casts within the same kind */
        NPY_SAME_KIND_CASTING=3,
        /* Allow any casts */
        NPY_UNSAFE_CASTING=4
} NPY_CASTING;

迭代器重写#

基于对代码的分析,似乎重构现有的迭代对象来实现这些优化是极其困难的。此外,对迭代器的一些使用需要修改内部值或标志,因此使用迭代器的代码无论如何都必须更改。因此,我们提议创建一个新的迭代器对象,该对象包含现有迭代器的功能并将其扩展以包含优化。

重写迭代器的高级目标包括:

  • 内存占用小,内存分配次数少。

  • 简单情况(如扁平数组)的开销应非常小。

  • 将单个和多个迭代合并到一个对象中。

应提供给用户代码的功能:

  • 以 C、Fortran 或“最快”(默认)顺序进行迭代。

  • 如果需要,跟踪 C 样式或 Fortran 样式的扁平索引(现有迭代器始终跟踪 C 样式索引)。这可以独立于迭代顺序完成。

  • 如果需要,跟踪坐标(现有迭代器需要手动更改内部迭代器标志来保证这一点)。

  • 跳过最后一个内部维度的迭代,以便它可以由内部循环处理。

  • 跳转到数组中的特定坐标。

  • 迭代任意子集轴(以支持例如一次约简多个轴)。

  • 如果提供了 NULL 输入,则自动分配输出参数的能力。这些输出应具有与迭代顺序匹配的内存布局,并且是 order='K' 支持的机制。

  • 自动复制和/或缓冲不满足类型/字节序/对齐要求的输入。无论进行何种缓冲或复制,调用者的迭代内部循环都应相同。

实现注意事项:

  • 用户代码绝不能触碰迭代器内部。这允许将来对内部内存布局进行重大更改,如果发现更高性能的实现策略。

  • 使用函数指针而不是宏进行迭代。这样,可以为常见情况(例如 ndim 小、不同标志设置、迭代数组数量少)创建特化。此外,可以规定一种迭代模式,该模式先复制函数指针,以便编译器可以将函数指针保留在寄存器中。

  • 动态创建内存布局,以最小化迭代器占用的缓存行数(对于 LP64,sizeof(PyArrayIterObject) 约为 2.5KB,像加法这样的二进制运算需要三个这样的对象用于 Multi-Iterator)。

  • 将 C API 对象与 Python 引用计数隔离,以便可以从 C 中自然使用它。Python 对象然后成为 C 迭代器的包装器。这类似于 PEP 3118 的 Py_buffer 和 memoryview 的设计分离。

提议的迭代器内存布局#

以下结构描述了迭代器的内存。所有项都打包在一起,这意味着 flags、ndim 和 niter 的不同值将产生略微不同的布局。

struct {
    /* Flags indicate what optimizations have been applied, and
     * affect the layout of this struct. */
    uint32 itflags;
    /* Number of iteration dimensions.  If FLAGS_HASCOORDS is set,
     * it matches the creation ndim, otherwise it may be smaller.  */
    uint16 ndim;
    /* Number of objects being iterated.  This is fixed at creation time. */
    uint16 niter;

    /* The number of times the iterator will iterate */
    intp itersize;

    /* The permutation is only used when FLAGS_HASCOORDS is set,
     * and is placed here so its position depends on neither ndim
     * nor niter. */
    intp perm[ndim];

    /* The data types of all the operands */
    PyArray_Descr *dtypes[niter];
    /* Backups of the starting axisdata 'ptr' values, to support Reset */
    char *resetdataptr[niter];
    /* Backup of the starting index value, to support Reset */
    npy_intp resetindex;

    /* When the iterator is destroyed, Py_XDECREF is called on all
       these objects */
    PyObject *objects[niter];

    /* Flags indicating read/write status and buffering
     * for each operand. */
    uint8 opitflags[niter];
    /* Padding to make things intp-aligned again */
    uint8 padding[];

    /* If some or all of the inputs are being buffered */
    #if (flags&FLAGS_BUFFERED)
    struct buffer_data {
        /* The size of the buffer, and which buffer we're on.
         * the i-th iteration has i = buffersize*bufferindex+pos
         */
        intp buffersize;
        /* For tracking position inside the buffer */
        intp size, pos;
        /* The strides for the pointers */
        intp stride[niter];
        /* Pointers to the data for the current iterator position.
         * The buffer_data.value ptr[i] equals either
         * axis_data[0].ptr[i] or buffer_data.buffers[i] depending
         * on whether copying to the buffer was necessary.
         */
        char* ptr[niter];
        /* Functions to do the copyswap and casting necessary */
        transferfn_t readtransferfn[niter];
        void *readtransferdata[niter];
        transferfn_t writetransferfn[niter];
        void *writetransferdata[niter];
        /* Pointers to the allocated buffers for operands
         * which the iterator determined needed buffering
         */
        char *buffers[niter];
    };
    #endif /* FLAGS_BUFFERED */

    /* Data per axis, starting with the most-frequently
     * updated, and in decreasing order after that. */
    struct axis_data {
        /* The shape of this axis */
        intp shape;
        /* The current coordinate along this axis */
        intp coord;
        /* The operand and index strides for this axis */
        intp stride[niter];
        #if (flags&FLAGS_HASINDEX)
            intp indexstride;
        #endif
        /* The operand pointers and index values for this axis */
        char* ptr[niter];
        #if (flags&FLAGS_HASINDEX)
            intp index;
        #endif
    }[ndim];
};

axis_data 结构数组按增加的增量更新速度排序。如果 perm 是恒等映射,这意味着它与 C 顺序相反。这样做是为了使最常访问的数据项更靠近结构的前面,其中包含公共属性,从而提高缓存一致性。它还简化了 iternext 调用,同时使 getcoord 和相关函数稍微复杂化。

提议的迭代器 API#

现有的迭代器 API 包括 PyArrayIter_Check、PyArray_Iter* 和 PyArray_ITER_* 等函数。多迭代器数组包括 PyArray_MultiIter*、PyArray_Broadcast 和 PyArray_RemoveSmallest。新的迭代器设计用一个单一的对象和相关的 API 替换了所有这些功能。新 API 的目标之一是,所有对现有迭代器的使用都可以在没有重大工作量的情况下被新迭代器替换。

选择的 C API 命名约定基于 numpy-refactor 分支中的约定,其中 libndarray 的数组名为 NpyArray,函数名为 NpyArray_*。迭代器命名为 NpyIter,函数名为 NpyIter_*

Python 暴露的迭代器命名为 np.nditer。这个迭代器的一个可能的发布策略是发布一个 1.X(1.6?)版本,添加迭代器但不被 NumPy 代码使用。然后,2.0 可以完全集成并发布。如果选择此策略,则应在 1.X 发布之前尽可能多地确定命名约定和 API。名称 np.iter 不能使用,因为它与 Python 内置的 iter 冲突。我建议在 Python 中使用 np.nditer 这个名称,因为它目前未使用。

除了为新迭代器设定的性能目标外,API 似乎还可以进行重构,以更好地支持一些常见的 NumPy 编程习惯。

通过将 UFunc 代码中的部分功能移到迭代器中,应该可以更容易地为希望在不完全符合 UFunc 范例的情况下模拟 UFunc 行为的扩展代码提供支持。特别是,模拟 UFunc 缓冲行为并非易事。

旧 -> 新迭代器 API 转换#

对于常规迭代器:

PyArray_IterNew

NpyIter_New

PyArray_IterAllButAxis

NpyIter_New + axes 参数 **或** 迭代器标志 NPY_ITER_NO_INNER_ITERATION

PyArray_BroadcastToShape

**不支持**(但如果需要,可以支持)

PyArrayIter_Check

需要在 Python 暴露中添加

PyArray_ITER_RESET

NpyIter_Reset

PyArray_ITER_NEXT

NpyIter_GetIterNext 返回的函数指针

PyArray_ITER_DATA

NpyIter_GetDataPtrArray

PyArray_ITER_GOTO

NpyIter_GotoCoords

PyArray_ITER_GOTO1D

NpyIter_GotoIndex

PyArray_ITER_NOTDONE

iternext 函数指针的返回值

对于多迭代器:

PyArray_MultiIterNew

NpyIter_MultiNew

PyArray_MultiIter_RESET

NpyIter_Reset

PyArray_MultiIter_NEXT

NpyIter_GetIterNext 返回的函数指针

PyArray_MultiIter_DATA

NpyIter_GetDataPtrArray

PyArray_MultiIter_NEXTi

**不支持**(始终锁定步进迭代)

PyArray_MultiIter_GOTO

NpyIter_GotoCoords

PyArray_MultiIter_GOTO1D

NpyIter_GotoIndex

PyArray_MultiIter_NOTDONE

iternext 函数指针的返回值

PyArray_Broadcast

NpyIter_MultiNew 处理

PyArray_RemoveSmallest

迭代器标志 NPY_ITER_NO_INNER_ITERATION

对于其他 API 调用:

PyArray_ConvertToCommonType

迭代器标志 NPY_ITER_COMMON_DTYPE

迭代器指针类型#

迭代器结构是内部生成的,但仍然需要一个类型来在将错误类型传递给 API 时提供警告和/或错误。我们通过 typedef 不完整的结构来实现这一点。

typedef struct NpyIter_InternalOnly NpyIter;

构造和析构#

NpyIter* NpyIter_New(PyArrayObject* op, npy_uint32 flags, NPY_ORDER order, NPY_CASTING casting, PyArray_Descr* dtype, npy_intp a_ndim, npy_intp *axes, npy_intp buffersize)

为给定的 NumPy 数组对象 op 创建一个迭代器。

可在 flags 中传递的标志是全局标志和每个操作数标志的任意组合,这些标志在 NpyIter_MultiNew 中进行了记录,但排除了 NPY_ITER_ALLOCATE

可以为 order 传递 NPY_ORDER 枚举值中的任何一个。为了高效迭代,NPY_KEEPORDER 是最佳选择,其他顺序强制执行特定的迭代模式。

可以为 casting 传递 NPY_CASTING 枚举值中的任何一个。这些值包括 NPY_NO_CASTINGNPY_EQUIV_CASTINGNPY_SAFE_CASTINGNPY_SAME_KIND_CASTINGNPY_UNSAFE_CASTING。为了允许转换发生,还必须启用复制或缓冲。

如果 dtype 不是 NULL,则它要求该数据类型。如果允许复制,则在数据可转换为所需类型时,它将创建临时副本。如果启用了 UPDATEIFCOPY,它还将在迭代器销毁时通过另一次转换将数据复制回去。

如果 a_ndim 大于零,则还必须提供 axes。在这种情况下,axes 是一个 a_ndim 大小的 op 轴数组。在 axes 中值为 -1 表示 newaxis。在 axes 数组中,轴不能重复。

如果 buffersize 为零,则使用默认缓冲区大小,否则它指定缓冲区的大小。推荐使用 2 的幂次方的缓冲区,例如 512 或 1024。

如果发生错误,则返回 NULL,否则返回分配的迭代器。

为了创建一个类似于旧迭代器的迭代器,这应该可以工作。

iter = NpyIter_New(op, NPY_ITER_READWRITE,
                    NPY_CORDER, NPY_NO_CASTING, NULL, 0, NULL);

如果你想编辑一个具有对齐 double 代码的数组,但顺序无关紧要,你将使用这个。

dtype = PyArray_DescrFromType(NPY_DOUBLE);
iter = NpyIter_New(op, NPY_ITER_READWRITE |
                    NPY_ITER_BUFFERED |
                    NPY_ITER_NBO,
                    NPY_ITER_ALIGNED,
                    NPY_KEEPORDER,
                    NPY_SAME_KIND_CASTING,
                    dtype, 0, NULL);
Py_DECREF(dtype);

NpyIter* NpyIter_MultiNew(npy_intp niter, PyArrayObject** op, npy_uint32 flags, NPY_ORDER order, NPY_CASTING casting, npy_uint32 *op_flags, PyArray_Descr** op_dtypes, npy_intp oa_ndim, npy_intp **op_axes, npy_intp buffersize)

为广播 op 中提供的 niter 个数组对象创建迭代器。

对于正常用法,请将 oa_ndim 设置为 0,并将 op_axes 设置为 NULL。有关这些参数的说明,请参阅下文,它们允许自定义手动广播以及重排和省略轴。

可以为 order 传递 NPY_ORDER 枚举值中的任何一个。为了高效迭代,NPY_KEEPORDER 是最佳选择,其他顺序强制执行特定的迭代模式。使用 NPY_KEEPORDER 时,如果还想确保迭代不沿轴反转,则应传递标志 NPY_ITER_DONT_NEGATE_STRIDES

可以为 casting 传递 NPY_CASTING 枚举值中的任何一个。这些值包括 NPY_NO_CASTINGNPY_EQUIV_CASTINGNPY_SAFE_CASTINGNPY_SAME_KIND_CASTINGNPY_UNSAFE_CASTING。为了允许转换发生,还必须启用复制或缓冲。

如果 op_dtypes 不是 NULL,它为每个 op[i] 指定一个数据类型或 NULL

oa_ndim 非零时,它指定将通过自定义广播进行迭代的维度数。如果提供了它,则还必须提供 op_axes。这两个参数允许您详细控制操作数数组的轴如何匹配和迭代。在 op_axes 中,您必须提供一个 niter 指针数组,指向 oa_ndim 大小的 npy_intp 类型数组。如果 op_axes 中的某个条目为 NULL,则将应用正常的广播规则。在 op_axes[j][i] 中存储的是 op[j] 的一个有效轴,或者 -1,表示 newaxis。在每个 op_axes[j] 数组中,轴不能重复。以下示例说明了正常广播如何应用于 3 维数组、2 维数组、1 维数组和标量。

npy_intp oa_ndim = 3;               /* # iteration axes */
npy_intp op0_axes[] = {0, 1, 2};    /* 3-D operand */
npy_intp op1_axes[] = {-1, 0, 1};   /* 2-D operand */
npy_intp op2_axes[] = {-1, -1, 0};  /* 1-D operand */
npy_intp op3_axes[] = {-1, -1, -1}  /* 0-D (scalar) operand */
npy_intp *op_axes[] = {op0_axes, op1_axes, op2_axes, op3_axes};

如果 buffersize 为零,则使用默认缓冲区大小,否则它指定缓冲区的大小。推荐使用 2 的幂次方的缓冲区,例如 512 或 1024。

如果发生错误,则返回 NULL,否则返回分配的迭代器。

可在 flags 中传递的标志(应用于整个迭代器)是:

NPY_ITER_C_INDEX, NPY_ITER_F_INDEX

使迭代器跟踪匹配 C 或 Fortran 顺序的索引。这些选项是互斥的。

NPY_ITER_COORDS

使迭代器跟踪数组坐标。这会阻止迭代器合并轴以产生更大的内部循环。

NPY_ITER_NO_INNER_ITERATION

使迭代器跳过内部循环的迭代,允许迭代器的用户来处理它。

此标志与 NPY_ITER_C_INDEXNPY_ITER_F_INDEXNPY_ITER_COORDS 不兼容。

NPY_ITER_DONT_NEGATE_STRIDES

这仅在为 order 参数指定 NPY_KEEPORDER 时影响迭代器。默认情况下,使用 NPY_KEEPORDER,迭代器会反转具有负步幅的轴,以便内存以正向遍历。此标志禁用此步骤。如果要使用轴的底层内存顺序,但又不想反转轴,则使用此标志。例如,numpy.ravel(a, order='K') 的行为就是如此。

NPY_ITER_COMMON_DTYPE

使迭代器将所有操作数转换为公共数据类型,该类型根据 ufunc 类型提升规则计算。每个操作数的标志必须设置,以便允许适当的转换,并且必须启用复制或缓冲。

如果公共数据类型是提前已知的,则不要使用此标志。而是为所有操作数设置所需的数据类型。

NPY_ITER_REFS_OK

指示可以接受并使用包含引用类型的数组(对象数组或包含对象类型的结构化数组)。如果启用了此标志,调用者必须确保检查 NpyIter_IterationNeedsAPI(iter) 是否为 true,在这种情况下,在迭代期间可能不会释放 GIL。

NPY_ITER_ZEROSIZE_OK

指示应允许大小为零的数组。由于典型的迭代循环不能自然地处理零大小的数组,因此您必须在进入迭代循环之前检查 IterSize 是否非零。

NPY_ITER_REDUCE_OK

允许可写操作数具有零步幅且尺寸大于一的维度。请注意,此类操作数必须是读/写。

启用缓冲时,还会切换到特殊的缓冲模式,该模式会根据需要减小循环长度,以免覆盖正在归约的值。

请注意,如果要对自动分配的输出进行归约,必须使用 NpyIter_GetOperandArray 获取其引用,然后在执行迭代循环之前将每个值设置为归约单元。在缓冲归约的情况下,这意味着还必须指定 NPY_ITER_DELAY_BUFALLOC 标志,然后在初始化分配的运算数后重置迭代器以准备缓冲。

NPY_ITER_RANGED

启用对 iterindex 范围 [0, NpyIter_IterSize(iter)) 的子范围的迭代支持。使用 NpyIter_ResetToIterIndexRange 函数来指定迭代范围。

此标志只能与 NPY_ITER_NO_INNER_ITERATION 结合使用,当 NPY_ITER_BUFFERED 启用时。这是因为没有缓冲时,内部循环的长度始终是最小迭代维度的长度,允许其被分割需要特殊处理,实际上使其更像缓冲版本。

NPY_ITER_BUFFERED

导致迭代器存储缓冲数据,并使用缓冲来满足数据类型、对齐和字节序要求。要缓冲运算数,请勿指定 NPY_ITER_COPYNPY_ITER_UPDATEIFCOPY 标志,因为它们会覆盖缓冲。缓冲对于使用迭代器的 Python 代码尤其有用,它允许一次处理更大的数据块,以分摊 Python 解释器开销。

如果与 NPY_ITER_NO_INNER_ITERATION 一起使用,由于跨距的布局方式,调用者的内部循环可能会获得比没有缓冲时更大的数据块。

请注意,如果运算数被赋予 NPY_ITER_COPYNPY_ITER_UPDATEIFCOPY 标志,将优先创建副本而不是缓冲。当数组被广播以至于需要复制元素以获得恒定跨距时,缓冲仍会发生。

在正常缓冲中,每个内部循环的大小等于缓冲大小,或者如果指定了 NPY_ITER_GROWINNER,则可能更大。如果启用了 NPY_ITER_REDUCE_OK 并且发生了归约,内部循环可能会变小,具体取决于归约的结构。

NPY_ITER_GROWINNER

当启用缓冲时,这允许内部循环的大小在不需要缓冲时增长。此选项最适用于执行数据通篇遍历的情况,而不是对每个内部循环使用小型缓存友好型临时值数组的情况。

NPY_ITER_DELAY_BUFALLOC

当启用缓冲时,这会将缓冲区的分配延迟到调用 NpyIter_Reset* 函数之一为止。此标志的存在是为了避免在创建缓冲迭代器的多个副本以进行多线程迭代时,浪费性地复制缓冲区数据。

此标志的另一个用途是用于设置归约操作。创建迭代器后,如果迭代器自动分配了归约输出(请确保使用 READWRITE 访问权限),则可以将其值初始化为归约单元。使用 NpyIter_GetOperandArray 获取对象。然后,调用 NpyIter_Reset 来分配缓冲区并用其初始值填充。

可能传递给 op_flags[i] 的标志,其中 0 <= i < niter

NPY_ITER_READWRITENPY_ITER_READONLYNPY_ITER_WRITEONLY

指示迭代器用户将如何读取或写入 op[i]。每个运算数必须指定其中一个标志。

NPY_ITER_COPY

如果 op[i] 不满足构造函数标志和参数指定的数据类型或对齐要求,则允许创建其副本。

NPY_ITER_UPDATEIFCOPY

触发 NPY_ITER_COPY,并且当数组运算数被标记为写入并被复制时,在迭代器销毁时会将副本中的数据复制回 op[i]

如果运算数被标记为仅写入并且需要副本,则会创建一个未初始化的临时数组,并在销毁时将其复制回 op[i],而不是执行不必要的复制操作。

NPY_ITER_NBONPY_ITER_ALIGNEDNPY_ITER_CONTIG

导致迭代器为 op[i] 提供本机字节序、根据 dtype 要求对齐、连续的数据,或它们的任何组合。

默认情况下,迭代器会生成指向提供的数组的指针,这些指针可能已对齐或未对齐,并具有任何字节序。如果未启用复制或缓冲,并且运算数数据不满足约束条件,则会引发错误。

连续约束仅适用于内部循环,后续的内部循环可能具有任意的指针更改。

如果请求的数据类型为非本机字节序,则 NBO 标志会覆盖它,并且请求的数据类型将转换为本机字节序。

NPY_ITER_ALLOCATE

这用于输出数组,并且要求设置 NPY_ITER_WRITEONLY 标志。如果 op[i] 为 NULL,则创建一个具有最终广播维度的新数组,其布局与迭代器的迭代顺序匹配。

op[i] 为 NULL 时,请求的数据类型 op_dtypes[i] 也可能为 NULL,在这种情况下,它将从标记为可读的数组的 dtype 中自动生成。生成 dtype 的规则与 UFuncs 相同。特别要注意的是对所选 dtype 的字节序的处理。如果只有一个输入,则使用输入的 dtype。否则,如果组合了多个输入 dtype,则输出将为本机字节序。

使用此标志分配后,调用者可以通过调用 NpyIter_GetOperandArray 并获取返回的 C 数组中的第 i 个对象来检索新数组。调用者必须对它调用 Py_INCREF 以声明对数组的引用。

NPY_ITER_NO_SUBTYPE

NPY_ITER_ALLOCATE 一起使用,此标志禁用分配数组子类型作为输出,强制其成为标准的 ndarray。

TODO:也许引入一个函数 NpyIter_GetWrappedOutput 并移除此标志会更好?

NPY_ITER_NO_BROADCAST

确保输入或输出完全匹配迭代维度。

NPY_ITER_WRITEABLE_REFERENCES

默认情况下,如果迭代器具有涉及 Python 引用的写操作数,则迭代器在创建时会失败。添加此标志表示使用迭代器的代码已意识到此可能性并正确处理。

NpyIter *NpyIter_Copy(NpyIter *iter)

创建给定迭代器的副本。此函数主要用于启用数据的多线程迭代。

TODO:将此移至关于多线程迭代的部分。

多线程迭代的推荐方法是首先使用 NPY_ITER_NO_INNER_ITERATIONNPY_ITER_RANGEDNPY_ITER_BUFFEREDNPY_ITER_DELAY_BUFALLOC 以及可能的 NPY_ITER_GROWINNER 标志创建迭代器。为每个线程创建此迭代器的副本(第一个迭代器除外)。然后,获取迭代索引范围 [0, NpyIter_GetIterSize(iter)) 并将其拆分为任务,例如使用 TBB parallel_for 循环。当一个线程获得要执行的任务时,它通过调用 NpyIter_ResetToIterIndexRange 并迭代完整范围来使用其迭代器副本。

在多线程代码或不持有 Python GIL 的代码中使用迭代器时,必须小心只调用在该上下文中安全的函数。NpyIter_Copy 不能在没有 Python GIL 的情况下安全调用,因为它会增加 Python 引用。通过将 errmsg 参数设置为非 NULL,可以安全地调用 Reset* 和其他一些函数,以便函数通过它传递错误而不是设置 Python 异常。

int NpyIter_UpdateIter(NpyIter *iter, npy_intp i, npy_uint32 op_flags, NPY_CASTING casting, PyArray_Descr *dtype) **未实现**

在迭代器中更新第 i 个运算数,使其可能具有新的数据类型或更严格的标志属性。这的一种用途是允许自动分配基于标准 NumPy 类型提升规则确定输出数据类型,然后使用此函数在处理过程中将输入以及可能的自动输出转换为不同的数据类型。

此操作只能在将 NPY_ITER_COORDS 作为标志传递给迭代器时进行。如果不需要坐标,请在不再需要调用 NpyIter_UpdateIter 时调用 NpyIter_RemoveCoords() 函数。

如果第 i 个运算数已被复制,则会引发错误。为避免这种情况,请将所有标志(读取/写入指示符除外)留给稍后将调用 NpyIter_UpdateIter 的任何运算数。

可以在 op_flags 中传递的标志是 NPY_ITER_COPYNPY_ITER_UPDATEIFCOPYNPY_ITER_NBONPY_ITER_ALIGNEDNPY_ITER_CONTIG

int NpyIter_RemoveAxis(NpyIter *iter, npy_intp axis)

从迭代中移除一个轴。这要求在创建迭代器时设置了 NPY_ITER_COORDS,并且在启用缓冲或跟踪索引时不起作用。此函数还会将迭代器重置到其初始状态。

这对于设置累积循环等很有用。迭代器可以首先用所有维度(包括累积轴)创建,以便正确创建输出。然后,可以移除累积轴,并以嵌套方式执行计算。

警告:此函数可能会更改迭代器的内部内存布局。必须重新获取任何缓存的函数或指针!

返回 NPY_SUCCEEDNPY_FAIL

int NpyIter_RemoveCoords(NpyIter *iter)

如果迭代器具有坐标,则此函数会剥离对它们的[支持],并进行进一步的迭代器优化(如果不需要坐标)。此函数还会将迭代器重置到其初始状态。

警告:此函数可能会更改迭代器的内部内存布局。必须重新获取任何缓存的函数或指针!

调用此函数后,NpyIter_HasCoords(iter) 将返回 false。

返回 NPY_SUCCEEDNPY_FAIL

int NpyIter_RemoveInnerLoop(NpyIter *iter)

如果使用了 UpdateIter/RemoveCoords,您可能希望指定 NPY_ITER_NO_INNER_ITERATION 标志。此标志不允许与 NPY_ITER_COORDS 同时使用,因此此函数用于在调用 NpyIter_RemoveCoords 之后启用该功能。此函数还会将迭代器重置到其初始状态。

警告:此函数会更改迭代器的内部逻辑。必须重新获取任何缓存的函数或指针!

返回 NPY_SUCCEEDNPY_FAIL

int NpyIter_Deallocate(NpyIter *iter)

释放迭代器对象。这还将释放所有创建的副本,并在必要时触发 UPDATEIFCOPY 行为。

返回 NPY_SUCCEEDNPY_FAIL

int NpyIter_Reset(NpyIter *iter, char **errmsg)

将迭代器重置回初始状态,位于迭代范围的开头。

返回 NPY_SUCCEEDNPY_FAIL。如果 errmsg 非 NULL,则在返回 NPY_FAIL 时不会设置 Python 异常。相反,*errmsg 被设置为错误消息。当 errmsg 非 NULL 时,可以在不持有 Python GIL 的情况下安全地调用该函数。

int NpyIter_ResetToIterIndexRange(NpyIter *iter, npy_intp istart, npy_intp iend, char **errmsg)

重置迭代器并将其限制在 iterindex 范围 [istart, iend)。有关如何为多线程迭代使用此功能,请参阅 NpyIter_Copy。这要求将 NPY_ITER_RANGED 标志传递给迭代器构造函数。

如果要同时重置 iterindex 范围和基指针,可以执行以下操作以避免额外的缓冲区复制(复制此代码时请确保添加返回码错误检查)。

/* Set to a trivial empty range */
NpyIter_ResetToIterIndexRange(iter, 0, 0);
/* Set the base pointers */
NpyIter_ResetBasePointers(iter, baseptrs);
/* Set to the desired range */
NpyIter_ResetToIterIndexRange(iter, istart, iend);

返回 NPY_SUCCEEDNPY_FAIL。如果 errmsg 非 NULL,则在返回 NPY_FAIL 时不会设置 Python 异常。相反,*errmsg 被设置为错误消息。当 errmsg 非 NULL 时,可以在不持有 Python GIL 的情况下安全地调用该函数。

int NpyIter_ResetBasePointers(NpyIter *iter, char **baseptrs, char **errmsg)

将迭代器重置回初始状态,但使用 baseptrs 中的值作为数据,而不是被迭代数组中的指针。此函数旨在与 op_axes 参数一起,由嵌套迭代代码使用两个或多个迭代器。

返回 NPY_SUCCEEDNPY_FAIL。如果 errmsg 非 NULL,则在返回 NPY_FAIL 时不会设置 Python 异常。相反,*errmsg 被设置为错误消息。当 errmsg 非 NULL 时,可以在不持有 Python GIL 的情况下安全地调用该函数。

TODO:将以下内容移至关于嵌套迭代器的特殊部分。

创建用于嵌套迭代的迭代器需要一些注意事项。所有迭代器运算数必须完全匹配,否则 NpyIter_ResetBasePointers 的调用将无效。这意味着不应随意使用自动复制和输出分配。仍然可以通过使用所有转换参数启用的一个迭代器来使用迭代器的自动数据转换和转换功能,然后使用 NpyIter_GetOperandArray 函数获取分配的运算数,并将它们传递给其余迭代器的构造函数。

警告:在为嵌套迭代创建迭代器时,代码不得在不同迭代器中使用同一个维度两次。如果这样做,嵌套迭代在迭代期间将产生越界指针。

警告:在为嵌套迭代创建迭代器时,缓冲只能应用于最内层的迭代器。如果将缓冲迭代器用作 baseptrs 的源,它将指向一个小缓冲区而不是数组,并且内部迭代将无效。

使用嵌套迭代器的模式如下:

NpyIter *iter1, *iter1;
NpyIter_IterNext_Fn iternext1, iternext2;
char **dataptrs1;

/*
 * With the exact same operands, no copies allowed, and
 * no axis in op_axes used both in iter1 and iter2.
 * Buffering may be enabled for iter2, but not for iter1.
 */
iter1 = ...; iter2 = ...;

iternext1 = NpyIter_GetIterNext(iter1);
iternext2 = NpyIter_GetIterNext(iter2);
dataptrs1 = NpyIter_GetDataPtrArray(iter1);

do {
    NpyIter_ResetBasePointers(iter2, dataptrs1);
    do {
        /* Use the iter2 values */
    } while (iternext2(iter2));
} while (iternext1(iter1));

int NpyIter_GotoCoords(NpyIter *iter, npy_intp *coords)

将迭代器调整为指向 coords 指向的 ndim 坐标。如果未跟踪坐标、坐标越界或禁用了内部循环迭代,则会返回错误。

返回 NPY_SUCCEEDNPY_FAIL

int NpyIter_GotoIndex(NpyIter *iter, npy_intp index)

将迭代器调整为指向指定的 index。如果迭代器使用 NPY_ITER_C_INDEX 标志构造,则 index 是 C 顺序索引;如果迭代器使用 NPY_ITER_F_INDEX 标志构造,则 index 是 Fortran 顺序索引。如果未跟踪索引、索引越界或禁用了内部循环迭代,则会返回错误。

返回 NPY_SUCCEEDNPY_FAIL

npy_intp NpyIter_GetIterSize(NpyIter *iter)

返回正在迭代的元素数量。这是形状中所有维度的乘积。

npy_intp NpyIter_GetReduceBlockSizeFactor(NpyIter *iter) **未实现**

这提供了一个因子,该因子必须整除用于范围迭代以安全地进行归约多线程处理的块大小。如果迭代器没有归约,则返回 1。

使用范围迭代进行归约多线程处理时,有两种可能的归约方法:

如果存在一个到小输出的大归约,为每个线程创建一个初始化为归约单元的临时数组,然后让每个线程将其归约到其临时数组中。完成后,将临时数组组合在一起。您可以通过观察到 NpyIter_GetReduceBlockSizeFactor 返回一个大的值(例如 NpyIter_GetIterSize 的一半或三分之一)来检测这种情况。您还应该检查输出是否很小以确保这一点。

如果有许多到大输出的小归约,并且归约维度是内部维度,则 NpyIter_GetReduceBlockSizeFactor 将返回一个小的数字,只要您选择的多线程处理块大小为 NpyIter_GetReduceBlockSizeFactor(iter)*n(其中 n 是一些整数),操作就是安全的。

坏的情况是当归约维度是迭代器中的最外层循环时。例如,如果您有一个形状为 (3,1000,1000) 的 C 顺序数组,并在维度 0 上进行归约,那么对于 NPY_KEEPORDERNPY_CORDER 迭代顺序,NpyIter_GetReduceBlockSizeFactor 将返回一个等于 NpyIter_GetIterSize 的大小。虽然这对 CPU 缓存不利,但将来可能会提供另一种顺序,也许是 NPY_REDUCEORDER,它将归约轴推到内部循环,但否则与 NPY_KEEPORDER 相同。

npy_intp NpyIter_GetIterIndex(NpyIter *iter)

获取迭代器的 iterindex,这是一个与迭代器迭代顺序匹配的索引。

void NpyIter_GetIterIndexRange(NpyIter *iter, npy_intp *istart, npy_intp *iend)

获取正在迭代的 iterindex 子范围。如果未指定 NPY_ITER_RANGED,则此函数始终返回范围 [0, NpyIter_IterSize(iter))

int NpyIter_GotoIterIndex(NpyIter *iter, npy_intp iterindex)

将迭代器调整为指向指定的 iterindex。IterIndex 是一个与迭代器迭代顺序匹配的索引。如果 iterindex 超出范围、启用了缓冲或禁用了内部循环迭代,则会返回错误。

返回 NPY_SUCCEEDNPY_FAIL

int NpyIter_HasInnerLoop(NpyIter *iter)

如果迭代器处理内部循环,则返回 1,否则返回 0,调用者需要处理。这由构造函数标志 NPY_ITER_NO_INNER_ITERATION 控制。

int NpyIter_HasCoords(NpyIter *iter)

如果迭代器使用 NPY_ITER_COORDS 标志创建,则返回 1,否则返回 0。

int NpyIter_HasIndex(NpyIter *iter)

如果迭代器使用 NPY_ITER_C_INDEXNPY_ITER_F_INDEX 标志创建,则返回 1,否则返回 0。

int NpyIter_IsBuffered(NpyIter *iter)

如果迭代器使用 NPY_ITER_BUFFERED 标志创建,则返回 1,否则返回 0。

int NpyIter_IsGrowInner(NpyIter *iter)

如果迭代器使用 NPY_ITER_GROWINNER 标志创建,则返回 1,否则返回 0。

npy_intp NpyIter_GetBufferSize(NpyIter *iter)

如果迭代器被缓冲,则返回正在使用的缓冲区的大小,否则返回 0。

npy_intp NpyIter_GetNDim(NpyIter *iter)

返回正在迭代的维度数。如果迭代器构造函数未请求坐标,则此值可能小于原始对象的维度数。

npy_intp NpyIter_GetNIter(NpyIter *iter)

返回正在迭代的对象数量。

npy_intp *NpyIter_GetAxisStrideArray(NpyIter *iter, npy_intp axis)

获取指定轴的跨距数组。需要迭代器正在跟踪坐标,并且未启用缓冲。

当您想以某种方式匹配运算数轴,然后使用 NpyIter_RemoveAxis 移除它们以手动处理其处理时,可以使用此功能。通过在移除轴之前调用此函数,您可以获得手动处理的跨距。

错误时返回 NULL

int NpyIter_GetShape(NpyIter *iter, npy_intp *outshape)

outshape 中返回迭代器的广播形状。这只能在支持坐标的迭代器上调用。

返回 NPY_SUCCEEDNPY_FAIL

PyArray_Descr **NpyIter_GetDescrArray(NpyIter *iter)

这会返回一个指向正在迭代的对象 niter 的数据类型描述符的指针。结果指向 iter,因此调用者不会获得描述符的任何引用。

此指针可以在迭代循环之前缓存,调用 iternext 不会改变它。

PyObject **NpyIter_GetOperandArray(NpyIter *iter)

这会返回一个指向正在迭代的运算数 niter PyObjects 的指针。结果指向 iter,因此调用者不会获得 PyObjects 的任何引用。

PyObject *NpyIter_GetIterView(NpyIter *iter, npy_intp i)

这会返回一个指向新 ndarray 视图的引用,该视图是 NpyIter_GetOperandArray() 返回的数组中第 i 个对象的视图,其维度和跨距与内部优化迭代模式匹配。对此视图的 C 顺序迭代等同于迭代器的迭代顺序。

例如,如果迭代器使用单个数组作为其输入创建,并且可以将所有轴重新排列然后将其压缩为单个跨距迭代,则此函数将返回一个一维数组视图。

void NpyIter_GetReadFlags(NpyIter *iter, char *outreadflags)

填充 niter 标志。如果 op[i] 可以读取,则设置 outreadflags[i] 为 1,否则设置为 0。

void NpyIter_GetWriteFlags(NpyIter *iter, char *outwriteflags)

填充 niter 标志。如果 op[i] 可以写入,则设置 outwriteflags[i] 为 1,否则设置为 0。

迭代函数#

NpyIter_IterNext_Fn NpyIter_GetIterNext(NpyIter *iter, char **errmsg)

返回一个迭代函数指针。此函数可能会计算出该函数指针的专用版本,而不是将其存储在迭代器结构中。因此,为了获得良好的性能,要求将函数指针保存在变量中,而不是在每次循环迭代时检索。

如果发生错误,则返回 NULL。如果 errmsg 非 NULL,则在返回 NPY_FAIL 时不会设置 Python 异常。相反,*errmsg 被设置为错误消息。当 errmsg 非 NULL 时,可以在不持有 Python GIL 的情况下安全地调用该函数。

典型的循环结构如下:

NpyIter_IterNext_Fn iternext = NpyIter_GetIterNext(iter, NULL);
char **dataptr = NpyIter_GetDataPtrArray(iter);

do {
    /* use the addresses dataptr[0], ... dataptr[niter-1] */
} while(iternext(iter));

当指定 NPY_ITER_NO_INNER_ITERATION 时,典型的内部循环结构如下:

NpyIter_IterNext_Fn iternext = NpyIter_GetIterNext(iter, NULL);
char **dataptr = NpyIter_GetDataPtrArray(iter);
npy_intp *stride = NpyIter_GetInnerStrideArray(iter);
npy_intp *size_ptr = NpyIter_GetInnerLoopSizePtr(iter), size;
npy_intp iiter, niter = NpyIter_GetNIter(iter);

do {
    size = *size_ptr;
    while (size--) {
        /* use the addresses dataptr[0], ... dataptr[niter-1] */
        for (iiter = 0; iiter < niter; ++iiter) {
            dataptr[iiter] += stride[iiter];
        }
    }
} while (iternext());

注意,我们使用的是迭代器内部的 dataptr 数组,而不是将值复制到局部临时变量。这是可能的,因为当调用 iternext() 时,这些指针将被新的值覆盖,而不是增量更新。

如果使用的是编译时固定缓冲区(NPY_ITER_BUFFEREDNPY_ITER_NO_INNER_ITERATION 标志均已设置),则内部大小也可以作为信号。当 iternext() 返回 false 时,大小保证变为零,从而启用以下循环结构。请注意,如果使用此结构,则不应将 NPY_ITER_GROWINNER 作为标志传递,因为它会在某些情况下导致更大的尺寸。

/* The constructor should have buffersize passed as this value */
#define FIXED_BUFFER_SIZE 1024

NpyIter_IterNext_Fn iternext = NpyIter_GetIterNext(iter, NULL);
char **dataptr = NpyIter_GetDataPtrArray(iter);
npy_intp *stride = NpyIter_GetInnerStrideArray(iter);
npy_intp *size_ptr = NpyIter_GetInnerLoopSizePtr(iter), size;
npy_intp i, iiter, niter = NpyIter_GetNIter(iter);

/* One loop with a fixed inner size */
size = *size_ptr;
while (size == FIXED_BUFFER_SIZE) {
    /*
     * This loop could be manually unrolled by a factor
     * which divides into FIXED_BUFFER_SIZE
     */
    for (i = 0; i < FIXED_BUFFER_SIZE; ++i) {
        /* use the addresses dataptr[0], ... dataptr[niter-1] */
        for (iiter = 0; iiter < niter; ++iiter) {
            dataptr[iiter] += stride[iiter];
        }
    }
    iternext();
    size = *size_ptr;
}

/* Finish-up loop with variable inner size */
if (size > 0) do {
    size = *size_ptr;
    while (size--) {
        /* use the addresses dataptr[0], ... dataptr[niter-1] */
        for (iiter = 0; iiter < niter; ++iiter) {
            dataptr[iiter] += stride[iiter];
        }
    }
} while (iternext());

NpyIter_GetCoords_Fn NpyIter_GetGetCoords(NpyIter *iter, char **errmsg)

返回一个用于获取迭代器坐标的函数指针。如果迭代器不支持坐标,则返回 NULL。建议在迭代循环之前将此函数指针缓存到本地变量中。

如果发生错误,则返回 NULL。如果 errmsg 非 NULL,则在返回 NPY_FAIL 时不会设置 Python 异常。相反,*errmsg 被设置为错误消息。当 errmsg 非 NULL 时,可以在不持有 Python GIL 的情况下安全地调用该函数。

char **NpyIter_GetDataPtrArray(NpyIter *iter)

这会返回一个指向 niter 数据指针的数组。如果未指定 NPY_ITER_NO_INNER_ITERATION,则每个数据指针都指向迭代器的当前数据项。如果没有指定内部迭代,它将指向内部循环的第一个数据项。

此指针可以在迭代循环之前缓存,调用 iternext 不会改变它。在不持有 Python GIL 的情况下安全调用此函数。

npy_intp *NpyIter_GetIndexPtr(NpyIter *iter)

这会返回一个指向正在跟踪的索引的指针,如果未跟踪索引,则返回 NULL。只有在构造期间指定了 NPY_ITER_C_INDEXNPY_ITER_F_INDEX 标志之一时,才能使用它。

当使用 NPY_ITER_NO_INNER_ITERATION 标志时,代码需要知道执行内部循环的参数。这些函数提供了该信息。

npy_intp *NpyIter_GetInnerStrideArray(NpyIter *iter)

返回一个指向 niter 跨距数组的指针,每个迭代对象一个,供内部循环使用。

此指针可以在迭代循环之前缓存,调用 iternext 不会改变它。在不持有 Python GIL 的情况下安全调用此函数。

npy_intp* NpyIter_GetInnerLoopSizePtr(NpyIter *iter)

返回指向内部循环应执行的迭代次数的指针。

此地址可以在迭代循环之前缓存,调用 iternext 不会改变它。在迭代过程中,该值本身可能会发生变化,尤其是在启用缓冲的情况下。在不持有 Python GIL 的情况下安全调用此函数。

void NpyIter_GetInnerFixedStrideArray(NpyIter *iter, npy_intp *out_strides)

获取一个在整个迭代过程中固定不变(或不会改变)的步幅数组。对于可能改变的步幅,值NPY_MAX_INTP将被放入步幅中。

一旦迭代器准备好进行迭代(如果使用了NPY_DELAY_BUFALLOC,则在重置之后),调用此函数以获取可用于选择快速内部循环函数的步幅。例如,如果步幅为0,这意味着内部循环可以始终将值加载到一个变量中,然后在整个循环中使用该变量;如果步幅等于元素大小,则可以为该操作数使用连续版本。

可以安全地调用此函数而不持有Python GIL。

示例#

一个使用迭代器的复制函数。order参数用于控制分配结果的内存布局。

如果输入是引用类型,此函数将失败。要解决此问题,代码必须更改为特别处理可写引用,并将NPY_ITER_WRITEABLE_REFERENCES添加到标志中。

/* NOTE: This code has not been compiled/tested */
PyObject *CopyArray(PyObject *arr, NPY_ORDER order)
{
    NpyIter *iter;
    NpyIter_IterNext_Fn iternext;
    PyObject *op[2], *ret;
    npy_uint32 flags;
    npy_uint32 op_flags[2];
    npy_intp itemsize, *innersizeptr, innerstride;
    char **dataptrarray;

    /*
     * No inner iteration - inner loop is handled by CopyArray code
     */
    flags = NPY_ITER_NO_INNER_ITERATION;
    /*
     * Tell the constructor to automatically allocate the output.
     * The data type of the output will match that of the input.
     */
    op[0] = arr;
    op[1] = NULL;
    op_flags[0] = NPY_ITER_READONLY;
    op_flags[1] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE;

    /* Construct the iterator */
    iter = NpyIter_MultiNew(2, op, flags, order, NPY_NO_CASTING,
                            op_flags, NULL, 0, NULL);
    if (iter == NULL) {
        return NULL;
    }

    /*
     * Make a copy of the iternext function pointer and
     * a few other variables the inner loop needs.
     */
    iternext = NpyIter_GetIterNext(iter);
    innerstride = NpyIter_GetInnerStrideArray(iter)[0];
    itemsize = NpyIter_GetDescrArray(iter)[0]->elsize;
    /*
     * The inner loop size and data pointers may change during the
     * loop, so just cache the addresses.
     */
    innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);
    dataptrarray = NpyIter_GetDataPtrArray(iter);

    /*
     * Note that because the iterator allocated the output,
     * it matches the iteration order and is packed tightly,
     * so we don't need to check it like the input.
     */
    if (innerstride == itemsize) {
        do {
            memcpy(dataptrarray[1], dataptrarray[0],
                                    itemsize * (*innersizeptr));
        } while (iternext(iter));
    } else {
        /* Should specialize this further based on item size... */
        npy_intp i;
        do {
            npy_intp size = *innersizeptr;
            char *src = dataaddr[0], *dst = dataaddr[1];
            for(i = 0; i < size; i++, src += innerstride, dst += itemsize) {
                memcpy(dst, src, itemsize);
            }
        } while (iternext(iter));
    }

    /* Get the result from the iterator object array */
    ret = NpyIter_GetOperandArray(iter)[1];
    Py_INCREF(ret);

    if (NpyIter_Deallocate(iter) != NPY_SUCCEED) {
        Py_DECREF(ret);
        return NULL;
    }

    return ret;
}

Python lambda UFunc示例#

为了展示新的迭代器如何允许在纯Python中定义高效的类UFunc函数,我们演示了luf函数,该函数使lambda表达式像UFunc一样工作。这与numexpr库非常相似,但只需要几行代码。

首先,这是luf函数的定义。

def luf(lamdaexpr, *args, **kwargs):
    """Lambda UFunc

        e.g.
        c = luf(lambda i,j:i+j, a, b, order='K',
                            casting='safe', buffersize=8192)

        c = np.empty(...)
        luf(lambda i,j:i+j, a, b, out=c, order='K',
                            casting='safe', buffersize=8192)
    """

    nargs = len(args)
    op = args + (kwargs.get('out',None),)
    it = np.nditer(op, ['buffered','no_inner_iteration'],
            [['readonly','nbo_aligned']]*nargs +
                            [['writeonly','allocate','no_broadcast']],
            order=kwargs.get('order','K'),
            casting=kwargs.get('casting','safe'),
            buffersize=kwargs.get('buffersize',0))
    while not it.finished:
        it[-1] = lamdaexpr(*it[:-1])
        it.iternext()

    return it.operands[-1]

然后,通过使用luf而不是直接的Python表达式,我们可以通过更好的缓存行为获得一些性能。

In [2]: a = np.random.random((50,50,50,10))
In [3]: b = np.random.random((50,50,1,10))
In [4]: c = np.random.random((50,50,50,1))

In [5]: timeit 3*a+b-(a/c)
1 loops, best of 3: 138 ms per loop

In [6]: timeit luf(lambda a,b,c:3*a+b-(a/c), a, b, c)
10 loops, best of 3: 60.9 ms per loop

In [7]: np.all(3*a+b-(a/c) == luf(lambda a,b,c:3*a+b-(a/c), a, b, c))
Out[7]: True

Python加法示例#

迭代器已经大部分编写并暴露给Python。为了了解它的行为,让我们看看我们可以用np.add ufunc做什么。即使不改变NumPy的核心,我们也可以使用迭代器来制作一个更快的add函数。

Python暴露提供了两个迭代接口,一个遵循Python迭代器协议,另一个镜像C风格的do-while模式。原生Python方法在大多数情况下更好,但如果您需要迭代器的坐标或索引,请使用C风格模式。

以下是我们如何使用Python迭代器协议编写iter_add函数。

def iter_add_py(x, y, out=None):
    addop = np.add

    it = np.nditer([x,y,out], [],
                [['readonly'],['readonly'],['writeonly','allocate']])

    for (a, b, c) in it:
        addop(a, b, c)

    return it.operands[2]

以下是同一个函数,但遵循C风格模式。

def iter_add(x, y, out=None):
    addop = np.add

    it = np.nditer([x,y,out], [],
                [['readonly'],['readonly'],['writeonly','allocate']])

    while not it.finished:
        addop(it[0], it[1], it[2])
        it.iternext()

    return it.operands[2]

关于此函数的一些值得注意的点

  • 将np.add缓存为局部变量以减少命名空间查找。

  • 输入是只读的,输出是只写的,如果输出为None,则会自动分配。

  • 使用np.add的out参数来避免额外的复制。

让我们创建一些测试变量,并对该函数以及内置的np.add进行计时。

In [1]: a = np.arange(1000000,dtype='f4').reshape(100,100,100)
In [2]: b = np.arange(10000,dtype='f4').reshape(1,100,100)
In [3]: c = np.arange(10000,dtype='f4').reshape(100,100,1)

In [4]: timeit iter_add(a, b)
1 loops, best of 3: 7.03 s per loop

In [5]: timeit np.add(a, b)
100 loops, best of 3: 6.73 ms per loop

慢一千倍,这显然不太好。迭代器的一个特性是no_inner_iteration标志,它旨在帮助加速内部循环。这与旧迭代器的PyArray_IterAllButAxis是相同的想法,但更智能一些。让我们修改iter_add来使用这个特性。

def iter_add_noinner(x, y, out=None):
    addop = np.add

    it = np.nditer([x,y,out], ['no_inner_iteration'],
                [['readonly'],['readonly'],['writeonly','allocate']])

    for (a, b, c) in it:
        addop(a, b, c)

    return it.operands[2]

性能有了显著提升。

In[6]: timeit iter_add_noinner(a, b)
100 loops, best of 3: 7.1 ms per loop

性能基本与内置函数一样好!事实证明,这是因为迭代器能够合并最后两个维度,从而产生了100次对10000个元素的加法。如果内循环不够大,性能就不会有如此显著的提升。让我们使用c而不是b来查看这一点。

In[7]: timeit iter_add_noinner(a, c)
10 loops, best of 3: 76.4 ms per loop

仍然比七秒好很多,但仍然比内置函数慢十倍以上。在这里,内循环有100个元素,循环了10000次。如果我们用C语言编程,我们的性能就已经和内置函数一样好了,但在Python中开销太大了。

这让我们想到迭代器的另一个特性,即它能够给我们提供对迭代内存的视图。它提供的视图是结构化的,因此在C顺序中处理它们(就像内置的NumPy代码一样)会产生与迭代器本身相同的访问顺序。实际上,我们正在使用迭代器来解决一个好的内存访问模式,然后使用其他的NumPy机制来有效地执行它。让我们再次修改iter_add

def iter_add_itview(x, y, out=None):
    it = np.nditer([x,y,out], [],
                [['readonly'],['readonly'],['writeonly','allocate']])

    (a, b, c) = it.itviews
    np.add(a, b, c)

    return it.operands[2]

现在性能与内置函数非常接近了。

In [8]: timeit iter_add_itview(a, b)
100 loops, best of 3: 6.18 ms per loop

In [9]: timeit iter_add_itview(a, c)
100 loops, best of 3: 6.69 ms per loop

现在让我们回到与新迭代器的最初动机类似的情况。这里是与C内存顺序不同的Fortran内存顺序下的相同计算。

In [10]: a = np.arange(1000000,dtype='f4').reshape(100,100,100).T
In [12]: b = np.arange(10000,dtype='f4').reshape(100,100,1).T
In [11]: c = np.arange(10000,dtype='f4').reshape(1,100,100).T

In [39]: timeit np.add(a, b)
10 loops, best of 3: 34.3 ms per loop

In [41]: timeit np.add(a, c)
10 loops, best of 3: 31.6 ms per loop

In [44]: timeit iter_add_itview(a, b)
100 loops, best of 3: 6.58 ms per loop

In [43]: timeit iter_add_itview(a, c)
100 loops, best of 3: 6.33 ms per loop

可以看到,内置函数的性能显著下降,但我们新编写的add函数保持了几乎相同的性能。作为最后的测试,让我们尝试将几个add函数链接在一起。

In [4]: timeit np.add(np.add(np.add(a,b), c), a)
1 loops, best of 3: 99.5 ms per loop

In [9]: timeit iter_add_itview(iter_add_itview(iter_add_itview(a,b), c), a)
10 loops, best of 3: 29.3 ms per loop

另外,只是为了检查它是否做了同样的事情,

In [22]: np.all(
   ....: iter_add_itview(iter_add_itview(iter_add_itview(a,b), c), a) ==
   ....: np.add(np.add(np.add(a,b), c), a)
   ....: )

Out[22]: True

图像合成示例回顾#

作为动机,我们有一个执行两个图像“叠加”合成操作的示例。现在让我们看看如何用新的迭代器来编写函数。

这里是原始函数之一,供参考,以及一些随机图像数据。

In [5]: rand1 = np.random.random(1080*1920*4).astype(np.float32)
In [6]: rand2 = np.random.random(1080*1920*4).astype(np.float32)
In [7]: image1 = rand1.reshape(1080,1920,4).swapaxes(0,1)
In [8]: image2 = rand2.reshape(1080,1920,4).swapaxes(0,1)

In [3]: def composite_over(im1, im2):
  ....:     ret = (1-im1[:,:,-1])[:,:,np.newaxis]*im2
  ....:     ret += im1
  ....:     return ret

In [4]: timeit composite_over(image1,image2)
1 loops, best of 3: 1.39 s per loop

这是同一个函数,重写为使用新的迭代器。注意添加一个可选的输出参数是多么容易。

In [5]: def composite_over_it(im1, im2, out=None, buffersize=4096):
  ....:     it = np.nditer([im1, im1[:,:,-1], im2, out],
  ....:                     ['buffered','no_inner_iteration'],
  ....:                     [['readonly']]*3+[['writeonly','allocate']],
  ....:                     op_axes=[None,[0,1,np.newaxis],None,None],
  ....:                     buffersize=buffersize)
  ....:     while not it.finished:
  ....:         np.multiply(1-it[1], it[2], it[3])
  ....:         it[3] += it[0]
  ....:         it.iternext()
  ....:     return it.operands[3]

In [6]: timeit composite_over_it(image1, image2)
1 loops, best of 3: 197 ms per loop

性能有了巨大的提升,甚至超过了之前使用纯NumPy和C顺序数组的最佳尝试!通过调整缓冲区大小,我们可以看到速度如何提高,直到达到内部循环中CPU缓存的极限。

In [7]: timeit composite_over_it(image1, image2, buffersize=2**7)
1 loops, best of 3: 1.23 s per loop

In [8]: timeit composite_over_it(image1, image2, buffersize=2**8)
1 loops, best of 3: 699 ms per loop

In [9]: timeit composite_over_it(image1, image2, buffersize=2**9)
1 loops, best of 3: 418 ms per loop

In [10]: timeit composite_over_it(image1, image2, buffersize=2**10)
1 loops, best of 3: 287 ms per loop

In [11]: timeit composite_over_it(image1, image2, buffersize=2**11)
1 loops, best of 3: 225 ms per loop

In [12]: timeit composite_over_it(image1, image2, buffersize=2**12)
1 loops, best of 3: 194 ms per loop

In [13]: timeit composite_over_it(image1, image2, buffersize=2**13)
1 loops, best of 3: 180 ms per loop

In [14]: timeit composite_over_it(image1, image2, buffersize=2**14)
1 loops, best of 3: 192 ms per loop

In [15]: timeit composite_over_it(image1, image2, buffersize=2**15)
1 loops, best of 3: 280 ms per loop

In [16]: timeit composite_over_it(image1, image2, buffersize=2**16)
1 loops, best of 3: 328 ms per loop

In [17]: timeit composite_over_it(image1, image2, buffersize=2**17)
1 loops, best of 3: 345 ms per loop

最后,为了再次检查它是否正常工作,我们可以比较这两个函数。

In [18]: np.all(composite_over(image1, image2) ==
    ...:        composite_over_it(image1, image2))
Out[18]: True

带有NumExpr的图像合成#

作为迭代器的测试,numexpr已经增强,允许使用迭代器而不是其内部广播代码。首先,让我们用numexpr实现合成操作。

In [22]: def composite_over_ne(im1, im2, out=None):
   ....:     ima = im1[:,:,-1][:,:,np.newaxis]
   ....:     return ne.evaluate("im1+(1-ima)*im2")

In [23]: timeit composite_over_ne(image1,image2)
1 loops, best of 3: 1.25 s per loop

这比纯NumPy操作要好,但效果不是很好。切换到numexpr的迭代器版本,我们得到了比使用迭代器的纯Python函数大得多的改进。请注意,这是在双核机器上。

In [29]: def composite_over_ne_it(im1, im2, out=None):
   ....:     ima = im1[:,:,-1][:,:,np.newaxis]
   ....:     return ne.evaluate_iter("im1+(1-ima)*im2")

In [30]: timeit composite_over_ne_it(image1,image2)
10 loops, best of 3: 67.2 ms per loop

In [31]: ne.set_num_threads(1)
In [32]: timeit composite_over_ne_it(image1,image2)
10 loops, best of 3: 91.1 ms per loop