NumPy C 代码解释#

狂热在于当你忘记目标时,你反而加倍努力。— 乔治·桑塔亚那

权威人士就是能告诉你更多你根本不在乎的信息的人。— 佚名

本页旨在解释一些新代码背后的逻辑。这些解释的目的是使人们能够比仅仅查看代码更容易地理解实现背后的思想。或许通过这种方式,算法可以得到更多人的改进、借鉴和/或优化。

内存模型#

ndarray 的一个基本方面是,数组被视为从某个位置开始的内存“块”。对这段内存的解释取决于步长信息。对于 \(N\) 维数组中的每个维度,一个整数(步长)决定了必须跳过多少字节才能到达该维度中的下一个元素。除非你有一个单段数组,否则在遍历数组时必须参考这个步长信息。编写接受步长的代码并不难,你只需使用 char* 指针,因为步长是以字节为单位的。还要记住,步长不必是元素大小的单位倍数。此外,如果数组的维度数为 0(有时称为 rank-0 数组),则步长维度变量为 NULL

除了 PyArrayObject 的步长(strides)和维度(dimensions)成员中包含的结构信息外,标志(flags)还包含有关数据如何访问的重要信息。特别是,当内存根据数据类型数组位于合适的边界时,NPY_ARRAY_ALIGNED 标志会被设置。即使你有一块连续的内存,你也不能仅仅假设解引用指向元素的特定数据类型指针是安全的。只有当 NPY_ARRAY_ALIGNED 标志被设置时,这才是安全的操作。在某些平台上它可以工作,但在其他平台上,如 Solaris,它会导致总线错误。如果你打算写入数组的内存区域,还应该确保设置了 NPY_ARRAY_WRITEABLE 标志。获取指向不可写内存区域的指针也是可能的。有时,当 NPY_ARRAY_WRITEABLE 标志未设置时写入内存区域只是粗鲁的行为。其他时候它可能导致程序崩溃(例如,一个只读内存映射文件的数据区域)。

数据类型封装#

数据类型ndarray 的重要抽象。操作将依赖数据类型来提供操作数组所需的关键功能。此功能由 PyArray_Descr 结构体的 f 成员指向的函数指针列表提供。通过这种方式,只需在 f 成员中提供带有合适函数指针的 PyArray_Descr 结构体,即可扩展数据类型的数量。对于内置类型,存在一些绕过此机制的优化,但数据类型抽象的目的是允许添加新的数据类型。

内置数据类型之一是 void 数据类型,它允许任意结构化类型,其中包含一个或多个字段作为数组元素。字段只是另一个数据类型对象以及当前结构化类型中的偏移量。为了支持任意嵌套字段,为 void 类型实现了几种数据类型访问的递归实现。一个常见的惯用法是循环遍历字典的元素,并根据给定偏移量处存储的数据类型对象执行特定操作。这些偏移量可以是任意数字。因此,必须认识到并考虑遇到未对齐数据的可能性(如有必要)。

N 维迭代器#

另请参阅

遍历数组

在许多 NumPy 代码中,一个非常常见的操作是需要遍历通用、带步长、N 维数组的所有元素。这种通用 N 维循环的操作被抽象为迭代器对象的概念。要编写一个 N 维循环,你只需从一个 ndarray 创建一个迭代器对象,操作迭代器对象结构中的 dataptr 成员,并对迭代器对象调用宏 PyArray_ITER_NEXT 以移动到下一个元素。下一个元素始终按 C 连续顺序排列。该宏的工作原理是首先特殊处理 C 连续、1 维和 2 维的情况,这些情况非常简单。

对于一般情况,迭代通过跟踪迭代器对象中的坐标计数器列表来工作。在每次迭代中,最后一个坐标计数器会增加(从 0 开始)。如果此计数器小于该维度中数组大小减一(一个预先计算和存储的值),则计数器增加,并且 dataptr 成员按照该维度中的步长增加,宏结束。如果到达维度的末尾,则最后一个维度的计数器重置为零,并且 dataptr 通过减去步长值乘以该维度中元素数量减一(这也预先计算并存储在迭代器对象的 backstrides 成员中)来移回该维度的开头。在这种情况下,宏不会结束,而是局部维度计数器递减,以便倒数第二个维度取代最后一个维度的作用,并且前面描述的测试再次在倒数第二个维度上执行。通过这种方式,dataptr 被适当地调整以适应任意步长。

PyArrayIterObject 结构体的 coordinates 成员维护当前的 N 维计数器,除非底层数组是 C 连续的,在这种情况下会跳过坐标计数。PyArrayIterObjectindex 成员跟踪迭代器的当前扁平索引。它由 PyArray_ITER_NEXT 宏更新。

广播#

另请参阅

广播

在 Numeric(NumPy 的前身)中,广播是在 ufuncobject.c 中深埋的几行代码中实现的。在 NumPy 中,广播的概念已被抽象化,以便可以在多个地方执行。广播由函数 PyArray_Broadcast 处理。此函数需要传入一个 PyArrayMultiIterObject(或其二进制等效物)。PyArrayMultiIterObject 跟踪广播的维度数量和每个维度的大小,以及广播结果的总大小。它还跟踪正在广播的数组数量,以及指向每个正在广播数组的迭代器的指针。

PyArray_Broadcast 函数接受已定义的迭代器,并使用它们来确定每个维度中的广播形状(若要在广播发生时同时创建迭代器,请使用 PyArray_MultiIterNew 函数)。然后,调整迭代器,使每个迭代器都认为它正在遍历一个具有广播大小的数组。这是通过调整迭代器的维度数量和每个维度中的形状来完成的。这之所以有效,是因为迭代器步长也得到了调整。广播只调整(或添加)长度为 1 的维度。对于这些维度,步长变量简单地设置为 0,这样当广播操作在扩展维度上执行时,该数组的迭代器的数据指针不会移动。

在 Numeric 中,广播总是使用 0 值步长实现扩展维度。在 NumPy 中,它的实现方式也完全相同。最大的区别在于,现在步长数组在 PyArrayIterObject 中被跟踪,广播结果中涉及的迭代器在 PyArrayMultiIterObject 中被跟踪,并且 PyArray_Broadcast 调用实现了通用广播规则

数组标量#

另请参阅

标量

数组标量提供了一个 Python 类型层次结构,允许数组中存储的数据类型与从数组中提取元素时返回的 Python 类型之间建立一对一的对应关系。对象数组是一个例外。对象数组是任意 Python 对象的异构集合。当你从对象数组中选择一个项时,你会得到原始的 Python 对象(而不是对象数组标量,后者虽然存在但很少用于实际目的)。

数组标量还提供与数组相同的方法和属性,旨在使相同的代码能够支持任意维度(包括 0 维度)。数组标量是只读的(不可变的),除了 void 标量,它也可以写入,以便结构化数组字段设置更自然地工作(a[0]['f1'] = value)。

索引#

所有 Python 索引操作 arr[index] 都通过首先准备索引并查找索引类型来组织。支持的索引类型有

  • 整数

  • 新轴

  • 切片

  • 省略号

  • 整数数组/类数组(高级)

  • 布尔值(单个布尔数组);如果索引中存在多个布尔数组或形状不完全匹配,则布尔数组将被转换为整数数组。

  • 0 维布尔值(以及整数);0 维布尔数组是一个特殊情况,必须在高级索引代码中处理。它们表示 0 维布尔数组必须被解释为整数数组。

以及标量数组的特殊情况,它表示整数数组被解释为整数索引,这很重要,因为整数数组索引会强制复制,但如果返回标量则会被忽略(完全整数索引)。准备好的索引保证是有效的,除了越界值和高级索引的广播错误。这包括为不完整的索引添加 Ellipsis,例如当使用单个整数对二维数组进行索引时。

下一步取决于找到的索引类型。如果所有维度都用整数索引,则返回或设置一个标量。单个布尔索引数组将调用专门的布尔函数。包含 Ellipsis切片但没有高级索引的索引将始终通过计算新的步长和内存偏移量来创建旧数组的视图。此视图可以返回,或者在赋值时使用 PyArray_CopyObject 填充。请注意,当数组为对象 dtype 时,PyArray_CopyObject 也可能在其他分支上对临时数组调用,以支持复杂的赋值。

高级索引#

迄今为止最复杂的情况是高级索引,它可能与典型的基于视图的索引结合使用,也可能不结合。在这里,整数索引被解释为基于视图的索引。在尝试理解这一点之前,你可能需要熟悉其微妙之处。高级索引代码有三个不同的分支和一个特殊情况

  • 有一个索引数组,它以及赋值数组都可以轻易迭代。例如,它们可以是连续的。此外,索引数组必须是 intp 类型,并且赋值中的值数组应为正确类型。这纯粹是一个快速路径。

  • 只有整数数组索引,因此不存在子数组。

  • 基于视图和高级索引混合。在这种情况下,基于视图的索引定义了一组子数组,这些子数组通过高级索引进行组合。例如,arr[[1, 2, 3], :] 是通过垂直堆叠子数组 arr[1, :]arr[2, :]arr[3, :] 创建的。

  • 存在一个子数组,但它恰好只有一个元素。这种情况可以像没有子数组一样处理,但在设置时需要一些注意。

确定适用哪种情况、检查广播以及确定所需的转置类型都在 PyArray_MapIterNew 中完成。设置完成后,有两种情况。如果没有子数组或者子数组只有一个元素,则无需子数组迭代,并准备一个迭代器,该迭代器迭代所有索引数组以及结果或值数组。如果存在子数组,则准备三个迭代器。一个用于索引数组,一个用于结果或值数组(减去其子数组),一个用于原始数组和结果/赋值数组的子数组。前两个迭代器提供(或允许计算)指向子数组开头的指针,然后允许重新启动子数组迭代。

当高级索引相邻时,可能需要转置。所有必要的转置都由 PyArray_MapIterSwapAxes 处理,并且必须由调用者处理,除非 PyArray_MapIterNew 被要求分配结果。

准备好之后,获取和设置相对简单,尽管需要考虑不同的迭代模式。除非在获取项时只有一个索引数组,否则会预先检查索引的有效性。否则,为了优化,它将在内部循环中自行处理。

通用函数#

通用函数是可调用对象,它们接受 \(N\) 个输入并产生 \(M\) 个输出,通过将逐元素工作的基本 1 维循环封装成完整易用的函数,无缝实现广播类型检查缓冲强制类型转换输出参数处理。新的通用函数通常用 C 语言创建,尽管存在从 Python 函数创建 ufunc 的机制(frompyfunc)。用户必须提供一个 1 维循环,该循环实现基本功能,接收输入标量值并将结果标量放置到适当的输出槽中,如实现中所述。

设置#

每个 ufunc 计算都涉及到一些与设置计算相关的开销。这种开销的实际意义在于,尽管 ufunc 的实际计算速度非常快,但你将能够编写比 ufunc 更快地处理小数组的数组和特定类型代码。特别是,使用 ufuncs 对 0 维数组执行许多计算会比其他基于 Python 的解决方案慢(悄悄导入的 scalarmath 模块的存在正是为了让数组标量具有基于 ufunc 的计算的外观和感觉,并且显著减少开销)。

当调用 ufunc 时,必须执行许多操作。从这些设置操作中收集到的信息存储在一个循环对象中。这个循环对象是一个 C 结构体(它可以成为 Python 对象,但由于仅在内部使用,因此不会作为 Python 对象进行初始化)。这个循环对象具有与 PyArray_Broadcast 一起使用所需的布局,以便广播可以以与其他代码段中处理广播相同的方式进行处理。

首先要做的是在线程特定的全局字典中查找缓冲区大小、错误掩码和相关错误对象的当前值。错误掩码的状态控制着当发现错误条件时会发生什么。应该注意的是,硬件错误标志的检查只在每个一维循环执行后才进行。这意味着如果输入和输出数组是连续的且类型正确,从而执行单个一维循环,那么直到数组的所有元素都被计算出来之后才可能检查这些标志。在线程特定的字典中查找这些值会花费时间,但对于除了非常小的数组之外的所有情况,这都可以轻易忽略。

检查线程特定的全局变量后,评估输入以确定 ufunc 应如何进行,并在必要时构建输入和输出数组。任何非数组的输入都将转换为数组(必要时使用上下文)。哪些输入是标量(因此转换为 0 维数组)也会被记录下来。

接下来,根据输入数组类型从 ufunc 可用的一维循环中选择一个合适的一维循环。通过尝试将输入的 DataTypes 签名与可用签名进行匹配来选择此一维循环。内置类型对应的签名存储在 ufunc 结构体的 ufunc.types 成员中。用户定义类型对应的签名存储在一个函数信息链表中,其头部元素作为 CObject 存储在以数据类型编号为键(参数列表中的第一个用户定义类型用作键)的 userloops 字典中。搜索签名,直到找到一个可以安全地将所有输入数组转换为该签名的签名(忽略不允许确定结果类型的任何标量参数)。这种搜索过程的含义是,当存储签名时,“较小类型”应放置在“较大类型”之下。如果没有找到一维循环,则报告错误。否则,argument_list 会用存储的签名进行更新 — 以防需要进行类型转换并修正一维循环假定的输出类型。

如果 ufunc 有 2 个输入和 1 个输出,并且第二个输入是 Object 数组,则会执行特殊情况检查,以便在第二个输入不是 ndarray、具有 __array_priority__ 属性且具有 __r{op}__ 特殊方法时返回 NotImplemented。通过这种方式,Python 会发出信号,让另一个对象有机会完成操作,而不是使用通用的对象数组计算。这允许(例如)稀疏矩阵覆盖乘法运算符的一维循环。

对于小于指定缓冲区大小的输入数组,会复制所有非连续、未对齐或字节序错误的数组,以确保对于小型数组只使用一个循环。然后,为所有输入数组创建数组迭代器,并将生成的迭代器集合广播到单个形状。

然后处理输出参数(如果有),并构造任何缺失的返回数组。如果任何提供的输出数组类型不正确(或未对齐)且小于缓冲区大小,则会构造一个新的输出数组,并设置特殊的 NPY_ARRAY_WRITEBACKIFCOPY 标志。在函数结束时,会调用 PyArray_ResolveWritebackIfCopy,以便将其内容复制回输出数组。然后处理输出参数的迭代器。

最后,决定如何执行循环机制,以确保输入数组的所有元素组合起来产生正确类型的输出数组。循环执行的选项包括单循环(针对连续、对齐且数据类型正确的数组)、步长循环(针对非连续但仍对齐且数据类型正确的数组),以及缓冲循环(针对未对齐或数据类型不正确的情况)。根据调用哪种执行方法,然后设置并计算循环。

函数调用#

本节描述了如何为三种不同类型的执行设置和执行基本通用函数计算循环。如果在编译期间定义了 NPY_ALLOW_THREADS,那么只要不涉及对象数组,Python 全局解释器锁(GIL)就会在调用循环之前释放。如果需要处理错误条件,则会重新获取它。硬件错误标志只在一维循环完成后才进行检查。

单循环#

这是所有情况中最简单的一种。ufunc 通过精确调用底层的一维循环一次来执行。这只有当输入和输出都具有正确类型(包括字节序)的对齐数据,并且所有数组都具有统一的步长(无论是连续的、0 维的还是 1 维的)时才可能。在这种情况下,一维计算循环只被调用一次,以计算整个数组。请注意,硬件错误标志只在整个计算完成后才检查。

步长循环#

当输入和输出数组对齐且类型正确,但步长不统一(非连续且为 2 维或更大)时,则采用第二种循环结构进行计算。这种方法将输入和输出参数的所有迭代器转换为除最大维度之外的所有维度进行迭代。内层循环由底层的一维计算循环处理。外层循环是转换后的迭代器上的标准迭代器循环。硬件错误标志在每个一维循环完成后检查。

缓冲循环#

这是处理输入和/或输出数组未对齐或数据类型不正确(包括字节交换)的情况的代码,这些情况与底层一维循环的预期不符。这些数组也被假定为非连续的。此代码与步长循环非常相似,除了内层一维循环被修改,以便对输入执行预处理,并对输出执行后处理,以 bufsize 块为单位(其中 bufsize 是用户可设置的参数)。底层的一维计算循环对已复制(如果需要)的数据进行调用。在这种情况下,设置代码和循环代码要复杂得多,因为它必须处理

  • 临时缓冲区的内存分配

  • 决定是否在输入和输出数据上使用缓冲区(未对齐和/或数据类型错误)

  • 对任何需要缓冲区的输入或输出数据进行复制和可能的类型转换。

  • Object 数组进行特殊处理,以便在需要复制和/或类型转换时正确处理引用计数。

  • 将内部一维循环分解成 bufsize 块(可能带有余数)。

同样,硬件错误标志在每个一维循环结束时检查。

最终输出操作#

Ufuncs 允许其他类似数组的类无缝地通过接口传递,即特定类的输入将导致输出也属于该类。其工作机制如下:如果任何输入不是 ndarray 且定义了 __array_wrap__ 方法,则具有最大 __array_priority__ 属性的类将决定所有输出的类型(除了任何已传入的输出数组)。输入数组的 __array_wrap__ 方法将被调用,ufunc 返回的 ndarray 作为其输入。 __array_wrap__ 函数支持两种调用风格。第一种是将 ndarray 作为第一个参数,将包含“上下文”的元组作为第二个参数。上下文是 (ufunc, arguments, output argument number)。这是第一次尝试的调用。如果发生 TypeError,则该函数只将 ndarray 作为第一个参数调用。

方法#

有三种 ufunc 方法需要类似于通用 ufunc 的计算。它们是 ufunc.reduceufunc.accumulateufunc.reduceat。这些方法中的每一个都需要一个设置命令,后跟一个循环。这些方法可能有四种循环样式:无元素、单元素、步长循环和缓冲循环。这些是与通用函数调用实现相同的基础循环样式,除了无元素和单元素情况是当输入数组对象分别有 0 和 1 个元素时发生的特殊情况。

设置#

所有这三种方法的设置函数都是 construct_reduce。此函数创建一个归约循环对象,并用完成循环所需的参数填充它。所有这些方法都只适用于接受 2 个输入并返回 1 个输出的 ufunc。因此,底层的一维循环被选择,假定签名为 [otype, otype, otype],其中 otype 是请求的归约数据类型。然后从(每线程)全局存储中检索缓冲区大小和错误处理信息。对于未对齐或数据类型不正确的小型数组,会进行复制,以便使用非缓冲代码段。然后,选择循环策略。如果数组中有 1 个元素或 0 个元素,则选择简单的循环方法。如果数组未对齐且数据类型正确,则选择步长循环。否则,必须执行缓冲循环。然后建立循环参数,并构造返回数组。输出数组的形状取决于方法是 reduceaccumulate 还是 reduceat。如果已提供输出数组,则会检查其形状。如果输出数组不是 C 连续、对齐且数据类型正确的,则会创建一个带有 NPY_ARRAY_WRITEBACKIFCOPY 标志的临时副本。通过这种方式,方法将能够使用行为良好的输出数组,但在函数完成时调用 PyArray_ResolveWritebackIfCopy 时,结果将被复制回真实的输出数组。最后,设置迭代器以循环遍历正确的(取决于提供给方法的轴值),然后设置例程返回到实际的计算例程。

归约#

所有 ufunc 方法都使用相同的底层一维计算循环,并调整输入和输出参数,以便进行适当的归约。例如,reduce 函数的关键在于调用一维循环时,输出和第二个输入都指向内存中的相同位置,并且步长都为 0。第一个输入指向输入数组,其步长由所选轴的适当步长给出。通过这种方式,执行的操作是

\begin{align*} o & = & i[0] \\ o & = & i[k]\textrm{<op>}o\quad k=1\ldots N \end{align*}

其中 \(N+1\) 是输入 \(i\) 中的元素数量,\(o\) 是输出,\(i[k]\)\(i\) 沿所选轴的第 \(k^{\textrm{th}}\) 个元素。此基本操作对于大于 1 维的数组重复执行,以便沿所选轴的每个 1 维子数组都进行归约。一个移除了所选维度的迭代器处理此循环。

对于缓冲循环,在调用循环函数之前必须注意复制和转换数据,因为底层循环需要正确数据类型(包括字节序)的对齐数据。缓冲循环必须在对不超过用户指定的 bufsize 大小的块调用循环函数之前处理此复制和转换。

累加#

accumulate 方法与 reduce 方法非常相似,因为输出和第二个输入都指向输出。不同之处在于,第二个输入指向当前输出指针后面一个步长的内存。因此,执行的操作是

\begin{align*} o[0] & = & i[0] \\ o[k] & = & i[k]\textrm{<op>}o[k-1]\quad k=1\ldots N. \end{align*}

输出与输入具有相同的形状,并且当所选轴中的形状为 \(N+1\) 时,每个 1 维循环操作 \(N\) 个元素。同样,缓冲循环在调用底层 1 维计算循环之前会负责复制和转换数据。

Reduceat#

reduceat 函数是 reduceaccumulate 函数的泛化。它通过索引指定输入数组的范围来执行 reduce 操作。在循环计算发生之前,会检查额外的 `indices` 参数,以确保每个输入在所选维度上不会对输入数组来说太大。循环实现使用与 reduce 代码非常相似的代码进行处理,重复次数与 `indices` 输入中的元素数量相同。特别是:传递给底层一维计算循环的第一个输入指针指向索引数组指示的正确位置处的输入数组。此外,传递给底层一维循环的输出指针和第二个输入指针指向内存中的相同位置。一维计算循环的大小被固定为当前索引与下一个索引之间的差值(当当前索引是最后一个索引时,则下一个索引被假定为数组沿所选维度的长度)。通过这种方式,一维循环将对指定的索引执行 reduce 操作。

未对齐或与输入和/或输出数据类型不匹配的循环数据类型通过缓冲代码处理,其中数据被复制到临时缓冲区并在必要时转换为正确的数据类型,然后调用底层一维函数。临时缓冲区以不超过用户可设置的缓冲区大小值(元素)大小创建。因此,循环必须足够灵活,能够多次调用底层一维计算循环,以不超过缓冲区大小的块完成总计算。