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 默认情况下从 UFunc 返回 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”类型或长度为 3 的额外维度。生成的内存布局既不是 C 连续也不是 Fortran 连续,但由于 ndarray 的灵活性,可以直接在 NumPy 中轻松使用。这似乎很理想,因为它使内存布局与典型的 C 或 C++ 图像代码兼容,同时在 Python 中提供自然的访问。获取像素 (x,y) 的颜色只是 'image[x,y]'。

事实证明,这种布局在 NumPy 中的性能非常差。这是一段创建两个黑色图像并在其上执行“叠加”合成操作的代码。

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连续时,现有代码已经利用了这一点,但是由于我们正在重新排序轴,因此我们可以更普遍地应用这种优化。

一旦迭代步长被排序为单调递减,任何可以合并的维度都并排在一起。如果对于所有操作数,按步长strides[i+1]递增shape[i+1]次与按步长strides[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一次处理多个值。

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

    • 现有代码中许多UFuncs已经对此进行了专门化。

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

实现细节#

除了内部循环专门化之外,所讨论的优化会显著影响ufunc_object.c以及用于进行广播的PyArrayIterObject/PyArrayMultiIterObject。通常情况下,应该能够在需要的地方模拟当前的行为,但我认为默认情况下应该生成和操作能够提供最佳性能的内存布局。

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

建议的'order='标志如下:

‘C’

C连续布局

‘F’

Fortran连续布局

‘A’

如果输入具有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,并且像加法这样的二元运算需要三个这样的迭代器用于多迭代器)。

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

建议的迭代器内存布局#

以下结构描述了迭代器内存。所有项目都打包在一起,这意味着标志、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 到新迭代器 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 除外。

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

任何 NPY_CASTING 枚举值都可以传递给 casting。这些值包括 NPY_NO_CASTINGNPY_EQUIV_CASTINGNPY_SAFE_CASTINGNPY_SAME_KIND_CASTINGNPY_UNSAFE_CASTING。要允许进行类型转换,还必须启用复制或缓冲。

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

如果 a_ndim 大于零,则还必须提供 axes。在这种情况下,axes 是一个大小为 a_ndimop 轴数组。 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。请参见以下内容,了解这些参数的描述,这些参数允许自定义手动广播以及重新排序和省略轴。

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

任何 NPY_CASTING 枚举值都可以传递给 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_ndimnpy_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_INDEXNPY_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类型提升规则计算得出。必须设置每个操作数的标志,以便允许适当的转换,并且必须启用复制或缓冲。

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

NPY_ITER_REFS_OK

指示可以接受和使用具有引用类型的数组(包含对象类型的对象数组或结构化数组)。如果启用此标志,则调用者必须确保检查NpyIter_IterationNeedsAPI(iter)是否为真,在这种情况下,它可能不会在迭代期间释放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_BUFFERED时,此标志才能与NPY_ITER_NO_INNER_ITERATION一起使用。这是因为如果没有缓冲,内部循环始终是内部迭代维的大小,并且允许将其分割将需要特殊处理,实际上使其更像缓冲版本。

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,在这种情况下,它将根据标记为可读的数组的dtypes自动生成。生成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引用。Reset*和其他一些函数可以通过将errmsg参数作为非NULL传递来安全地调用,以便函数通过它返回错误而不是设置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 非空,则当返回 NPY_FAIL 时不会设置 Python 异常。而是将 *errmsg 设置为错误消息。当 errmsg 非空时,可以在不持有 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 非空,则当返回 NPY_FAIL 时不会设置 Python 异常。而是将 *errmsg 设置为错误消息。当 errmsg 非空时,可以在不持有 Python GIL 的情况下安全地调用此函数。

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

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

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

待办事项:将以下内容移动到嵌套迭代器的特殊部分。

创建用于嵌套迭代的迭代器需要一些注意。所有迭代器操作数必须完全匹配,否则对 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数据类型 Descr 的指针。结果指向iter,因此调用者不会获得对 Descr 的任何引用。

此指针可以在迭代循环之前缓存,调用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,则表示内部循环始终可以一次将它的值加载到变量中,然后在整个循环中使用该变量;或者如果步幅等于 itemsize,则可以使用该操作数的连续版本。

可以在不持有 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的核心,我们也能使用迭代器来创建一个更快的加法函数。

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

慢了1000倍,这显然不太好。迭代器的一个旨在帮助加快内部循环速度的功能是标志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

它仍然比7秒好得多,但仍然比内置函数慢十倍以上。在这里,内部循环有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

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

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

正如您所看到的,内置函数的性能显著下降,但我们新编写的加法函数基本保持了相同的性能。作为最后的测试,让我们尝试将多个加法链接在一起。

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