NEP 11 — 延迟 UFunc 计算#

作者:

Mark Wiebe <mwwiebe@gmail.com>

内容类型:

text/x-rst

创建日期:

2010-11-30

状态:

延迟

摘要#

本 NEP 描述了一个向 NumPy 的 UFunc 添加延迟计算的提案。这将允许像“a[:] = b + c + d + e”这样的 Python 表达式一次性遍历所有变量进行计算，而无需临时数组。由此产生的性能可能与numexpr库相当，但语法更自然。

这个想法与 UFunc 错误处理和 UPDATEIFCOPY 标志有一些交互作用，影响了设计和实现，但结果允许在几乎不需要 Python 用户付出努力的情况下使用延迟计算。

动机#

NumPy 的 UFunc 执行方式导致大型表达式的性能不佳，因为分配了多个临时变量，并且对输入进行了多次扫描。numexpr库可以通过在小的缓存友好的块中执行计算以及对每个元素计算整个表达式来胜过 NumPy。这导致对每个输入进行一次扫描，这对缓存来说要好得多。

关于如何在不更改 Python 代码的情况下在 NumPy 中获得这种行为的想法，可以考虑 C++ 中的表达式模板技术。这些可以用来相当随意地重新排列使用向量或其他数据结构的表达式，例如

A = B + C + D;

可以转换成等同于

for(i = 0; i < A.size; ++i) {
    A[i] = B[i] + C[i] + D[i];
}

这是通过返回一个知道如何计算结果的代理对象而不是返回实际对象来实现的。使用现代 C++ 优化编译器，生成的机器代码通常与手工编写的循环相同。例如，请参见Blitz++ 库。一个最近创建的用于帮助编写表达式模板的库是Boost Proto。

通过在 Python 中使用返回代理对象相同的思想，我们可以动态地完成同样的事情。返回的对象是一个未分配缓冲区的 ndarray，并且具有足够的知识来在需要时计算自身。当最终计算“延迟数组”时，我们可以使用由所有操作数延迟数组组成的表达式树，有效地创建一个新的 UFunc 来动态计算。

示例 Python 代码#

以下是它在 NumPy 中的用法。

# a, b, c are large ndarrays

with np.deferredstate(True):

    d = a + b + c
    # Now d is a 'deferred array,' a, b, and c are marked READONLY
    # similar to the existing UPDATEIFCOPY mechanism.

    print d
    # Since the value of d was required, it is evaluated so d becomes
    # a regular ndarray and gets printed.

    d[:] = a*b*c
    # Here, the automatically combined "ufunc" that computes
    # a*b*c effectively gets an out= parameter, so no temporary
    # arrays are needed whatsoever.

    e = a+b+c*d
    # Now e is a 'deferred array,' a, b, c, and d are marked READONLY

    d[:] = a
    # d was marked readonly, but the assignment could see that
    # this was due to it being a deferred expression operand.
    # This triggered the deferred evaluation so it could assign
    # the value of a to d.

但是，可能存在一些意外行为。

with np.deferredstate(True):

    d = a + b + c
    # d is deferred

    e[:] = d
    f[:] = d
    g[:] = d
    # d is still deferred, and its deferred expression
    # was evaluated three times, once for each assignment.
    # This could be detected, with d being converted to
    # a regular ndarray the second time it is evaluated.

我相信在文档中应该推荐的用法是将延迟状态保留为默认值，除非评估可以从中受益的大型表达式。

# calculations

with np.deferredstate(True):
    x = <big expression>

# more calculations

这将避免始终保持延迟使用为 True 时可能出现的意外情况，例如在稍后使用延迟表达式时出现意外时间的浮点警告或异常。“为什么我的打印语句抛出除零错误？”之类的用户问题可以通过推荐这种方法来避免。

建议的延迟计算 API#

为了使延迟计算工作，C API 需要知道它的存在，并且能够在必要时触发计算。ndarray 将增加两个新的标志。

NPY_ISDEFERRED

指示此 ndarray 实例的表达式计算已被延迟。

NPY_DEFERRED_WASWRITEABLE

只有当PyArray_GetDeferredUsageCount(arr) > 0时才能设置。它指示当arr首次在延迟表达式中使用时，它是一个可写数组。如果设置了此标志，则调用PyArray_CalculateAllDeferred()将使arr再次可写。

注意

问题

NPY_DEFERRED 和 NPY_DEFERRED_WASWRITEABLE 应该对 Python 可见，还是应该在必要时从 python 访问标志触发 PyArray_CalculateAllDeferred？

API 将通过许多函数进行扩展。

int PyArray_CalculateAllDeferred()

此函数强制执行所有当前延迟的计算。

例如，如果错误状态设置为忽略所有，并且 np.seterr({all='raise'})，这将改变对已延迟表达式的处理方式。因此，在更改错误状态之前，应该计算所有现有的延迟数组。

int PyArray_CalculateDeferred(PyArrayObject* arr)

如果“arr”是延迟数组，则为其分配内存并计算延迟表达式。如果“arr”不是延迟数组，则只需返回成功。返回 NPY_SUCCESS 或 NPY_FAILURE。

int PyArray_CalculateDeferredAssignment(PyArrayObject* arr, PyArrayObject* out)

如果“arr”是延迟数组，则将延迟表达式计算到“out”中，“arr”仍然是延迟数组。如果“arr”不是延迟数组，则将其值复制到 out。返回 NPY_SUCCESS 或 NPY_FAILURE。

int PyArray_GetDeferredUsageCount(PyArrayObject* arr)

返回使用此数组作为操作数的延迟表达式的数量。

Python API 将如下扩展。

numpy.setdeferred(state)

启用或禁用延迟计算。True 表示始终使用延迟计算。False 表示从不使用延迟计算。None 表示如果错误处理状态设置为忽略所有内容，则使用延迟计算。在 NumPy 初始化时，延迟状态为 None。

返回之前的延迟状态。

numpy.getdeferred()

返回当前延迟状态。

numpy.deferredstate(state)

用于延迟状态处理的上下文管理器，类似于numpy.errstate。

错误处理#

错误处理对于延迟计算来说是一个棘手的问题。如果 NumPy 错误状态为 {all='ignore'}，则引入延迟计算作为默认值可能是合理的，但是如果 UFunc 可以引发错误，则后面的“打印”语句引发异常而不是实际导致错误的操作将非常奇怪。

一个可能的好方法是默认情况下仅在错误状态设置为忽略所有情况时启用延迟计算，但允许用户使用“setdeferred”和“getdeferred”函数进行控制。True 表示始终使用延迟计算，False 表示从不使用它，而 None 表示仅在安全时使用它（即错误状态设置为忽略所有）。

与 UPDATEIFCOPY 的交互#

NPY_UPDATEIFCOPY文档说明

数据区域表示一个（行为良好的）副本，当此数组被删除时，其信息应传输回原始数组。

这是一个特殊标志，如果此数组表示由于用户在 PyArray_FromAny 中需要某些标志而创建的副本（并且用户要求在这种情况下设置此标志），则会设置此标志。然后，base 属性指向“行为不端”的数组（设置为只读）。当设置了此标志的数组被释放时，它将把其内容复制回“行为不端”的数组（如果需要则进行转换），并将“行为不端”的数组重置为 NPY_WRITEABLE。如果“行为不端”的数组一开始不是 NPY_WRITEABLE，那么 PyArray_FromAny 将返回错误，因为 NPY_UPDATEIFCOPY 将无法实现。

当前UPDATEIFCOPY的实现假设它是唯一以这种方式修改可写标志的机制。这些机制必须相互协调才能正常工作。以下是一个它们可能出错的示例。

  1. 使用UPDATEIFCOPY创建‘arr’的临时副本（‘arr’变为只读）。

  2. 在延迟表达式中使用‘arr’（延迟使用计数变为1，`NPY_DEFERRED_WASWRITEABLE`**未**设置，因为‘arr’是只读的）。

  3. 销毁临时副本，导致‘arr’变为可写。

  4. 写入‘arr’会破坏延迟表达式的值。

为了解决这个问题，我们使这两种状态互斥。

  • UPDATEIFCOPY的使用会检查`NPY_DEFERRED_WASWRITEABLE`标志，如果已设置，则调用`PyArray_CalculateAllDeferred`来刷新所有延迟计算后再继续。

  • ndarray 获取一个新的标志`NPY_UPDATEIFCOPY_TARGET`，指示数组将在将来的某个时间点被更新并变为可写。如果延迟计算机制在任何操作数中看到此标志，它将触发立即计算。

其他实现细节#

创建延迟数组时，它会引用UFunc的所有操作数以及UFunc本身。“DeferredUsageCount”为每个操作数递增，并在计算延迟表达式或销毁延迟数组时递减。

跟踪一个全局列表，其中包含对所有延迟数组的弱引用，按创建顺序排列。当调用`PyArray_CalculateAllDeferred`时，首先计算最新的延迟数组。这可能会释放对延迟表达式树中包含的其他延迟数组的引用，从而无需计算这些数组。

进一步优化#

与其在任何错误未设置为“忽略”时保守地禁用延迟计算，不如让每个UFunc提供它生成的可能错误的集合。然后，如果所有这些错误都设置为“忽略”，即使其他错误未设置为忽略，也可以使用延迟计算。

一旦显式存储表达式树，就可以对其进行转换。例如，`add(add(a,b),c)`可以转换为`add3(a,b,c)`，或者`add(multiply(a,b),c)`可以使用CPU融合乘加指令（如果可用）转换为`fma(a,b,c)`。

虽然我将延迟计算框架定义为仅用于UFunc，但它可以扩展到其他函数，例如`dot()`。例如，可以重新排序链式矩阵乘法以最大限度地减少中间结果的大小，或者窥孔式优化器通道可以搜索与优化的BLAS/其他高性能库调用匹配的模式。

对于非常大的数组上的操作，将像LLVM这样的JIT集成到此系统中可能是一个很大的好处。UFunc和其他操作将提供位代码，这些位代码可以一起内联并由LLVM优化器进行优化，然后执行。实际上，迭代器本身也可以用位代码表示，允许LLVM在进行优化时考虑整个迭代。