NEP 43 — 增强UFunc的可扩展性#

标题:

增强UFunc的可扩展性

作者:

Sebastian Berg

状态:

草稿

类型:

标准

创建时间:

2020-06-20

注意

此NEP是系列中的第四篇

  • NEP 40 解释了NumPy的dtype实现的不足。

  • NEP 41 概述了我们建议的替代方案。

  • NEP 42 描述了新设计中与数据类型相关的API。

  • NEP 43(本文档)描述了新设计中通用函数(ufunc)的API。

摘要#

前一篇NEP 42提议创建用户可以在NumPy本身之外定义的新DType。NEP 42的实现将使用户能够创建具有自定义dtype和存储值的数组。本NEP概述了NumPy将来如何处理具有自定义dtype的数组。操作NumPy数组最重要的函数是所谓的“通用函数”(ufunc),包括所有数学函数,例如np.addnp.multiply,甚至np.matmul。这些ufunc必须高效地处理具有不同数据类型的多个数组。

本NEP提议扩展ufunc的设计。它区分了可以操作多种不同dtype(如浮点数或整数)的ufunc,以及一个定义特定dtype高效操作的新ArrayMethod

注意

私有和外部API的详细信息可能会根据用户评论和实现约束而更改。底层原理和选择不应发生重大变化。

动机与范围#

本NEP的目标是扩展通用函数以支持NEPs 41和42中详细介绍的新DType系统。虽然主要动机是支持新的用户定义DType,但这也将大大简化NumPy字符串或结构化DType的通用函数的定义。到目前为止,这些DType不受NumPy任何函数(如np.addnp.equal)的支持,原因是其参数化性质(参见NEP 41和42)带来的困难,例如字符串长度。

数组上的函数必须处理一系列不同的步骤,这些步骤将在“UFunc调用涉及的步骤”部分更详细地描述。最重要的步骤是

  • 组织定义特定DType的ufunc调用所需的所有功能。这通常被称为“内部循环”(inner-loop)。

  • 处理找不到精确匹配实现的情况。例如,当int32float64相加时,int32将被转换为float64。这需要一个独立的“提升”(promotion)步骤。

在组织和定义了这些之后,我们需要

  • 定义用户API以自定义以上两点。

  • 允许方便地重用现有功能。例如,表示物理单位(如米)的DType应该能够回退到NumPy现有的数学实现。

本NEP详细说明了这些要求将在NumPy中如何实现。

  • 当前属于ufunc定义的所有特定于DTyper的功能将被定义为一个新的ArrayMethod对象。这个ArrayMethod对象将是描述操作数组的任何函数的新首选方式。

  • Ufuncs将分派到ArrayMethod,并可能使用提升来找到要使用的正确ArrayMethod。这将在“提升和分派”部分进行描述。

每个部分都将概述一个新的C-API。预期的未来Python API将非常相似,并且为了可读性,C-API将以Python代码的形式呈现。

NEP提议对NumPy ufunc的内部结构进行大规模但必要的重构。这种现代化不会直接影响最终用户,不仅是新DType的必要步骤,而且本身也是一项维护工作,预计将有助于未来ufunc机制的改进。

虽然提议的最重要的重构是新的ArrayMethod对象,但长期的最大考虑是用于提升和分派的API选择。

向后兼容性#

一般的向后兼容性问题也已在NEP 41中列出。

绝大多数用户除了NumPy发布中的典型变化之外,不应看到任何其他变化。以下是受提议更改影响的三个主要用户或用例

  1. Numba包使用对NumPy C-loops的直接访问,并直接修改NumPy ufunc结构以满足自身目的。

  2. Astropy使用其自己的“类型解析器”,这意味着从现有类型解析到新的默认Promoter的默认切换需要谨慎。

  3. 目前可以为dtype*实例*注册循环。这对于结构化dtype在理论上很有用,并且是这里提出的DType解析步骤*之后*发生的解析步骤。

本NEP将尽力在最大程度上保持向后兼容性。然而,这两个项目都已表示愿意适应破坏性更改。

NumPy能够提供向后兼容性的主要原因是

  • 现有的内部循环可以被封装,在调用时增加了间接性,但保持了完全的向后兼容性。get_loop函数在这种情况下可以搜索存储在ufunc直接上的现有内部循环函数,以保持与潜在的直接结构访问的完全兼容性。

  • 可以调用旧的类型解析器作为回退(可能缓存结果)。解析器可能需要被调用两次(一次用于DType解析,一次用于resolve_descriptor实现)。

  • 在大多数情况下,回退到旧的类型解析器将处理为这种结构化dtype实例定义的循环。这是因为如果没有其他的np.Void实现,旧的回退至少会保留旧的行为。

掩码类型解析器(masked type resolvers)*不会*保留支持,但没有已知的用户(包括NumPy本身,它只使用默认版本)。

此外,对于*调用*(而不是提供正常或掩码类型解析器)将不尝试兼容性。因为NumPy只将其作为回退使用。没有已知的用户使用这种(未记录的)可能性。

尽管上述更改可能会破坏一些工作流程,但我们认为长期的改进远远超过了这一点。此外,像astropy和Numba这样的包能够适应,因此最终用户可能需要更新他们的库,而不是他们的代码。

用法与影响#

本NEP重构了NumPy内部和外部实现者如何定义NumPy数组上的操作。NEP主要涉及那些为自定义DType扩展ufunc或创建自定义ufunc的人。它不旨在最终确定所有潜在用例,而是重构NumPy使其具有可扩展性,并允许在出现新问题或功能请求时进行处理。

概述和最终用户API#

为了概述本NEP提议如何构建ufunc,以下内容描述了拟议重构可能向最终用户暴露的内容。

当考虑具有单个输入的ufunc时,通用函数非常类似于在数组的DType上定义的Python方法。

res = np.positive(arr)

可以(概念上)实现为

positive_impl = arr.dtype.positive
res = positive_impl(arr)

然而,与方法不同的是,positive_impl并不存储在dtype本身上。它而是np.positive对于特定DType的实现。当前的NumPy部分暴露了这种“实现选择”,使用了通用函数中的dtype(或更确切地说是signature)属性,尽管这些很少被使用。

np.positive(arr, dtype=np.float64)

强制NumPy使用专门为Float64 DType编写的positive_impl

本NEP通过创建一个新的对象来表示positive_impl,使这种区分更加明确。

positive_impl = np.positive.resolve_impl((type(arr.dtype), None))
# The `None` represents the output DType which is automatically chosen.

虽然创建positive_impl对象和resolve_impl方法是本NEP的一部分,但以下代码

res = positive_impl(arr)

可能不会立即实现,并且不是重新设计的核心。

通常,NumPy通用函数可以接受多个输入。这需要通过考虑所有输入来查找实现,并使ufunc相对于输入DTypes成为“多方法”(multi-methods)。

add_impl = np.add.resolve_impl((type(arr1.dtype), type(arr2.dtype), None))

本NEP定义了positive_impladd_impl如何表示为一个新的ArrayMethod,可以在NumPy外部实现。此外,它定义了resolve_impl将如何实现和解决分派和提升。

在回顾了“UFunc调用涉及的步骤”部分后,这种划分的原因可能会更清楚。

定义新的ufunc实现#

以下是一个模拟,演示如何将新的实现(在此示例中为定义字符串相等性)添加到ufunc中。

class StringEquality(BoundArrayMethod):
    nin = 1
    nout = 1
    # DTypes are stored on the BoundArrayMethod and not on the internal
    # ArrayMethod, to reference cyles.
    DTypes = (String, String, Bool)

    def resolve_descriptors(self: ArrayMethod, DTypes, given_descrs):
        """The strided loop supports all input string dtype instances
        and always returns a boolean. (String is always native byte order.)

        Defining this function is not necessary, since NumPy can provide
        it by default.

        The `self` argument here refers to the unbound array method, so
        that DTypes are passed in explicitly.
        """
        assert isinstance(given_descrs[0], DTypes[0])
        assert isinstance(given_descrs[1], DTypes[1])
        assert given_descrs[2] is None or isinstance(given_descrs[2], DTypes[2])

        out_descr = given_descrs[2]  # preserve input (e.g. metadata)
        if given_descrs[2] is None:
            out_descr = DTypes[2]()

        # The operation is always "no" casting (most ufuncs are)
        return (given_descrs[0], given_descrs[1], out_descr), "no"

    def strided_loop(context, dimensions, data, strides, innerloop_data):
        """The 1-D strided loop, similar to those used in current ufuncs"""
        # dimensions: Number of loop items and core dimensions
        # data: Pointers to the array data.
        # strides: strides to iterate all elements
        n = dimensions[0]  # number of items to loop over
        num_chars1 = context.descriptors[0].itemsize
        num_chars2 = context.descriptors[1].itemsize

        # C code using the above information to compare the strings in
        # both arrays.  In particular, this loop requires the `num_chars1`
        # and `num_chars2`.  Information which is currently not easily
        # available.

np.equal.register_impl(StringEquality)
del StringEquality  # may be deleted.

此定义足以创建一个新循环,并且该结构允许将来的扩展;这已经是实现NumPy内部转换所必需的。我们在这里使用BoundArrayMethodcontext结构。这些将在后面详细描述和说明。简而言之:

  • context是对Python传递给方法的self的泛化。

  • BoundArrayMethod等同于Python中class.method是方法,而class().method返回“绑定”方法的区别。

自定义分派和提升#

当调用np.positive.resolve_impl()时,查找正确的实现主要是实现细节。但在某些情况下,当没有实现与请求的DTypes完全匹配时,可能需要影响此过程。

np.multiple.resolve_impl((Timedelta64, Int8, None))

将没有精确匹配,因为NumPy只有一个用于将Timedelta64乘以Int64的实现。在简单情况下,NumPy将使用默认的提升步骤来尝试查找正确的实现,但为了实现上述步骤,我们将允许以下操作:

def promote_timedelta_integer(ufunc, dtypes):
    new_dtypes = (Timdelta64, Int64, dtypes[-1])
    # Resolve again, using Int64:
    return ufunc.resolve_impl(new_dtypes)

np.multiple.register_promoter(
    (Timedelta64, Integer, None), promote_timedelta_integer)

其中Integer是一个抽象DType(参见NEP 42)。

UFunc调用涉及的步骤#

在进入更详细的API选择之前,回顾NumPy中通用函数调用涉及的步骤很有帮助。

UFunc调用分为以下步骤:

  1. 处理__array_ufunc__协议

  2. 提升和分派

    • 给定所有输入的DTypes,找到正确的实现。例如,float64int64或用户定义的DType的实现。

    • 当不存在精确实现时,必须执行*提升*。例如,将float32float64相加是通过首先将float32转换为float64来实现的。

  3. 参数化dtype解析

    • 通常,每当输出DType是参数化的,就必须找到(解析)参数。

    • 例如,如果一个循环相加两个字符串,则必须定义正确的输出(可能还有输入)dtype。S5 + S4 -> S9,而一个upper函数具有签名S5 -> S5

    • 当它们不是参数化的时,将提供一个默认实现,该实现填充默认的dtype实例(例如,确保本地字节序)。

  4. 准备迭代

    • 此步骤主要由NpyIter(迭代器)内部处理。

    • 分配所有输出和临时缓冲区,以执行转换。这*需要*步骤3中解析的dtype。

    • 找到最佳的迭代顺序,其中包括高效实现广播的信息。例如,将单个值添加到数组会重复相同的该值。

  5. 设置并获取C级函数

    • 如有必要,分配临时工作空间。

    • 找到C实现的、轻量级的内部循环函数。查找内部循环函数可以允许将来进行专门化实现。例如,转换(casting)目前优化了连续转换,而归约(reductions)具有目前在内部循环函数本身内部处理的优化。

    • 指示内部循环是否需要Python API,或者是否可以释放GIL(以允许多线程)。

    • 清除浮点异常标志。

  6. 执行实际计算

    • 运行特定DType的内部循环函数。

    • 内部循环可能需要访问额外的数据,例如dtype或在上一步中设置的附加数据。

    • 内部循环函数可能会被调用不定次数。

  7. 完成

    • 释放步骤5中分配的任何临时工作空间。

    • 检查浮点异常标志。

    • 返回结果。

ArrayMethod提供了一个概念,用于将步骤3到6以及部分7进行分组。但是,新的ufunc或ArrayMethod的实现者通常不需要自定义步骤4或6中的行为,NumPy可以并且确实提供了这些。对于ArrayMethod实现者,需要自定义的核心步骤是步骤3和步骤5。这些提供了自定义的内部循环函数以及可能的内部循环特定的设置。进一步的自定义是可能的,并有望成为未来的扩展。

步骤2是提升和分派,将通过新的API进行重构,以便在需要时允许自定义此过程。

步骤1列出是为了完整性,本NEP不受其影响。

以下草图概述了步骤2到6,并重点介绍了dtype的处理方式,以及哪些部分是可自定义的(“已注册”)而哪些是由NumPy处理的。

_images/nep43-sketch.svg

ArrayMethod#

本NEP的一个核心提案是创建ArrayMethod作为一个对象,用于描述每个特定于一组DType的实现。我们使用class语法来描述创建新的ArrayMethod对象所需的信息。

class ArrayMethod:
    name: str  # Name, mainly useful for debugging

    # Casting safety information (almost always "safe", necessary to
    # unify casting and universal functions)
    casting: Casting = "no"

    # More general flags:
    flags: int

    def resolve_descriptors(self,
            Tuple[DTypeMeta], Tuple[DType|None]: given_descrs) -> Casting, Tuple[DType]:
        """Returns the safety of the operation (casting safety) and the
        """
        # A default implementation can be provided for non-parametric
        # output dtypes.
        raise NotImplementedError

    @staticmethod
    def get_loop(Context : context, strides, ...) -> strided_loop_function, flags:
        """Returns the low-level C (strided inner-loop) function which
        performs the actual operation.

        This method may initially private, users will be able to provide
        a set of optimized inner-loop functions instead:

        * `strided_inner_loop`
        * `contiguous_inner_loop`
        * `unaligned_strided_loop`
        * ...
        """
        raise NotImplementedError

    @staticmethod
    def strided_inner_loop(
            Context : context, data, dimensions, strides, innerloop_data):
        """The inner-loop (equivalent to the current ufunc loop)
        which is returned by the default `get_loop()` implementation."""
        raise NotImplementedError

其中Context提供关于函数调用的主要静态信息。

class Context:
    # The ArrayMethod object itself:
    ArrayMethod : method

    # Information about the caller, e.g. the ufunc, such as `np.add`:
    callable : caller = None
    # The number of input arguments:
    int : nin = 1
    # The number of output arguments:
    int : nout = 1
    # The actual dtypes instances the inner-loop operates on:
    Tuple[DType] : descriptors

    # Any additional information required. In the future, this will
    # generalize or duplicate things currently stored on the ufunc:
    #  - The ufunc signature of generalized ufuncs
    #  - The identity used for reductions

以及flags存储的属性,用于指示:

  • ArrayMethod是否支持未对齐的输入和输出数组。

  • 内部循环函数是否需要Python API(GIL)。

  • NumPy是否必须检查浮点误差CPU标志。

注意:预期会根据需要添加更多信息。

调用Context#

“context”对象类似于Python的self,它被传递给所有方法。为了理解“context”对象为何必要及其内部结构,记住Python方法可以以下方式编写(另请参阅__get__的文档)会很有帮助:

class BoundMethod:
    def __init__(self, instance, method):
        self.instance = instance
        self.method = method

    def __call__(self, *args, **kwargs):
        return self.method.function(self.instance, *args, **kwargs)


class Method:
    def __init__(self, function):
        self.function = function

    def __get__(self, instance, owner=None):
        assert instance is not None  # unsupported here
        return BoundMethod(instance, self)

以下method1method2的行为是相同的。

def function(self):
    print(self)

class MyClass:
    def method1(self):
        print(self)

    method2 = Method(function)

两者都将打印相同的结果。

>>> myinstance = MyClass()
>>> myinstance.method1()
<__main__.MyClass object at 0x7eff65436d00>
>>> myinstance.method2()
<__main__.MyClass object at 0x7eff65436d00>

这里self.instance将是Context传递的所有信息。Context是一个泛化,并且必须传递额外的信息。

  • 与操作单个类实例的方法不同,ArrayMethod操作多个输入数组,因此多个dtype。

  • 上面BoundMethod__call__仅包含一个函数调用。但是ArrayMethod必须调用resolve_descriptors,然后将该信息传递给内部循环函数。

  • Python函数没有除其外部范围定义的之外的状态。在C中,Context能够根据需要提供额外的状态。

就像Python需要区分方法和绑定方法一样,NumPy也将有一个BoundArrayMethod。它存储了Context的所有恒定信息,例如:

幸运的是,大多数用户甚至ufunc实现者都不需要担心这些内部细节;就像很少有Python用户需要了解__get__ dunder方法一样。Context对象或C结构向快速C函数提供所有必要数据,而NumPy API则根据需要创建新的ArrayMethodBoundArrayMethod

ArrayMethod规范#

这些规范提供了一个最小的初始C-API,将来应进行扩展,例如以允许专门化的内部循环。

简而言之,NumPy目前依赖于跨步(strided)的内部循环,并且这将是最初允许定义ufunc的唯一方法。我们期望将来会添加一个setup函数或公开get_loop

UFuncs需要与casting相同的信息,给出以下定义(也参见NEP 42中的CastingImpl):

  • 一个要传递给resolve函数和内部循环的新结构。

    typedef struct {
        PyObject *caller;  /* The ufunc object */
        PyArrayMethodObject *method;
    
        int nin, nout;
    
        PyArray_DTypeMeta **dtypes;
        /* Operand descriptors, filled in by resolve_desciptors */
        PyArray_Descr **descriptors;
    
        void *reserved;  // For Potential in threading (Interpreter state)
    } PyArrayMethod_Context
    

    此结构可能会被附加以包含NumPy未来版本中的附加信息,并包含所有恒定的循环元数据。

    我们可以为该结构添加版本号,尽管为ArrayMethod本身添加版本号可能更简单。

  • 与casting类似,ufunc可能需要找到正确的循环dtype或指示某个循环仅能处理特定实例的涉及DType(例如,仅本地字节序)。这由resolve_descriptors函数处理(与CastingImplresolve_descriptors相同)。

    NPY_CASTING
    resolve_descriptors(
            PyArrayMethodObject *self,
            PyArray_DTypeMeta *dtypes,
            PyArray_Descr *given_dtypes[nin+nout],
            PyArray_Descr *loop_dtypes[nin+nout]);
    

    该函数根据给定的given_dtypes填充loop_dtypes。这需要填充输出的描述符。通常还需要找到输入的描述符,例如以确保内部循环需要时所需的本地字节序。

    在大多数情况下,ArrayMethod将具有非参数化的输出DTypes,因此可以提供默认实现。

  • 一个附加的void *user_data通常会被类型化以扩展现有的NpyAuxData *结构。

    struct {
        NpyAuxData_FreeFunc *free;
        NpyAuxData_CloneFunc *clone;
        /* To allow for a bit of expansion without breaking the ABI */
       void *reserved[2];
    } NpyAuxData;
    

    该结构目前主要用于NumPy内部的casting机制,并且目前必须提供freeclone,尽管这可以放宽。

    与NumPy的casts不同,绝大多数ufunc目前不需要这种额外的临时空间,但可能需要简单的标志功能,例如用于实现警告(请参阅下面的错误和警告处理)。为了简化这一点,当user_data未设置时,NumPy将传递一个零初始化的npy_intp *。*注意,也可以将此作为Context的一部分传递*。

  • 可选的get_loop函数最初将是私有的,以避免最终确定API,这需要与casting相关的设计选择。

    innerloop *
    get_loop(
        PyArrayMethod_Context *context,
        int aligned, int move_references,
        npy_intp *strides,
        PyArray_StridedUnaryOp **out_loop,
        NpyAuxData **innerloop_data,
        NPY_ARRAYMETHOD_FLAGS *flags);
    

    NPY_ARRAYMETHOD_FLAGS可以指示是否需要Python API以及是否必须检查浮点误差。move_references目前在NumPy casting中用于内部用途。

  • 内部循环函数

    int inner_loop(PyArrayMethod_Context *context, ..., void *innerloop_data);
    

    将具有与当前内部循环相同的签名,但有以下更改:

    • 一个返回值,用于在返回-1而不是0时指示错误。当返回-1时,必须设置Python错误。

    • 新的第一个参数PyArrayMethod_Context *用于以方便的方式传入关于ufunc或描述符的潜在必需信息。

    • void *innerloop_data将是NpyAuxData **innerloop_data,由get_loop设置。如果get_loop未设置innerloop_data,则会传递一个npy_intp *(参见下面的错误处理部分以了解动机)。

    *注意:* 由于get_loop预计是私有的,因此innerloop_data的确切实现可以在最终公开之前进行修改。

创建新的BoundArrayMethod将使用PyArrayMethod_FromSpec()函数。一个快捷方式将允许使用PyUFunc_AddImplementationFromSpec()直接注册到ufunc。该规范预计包含以下内容(这可能会在未来扩展):

typedef struct {
    const char *name;  /* Generic name, mainly for debugging */
    int nin, nout;
    NPY_CASTING casting;
    NPY_ARRAYMETHOD_FLAGS flags;
    PyArray_DTypeMeta **dtypes;
    PyType_Slot *slots;
} PyArrayMethod_Spec;

讨论和替代方案#

上述将ArrayMethodContext分离,以及对BoundArrayMethod的附加要求,是必要的分离,反映了Python中方法和绑定方法的实现。

这一要求的一个原因是,它允许在许多情况下存储ArrayMethod对象而不必保留对DTypes的引用,如果DTypes是动态创建(和删除)的,这可能很重要。(这是一个复杂的主题,在当前的Python中没有完整的解决方案,但该方法解决了与casting相关的问题。)

似乎没有其他方法可以实现这种结构。将特定于DType的步骤与通用的ufunc分派和提升分开,对于允许将来的扩展和灵活性绝对是必要的。此外,它允许统一casting和ufunc。

由于ArrayMethodBoundArrayMethod的结构将是不透明的且可扩展的,因此除了选择将它们作为Python对象之外,几乎没有长期的设计影响。

resolve_descriptors#

resolve_descriptors方法可能是本NEP的主要创新,它也是NEP 42中casting实现的核心。

通过确保每个ArrayMethod都提供resolve_descriptors,我们为UFunc调用涉及的步骤中的步骤3定义了一个统一、清晰的API。此步骤是分配输出数组所必需的,并且必须在准备casting之前进行。

虽然返回的casting安全性(NPY_CASTING)对于通用函数几乎总是“no”,但包含它有两个主要优点:

  • -1表示发生了错误。如果设置了Python错误,它将被引发。如果没有设置Python错误,则将被视为“不可能”的转换,并将设置自定义错误。(这对np.can_cast()函数很重要,该函数应引发前者并返回False,而后者对典型的ufunc并不重要)。*这一点正在考虑中,我们可能会使用-1表示一般错误,并使用不同的返回值表示不可能的转换*。

  • 返回casting安全性是NEP 42中casting的核心,并允许在那里 unmodified地使用ArrayMethod

  • 未来可能希望实现快速但不安全的实现。例如,对于int64 + int64 -> int32,从casting角度来看这是不安全的。目前,这将使用int64 + int64 -> int64,然后转换为int32。一个跳过转换的实现必须表明它有效地包含了“相同类型”的转换,因此不被认为是“no”。

get_loop方法#

目前,NumPy ufuncs通常只提供一个跨步循环,因此get_loop方法似乎是不必要的。因此,我们计划最初将get_loop设为私有函数。

然而,get_loop对于casting是必需的,其中专门化的循环甚至超出跨步和连续循环。因此,get_loop函数必须完全替代内部的PyArray_GetDTypeTransferFunction

将来,get_loop可能会公开,或者会公开一个新的setup函数以允许更多控制,例如允许分配工作内存。此外,我们可以扩展get_loop并允许ArrayMethod实现者也控制外部迭代,而不仅仅是1-D内部循环。

扩展内部循环签名#

扩展内部循环签名是本NEP的另一个核心且必要的部分。

传递Context

传递Context可能会通过向context结构添加新字段来允许将来扩展签名。此外,它提供了对内部循环操作的dtype实例的直接访问。这对于参数化dtype是必要信息,因为例如比较两个字符串需要知道两个字符串的长度。Context还可以包含可能有用信息,例如原始ufunc,这在报告错误时可能很有帮助。

原则上,传递Context不是必需的,因为所有信息都可以包含在innerloop_data中并在get_loop函数中设置。在本NEP中,我们建议传递该结构以简化参数化DType的循环创建。

传递用户数据

当前的类型转换实现使用现有的 NpyAuxData * 来传递 get_loop 定义的附加数据。当然,除了使用此结构外,还有其他替代方案,但它提供了一个简单的解决方案,该方案已在 NumPy 和公共 API 中使用。

NpyAyxData * 是一个轻量级、已分配的结构,并且由于它已在 NumPy 中用于此目的,因此似乎是一个自然的选择。为了简化某些用例(请参阅下面的“错误处理”),当未提供 innerloop_data 时,我们将改为传递 npy_intp *innerloop_data = 0

注意: 由于 get_loop 最初预计是私有的,因此我们可以获得 innerloop_data 的经验,然后再将其公开为公共 API。

返回值

用于指示错误的返回值是 NumPy 中一个重要但目前缺失的功能。CPU 信号浮点错误的方式使错误处理更加复杂。这两者将在下一节中讨论。

错误处理#

我们预计未来的内部循环通常会在找到错误后立即设置 Python 错误。当在未锁定 GIL 的情况下运行内部循环时,这会很复杂。在这种情况下,函数将必须锁定 GIL,设置 Python 错误并返回 -1 以指示发生了错误:

int
inner_loop(PyArrayMethod_Context *context, ..., void *innerloop_data)
{
    NPY_ALLOW_C_API_DEF

    for (npy_intp i = 0; i < N; i++) {
        /* calculation */

        if (error_occurred) {
            NPY_ALLOW_C_API;
            PyErr_SetString(PyExc_ValueError,
                "Error occurred inside inner_loop.");
            NPY_DISABLE_C_API
            return -1;
        }
    }
    return 0;
}

浮点错误是特殊的,因为它们需要检查硬件状态,如果它在内部循环函数本身中完成,成本太高。因此,NumPy 将在 ArrayMethod 标记这些错误时处理它们。一个 ArrayMethod 永远不应该引起浮点错误标志被设置,如果它标记这些标志不应该被检查。这可能会干扰调用多个函数;特别是当需要类型转换时。

另一种解决方案是允许仅在 NumPy 稍后进行最终化步骤时设置错误,此时 NumPy 还会检查浮点错误标志。

我们目前决定不采用这种模式。它似乎更复杂且通常不必要。虽然在循环中安全地获取 GIL 可能需要将来传递额外的 PyThreadStatePyInterpreterState(用于子解释器支持),但这可以接受并且可以预见。稍后设置错误会增加复杂性:例如,如果一个操作被暂停(目前可能发生类型转换),错误检查需要每次发生时显式运行。

我们预计立即设置错误是最简单、最方便的解决方案,更复杂的解决方案可能是未来的扩展。

处理警告稍微复杂一些:对于每个函数调用(即,对于整个数组),应该只发出一次警告,即使它可能被多次发出。为了简化此类用例,我们将默认传递 npy_intp *innerloop_data = 0,可用于存储标志(或其他简单的持久数据)。例如,我们可以想象一个整数乘法循环,在发生溢出时发出警告。

int
integer_multiply(PyArrayMethod_Context *context, ..., npy_intp *innerloop_data)
{
    int overflow;
    NPY_ALLOW_C_API_DEF

    for (npy_intp i = 0; i < N; i++) {
        *out = multiply_integers(*in1, *in2, &overflow);

        if (overflow && !*innerloop_data) {
            NPY_ALLOW_C_API;
            if (PyErr_Warn(PyExc_UserWarning,
                    "Integer overflow detected.") < 0) {
                NPY_DISABLE_C_API
                return -1;
            }
            *innerloop_data = 1;
            NPY_DISABLE_C_API
    }
    return 0;
}

待办事项:innerloop_data 未设置时,传递 `npy_intp` 临时空间的概念似乎很方便,但我对此不确定,因为我不知道任何类似的先例。这个“临时空间”原则上也可以是 `context` 的一部分。

重用现有循环/实现#

对于许多 DType,上述添加 C 级附加循环的定义将足够,并且只需要一个 strided 循环实现。如果循环适用于参数化 DType,则 *必须* 另外提供 `resolve_descriptors` 函数。

然而,在某些用例中,希望回调到现有实现。在 Python 中,这可以通过简单地调用原始 ufunc 来实现。

为了在 C 中获得更好的性能,并针对大型数组,我们希望尽可能直接地重用现有的 ArrayMethod,以便其内部循环函数可以直接使用,而不会产生额外的开销。因此,我们将允许从现有的 ArrayMethod 创建一个新的、包装的 ArrayMethod

这个包装的 ArrayMethod 将有两个额外的方法:

  • view_inputs(Tuple[DType]: input_descr) -> Tuple[DType] 将用户输入描述符替换为匹配包装循环的描述符。必须能够将输入视为输出。例如,对于 Unit[Float64]("m") + Unit[Float32]("km"),它将返回 float64 + int32。原始的 resolve_descriptors 将此转换为 float64 + float64

  • wrap_outputs(Tuple[DType]: input_descr) -> Tuple[DType] 将解析后的描述符替换为所需的实际循环描述符。原始的 resolve_descriptors 函数将在两次调用之间调用,因此输出描述符可能不会在第一次调用中设置。在上面的示例中,它将使用返回的 float64(可能已更改字节序),并进一步解析物理单位,从而形成最终签名:

    Unit[Float64]("m") + Unit[Float64]("m") -> Unit[Float64]("m")
    

    UFunc 机制将负责将“km”输入转换为“m”。

`view_inputs` 方法允许将正确的输入传递给原始的 `resolve_descriptors` 函数,而 `wrap_outputs` 确保使用正确的描述符进行输出分配和输入缓冲类型转换。

这一个重要的用例是抽象的 Unit DType,每个数字 dtype 都有子类(可以动态创建)。

Unit[Float64]("m")
# with Unit[Float64] being the concrete DType:
isinstance(Unit[Float64], Unit)  # is True

这样一个 Unit[Float64]("m") 实例相对于类型提升具有明确定义的签名。 `Unit` DType 的作者可以通过包装现有的数学函数并使用上面两个附加方法来实现大部分必要的逻辑。使用提升步骤,这将允许为抽象 `Unit` DType 与 `ufunc` 注册一个单一的提升器。然后,提升器可以在提升时动态添加包装好的具体 `ArrayMethod`,NumPy 可以在第一次调用后缓存(或存储)它。

替代用例

另一个用例是 Unit(float64, "m") DType,其中数值类型是 DType 参数的一部分。这种方法是可行的,但需要一个包装现有循环的自定义 ArrayMethod。它还必须始终需要两个分派步骤(一个到 `Unit` DType,第二个到数值类型)。

此外,高效的实现将需要能够从另一个 `ArrayMethod` 获取和重用内部循环函数。(这对于 Numba 等用户来说可能很重要,但目前尚不确定它是否应该是常见模式,并且它本身无法从 Python 访问。)

提升和分派#

NumPy ufuncs 是多方法,因为它们同时操作(或使用)多个 DType。虽然输入(和输出)dtype 附加到 NumPy 数组,但 `ndarray` 类型本身不携带要应用于数据的信息。

例如,给定输入

int_arr = np.array([1, 2, 3], dtype=np.int64)
float_arr = np.array([1, 2, 3], dtype=np.float64)
np.add(int_arr, float_arr)

必须找到正确的 ArrayMethod 来执行操作。理想情况下,存在一个确切的匹配定义,例如,对于 np.add(int_arr, int_arr), `ArrayMethod[Int64, Int64, out=Int64]` 精确匹配并可以使用。然而,对于 np.add(int_arr, float_arr),没有直接匹配,需要一个提升步骤。

提升和分派过程#

通常,通过搜索所有输入 DType 的确切匹配来找到 ArrayMethod。输出 DType *不应* 影响计算,但如果多个已注册的 ArrayMethod 完全匹配,则将使用输出 DType 来查找更好的匹配。这将允许当前区分 `np.equal` 循环,这些循环定义了 Object, Object -> Bool(默认)和 Object, Object -> Object

最初,ArrayMethod 将仅为*具体* DType 定义,并且由于这些不能被子类化,因此保证了精确匹配。将来,我们预计 ArrayMethod 也可以为*抽象* DType 定义。在这种情况下,最佳匹配的查找方式如下所述。

提升

如果存在匹配的 ArrayMethod,则分派很直接。但是,当不存在时,需要其他定义来实现这种“提升”。

  • 默认情况下,任何 UFunc 都有一个提升,该提升使用所有输入的公共 DType 并再次分派。这对大多数数学函数定义良好,但如果需要,可以禁用或自定义。例如,int32 + float64 会再次尝试使用 float64 + float64,这是公共 DType。

  • 用户可以像注册新 ArrayMethod 一样注册新Promoter。这些将使用抽象 DType 来匹配各种签名。promoter 函数的返回值应为新的 ArrayMethodNotImplemented。它必须对具有相同输入的多个调用保持一致,以允许缓存结果。

promoter 函数的签名将是

promoter(np.ufunc: ufunc, Tuple[DTypeMeta]: DTypes): -> Union[ArrayMethod, NotImplemented]

请注意,DType 可能包含输出的 DType,但通常输出 DType *不会*影响选择哪个 `ArrayMethod`。

在大多数情况下,添加自定义 promotion 函数是不必要的。需要此的一个示例是与单位乘法:在 NumPy 中,`timedelta64` 可以与大多数整数相乘,但 NumPy 只为 `timedelta64 * int64` 定义了一个循环(`ArrayMethod`),因此与 `int32` 相乘将失败。

为了允许这种情况,可以为 `(Timedelta64, Integral, None)` 注册以下 promoter:

def promote(ufunc, DTypes):
    res = list(DTypes)
    try:
        res[1] = np.common_dtype(DTypes[1], Int64)
    except TypeError:
        return NotImplemented

    # Could check that res[1] is actually Int64
    return ufunc.resolve_impl(tuple(res))

在这种情况下,就像需要 `Timedelta64 * int64` 和 `int64 * timedelta64` `ArrayMethod` 一样,需要注册第二个 promoter 来处理整数先传入的情况。

ArrayMethod 和 Promoter 的分派规则

Promoter 和 `ArrayMethod` 的发现是通过查找 DType 类层次结构定义的最优匹配。最优匹配定义为:

  • 签名匹配所有输入 DType,因此 `issubclass(input_DType, registered_DType)` 返回 true。

  • 任何其他 promoter 或 `ArrayMethod` 在任何输入上都不比当前更精确:`issubclass(other_DType, this_DType)` 为 true(这可能包括两者相同)。

  • 此 promoter 或 `ArrayMethod` 在至少一个输入或输出 DType 上更精确。

如果返回 `NotImplemented` 或两个 promoter 匹配输入相同,则会报错。当现有 promoter 不够精确以支持新功能时,必须添加新的 promoter。为确保此 promoter 具有优先权,可能需要定义新的抽象 DType 作为现有 DType 的更精确的子类。

上述规则支持在提供输出或指定完整循环时的专门化。这通常不是必需的,但允许解析 `np.logic_or` 等,它们同时具有 `Object, Object -> Bool` 和 `Object, Object -> Object` 循环(默认使用第一个)。

讨论和替代方案#

与其解析并返回一个新的实现,我们还可以返回一组新的 DType 来进行分派。这样做可行,但缺点是无法分派到定义在不同 ufunc 上的循环,也无法动态创建新的 `ArrayMethod`。

被否决的替代方案

在上述方案中,promoter 使用多重分派风格的类型解析,而当前的 UFunc 机制使用有序层次结构中的第一个“安全”循环(另请参阅 NEP 40)。

虽然“安全”类型转换规则不够严格,但我们可以设想使用新的“提升”类型转换规则,或使用公共 DType 逻辑来通过根据需要向上提升输入来找到最佳匹配循环。

这种方法的一个缺点是,向上提升可能导致结果超出用户的预期:目前(并且将继续作为后备方案支持),任何仅定义了 float64 循环的 ufunc,通过*向上提升*也能用于 float16 和 float32。

>>> from scipy.special import erf
>>> erf(np.array([4.], dtype=np.float16))  # float16
array([1.], dtype=float32)

将 `erf` 函数更改为返回 float16 结果而不更改以下代码的结果是不可能的。总的来说,我们认为自动向上提升不应发生在可以定义精度较低的循环的情况下,*除非* ufunc 作者有意通过提升来实现。

这个考虑意味着向上提升必须受到某种附加方法的限制。

替代方案 1

假设不打算进行通用向上提升,必须定义一个规则来限制从 `float16 -> float32` 向上提升输入,方法是使用 DType 的通用逻辑或 UFunc 本身(或两者的组合)。UFunc 本身无法轻松做到这一点,因为它不知道所有注册了循环的 DType。考虑以下两个示例:

第一个(应被拒绝)

  • 输入:float16 * float16

  • 现有循环:float32 * float32

第二个(应被接受)

  • 输入:timedelta64 * int32

  • 现有循环:timedelta64 * int16

这需要:

  1. `timedelta64` 以某种方式发出信号,表示如果 `int64` 向上提升参与操作,则总是支持它。

  2. `float32 * float32` 循环拒绝向上提升。

实现第一种方法需要发出信号,表示向上提升在特定上下文中是可接受的。这需要额外的钩子,并且对于复杂 DType 来说可能并不简单。

对于第二种方法,在大多数情况下,简单的 `np.common_dtype` 规则将适用于初始分派,但即使这样,对于同质循环也只是勉强适用。此选项需要添加一个函数来检查输入是否是每个循环的有效向上提升,这似乎有问题。在许多情况下,可以提供一个默认值(同质签名)。

替代方案 2

另一种“提升”步骤是确保在首次找到正确的输出 DType 后,*输出* DType 与循环匹配。如果输出 DType 已知,则找到安全循环将变得容易。在大多数情况下,这很有效,正确的输出 DType 只是

np.common_dtype(*input_DTypes)

或者某个固定的 DType(例如,逻辑函数的 Bool)。

然而,例如,在上面的 `timedelta64 * int32` 情况下,这是无效的,因为先验无法知道此输出的“预期”结果类型实际上是 `timedelta64`(`np.common_dtype(Datetime64, Int32)` 失败)。这需要关于 `timedelta64` 精度为 int64 的一些额外知识。由于 ufunc 可以有任意数量的(相关)输入,因此至少需要 `Datetime64`(以及所有 DType)上有一个额外的 `__promoted_dtypes__` 方法。

另一个限制是掩码 DType。当涉及掩码时,逻辑函数不会返回布尔结果,因此需要原始 ufunc 作者在这种方案中预测掩码 DType。类似地,为复数值定义的某些函数返回实数,而另一些则返回复数。如果原始作者没有预测复数,那么对于稍后添加的复数循环,提升可能会不正确。

我们认为 promoter 是最好的解决方案,尽管它们允许巨大的理论复杂性。

  1. 提升允许动态添加新循环。例如,可以定义一个抽象 Unit DType,它动态创建类来包装其他现有 DType。使用单个 promoter,此 DType 可以动态包装现有的 `ArrayMethod`,使其能够在一个查找中而不是两个查找中找到正确的循环。

  2. 提升逻辑通常会保守行事:除非添加了 promoter,否则新添加的循环不会被误用。

  3. 它们将仔细思考逻辑是否正确的负担放在了向 UFunc 添加新循环的程序员身上。(与替代方案 2 相比)。

  4. 如果现有提升不正确,可以编写一个 promoter 来限制或精炼通用规则。通常,promotion 规则不应返回*不正确*的 promotion,但如果现有 promotion 逻辑对新添加的循环失败或不正确,则该循环可以添加新的 promoter 来精炼逻辑。

让每个循环验证没有发生向上提升的选项可能是最好的替代方案,但它不包括动态添加新循环的能力。

通用 promoter 的主要缺点是它们允许可能非常大的复杂性。第三方库*可能*会向 NumPy 添加不正确的 promotion,但这已经可以通过添加新的不正确循环来完成。总的来说,我们认为我们可以信赖下游项目谨慎负责地使用这种力量和复杂性。

用户指南#

通常,向 UFunc 添加 promoter 必须非常小心。promoter 永远不应影响可以由其他数据类型合理定义的循环。定义一个假设的 `erf(UnitFloat16)` 循环不应导致 `erf(float16)`。总的来说,promoter 应满足以下要求:

  • 定义新 promotion 规则时要保守。不正确的结果比意外的错误更危险。

  • 添加的(抽象)DType 之一应专门匹配您项目定义的 DType(或 DType 系列)。切勿添加超出正常公共 DType 规则的 promotion 规则!为 `int16 + uint16 -> int24` 定义循环,如果您编写了一个 `int24` dtype,这是*不合理*的。此操作的结果先前已定义为 `int32`,并将以此假设使用。

  • promoter(或循环)不应影响现有循环结果。这包括添加更快但精度较低的循环/promoter 以替换现有的循环/promoter。

  • 尝试在所有 promotion(和类型转换)相关逻辑中保持清晰、线性的层次结构。NumPy 本身破坏了整数和浮点数的逻辑(它们不严格线性,因为 int64 不能提升到 float32)。

  • 循环和 promoter 可以由任何项目添加,可以是:

    • 定义 ufunc 的项目

    • 定义 DType 的项目

    • 第三方项目

    尝试找出添加循环的最佳项目。如果定义 ufunc 的项目和定义 DType 的项目都没有添加循环,可能会出现多个定义(被拒绝)的问题,并且应注意确保循环行为始终比错误更可取。

在某些情况下,这些规则的例外情况可能是合理的,但总的来说,我们要求您极其谨慎,如有疑问,请创建新的 UFunc。这会清楚地告知用户不同的规则。如有疑问,请在 NumPy 邮件列表或问题跟踪器上提问!

实现#

此 NEP 的实现将涉及对当前 ufunc 机制(以及类型转换)的大规模重构和重新组织。

实现不幸地需要对 UFunc 机制进行大量维护,因为实际的 UFunc 循环调用以及初始分派步骤都必须进行修改。

总的来说,正确的 `ArrayMethod`,包括由 promoter 返回的,将被缓存(或存储)在一个哈希表中以供高效查找。

讨论#

存在大量可能的实现,在各种地方进行了许多讨论,以及初步的想法和设计文档。这些列在 NEP 40 的讨论中,为简洁起见在此不重复。

一项涉及上述许多要点并指向类似解决方案的长篇讨论可以在 github issue 12518 “What should be the calling convention for ufunc inner loop signatures?” 中找到。

参考文献#

有关更多讨论和参考文献,请参阅 NEP 40 和 41。