NumPy源码解析

Numpy源码解析

PyArray_Type & PyArrayObject

numpy中的ndarray对象实际对应底层C语言中的 PyArrayObject 对象，PyArrayObject结构体的前面部分同 PyObject 兼容，所以可以被cast成 PyObject，当做一般的python对象处理。这也是创建Python新对象的通用方式。

PyArrayObject定义如下：

typedef struct tagPyArrayObject_fields {
    PyObject ob_base;
    char *data;
    int nd;
    npy_intp *dimensions;
    npy_inpt *strides;
    PyObject *base;
    PyArray_Descr *desc;
    int flags;
    PyObject *weakreflist;
    void *_buffer_info;
    PyObject *mem_handler;
} PyArrayObject_fields;

typedef struct tagPyArrayObject {
    PyObject ob_base;
} PyArrayObject;

注： 没有直接将新的属性放置到 PyArrayObject上，而是使用了一个内部的数据结构 PyArrayObject_fields。猜测应该是进一步对上层用户封装细节，避免直接使用PyArrayObject类型操作底层数据。

PyArray_Type 是一个PyTypeObject类对象，代表ndarray的类型，采用单例模型，全局只有一个对象。内部成员一般是描述该类型属性和方法的函数指针。
PyArray_Type的定义如下：

PyTypeObject PyArray_Type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "numpy.ndarray",
    .tp_basicsize = sizeof(PyArrayObject_fields),
    /* methods */
    .tp_dealloc = (destructor)array_dealloc,
    .tp_repr = (reprfunc)array_repr,
    .tp_as_number = &array_as_number,
    .tp_as_sequence = &array_as_sequence,
    .tp_as_mapping = &array_as_mapping,
    .tp_str = (reprfunc)array_str,
    .tp_as_buffer = &array_as_buffer,
    .tp_flags =(Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE),

    .tp_richcompare = (richcmpfunc)array_richcompare,
    .tp_weaklistoffset = offsetof(PyArrayObject_fields, weakreflist),
    .tp_iter = (getiterfunc)array_iter,
    .tp_methods = array_methods,
    .tp_getset = array_getsetlist,
    .tp_new = (newfunc)array_new,
};

注： 在numpy中，该对象是 __attribute__((visibility("hidden"))) 的，无法外部直接获取。

ufunc

描述

universal functions (ufunc) 官方定义，中文官网：

通用函数（或简称为ufunc） 是一种以逐元素方式操作ndarrays的函数，支持数组广播，类型转换和其他一些标准功能。也就是说，ufunc是一个函数的 “ 矢量化 ” 包装器，它接受固定数量的特定输入并产生固定数量的特定输出。

在NumPy中，通用函数是numpy.ufunc类的实例 。 许多内置函数都是在编译的C代码中实现的。 基本的对标量的ufuncs操作，但也有一种通用类型，基本元素是子数组（向量，矩阵等）， 广播是在其他维度上完成的。也可以使用numpy.frompyfunc工厂函数生成自定义numpy.ufunc实例。

numpy中绝大部分API都是通过ufunc机制实现的。

numpy底层调用的各类c函数，会被封装成ufunc对象。ufunc本身就是个Python对象，但大部分内置 ufunc对象是通过底层C代码定义的。
numpy在编译阶段判断当前系统硬件信息和依赖满足情况，自动编译最优的底层实现，然后封装到ufunc中.

ufunc可以分为两类：

• scalar ufunc: 基础的ufunc通常针对ndarray元素(scalar)进行操作
• generalized ufunc: 可以针对 sub-array 进行操作

ufunc类型预先定义了5个方法：

方法	描述
ufunc.reduce(array, [, axis, dtype, out, …])	沿指定轴使用reduce 一个个减少array的维度(通过累积成1个值降维)
ufunc.accumulate(array, [, axis, dtype, out])	沿指定轴使用该操作到所有元素的累积结果(每个元素前的累积，维度不变)
ufunc.reduceat(array, indices[, axis, …])	在指定轴上使用指定切片执行（局部）缩减
ufunc.outer(A, B, /, **kwargs)	将ufunc op应用于所有（a，b）对，其中A中的a和B中的b
ufunc.at(a, indices[, b])	对 index 指定的元素在操作数 a 上执行无缓冲的就地操作

具体讲解可参考numpy矩阵和通用函数第5.7节

说明： 如果某个输入参数定义了 __array_ufunc__，则优先调用该参数定义的ufunc方法，即ufunc被overrides了。

ufunc调用过程中会对shape不匹配的输入进行 Broadcasting；ufunc调用过程中，使用 Type casting rules

内部缓存机制

在内部，缓冲区用于未对齐的数据，交换的数据以及必须从一种数据类型转换为另一种数据类型的数据。内部缓冲区的大小可以基于每个线程设置。最多可以 创建指定大小$2(n_{inputs}+n_{outputs})$的缓冲区来处理来自ufunc的所有输入和输出的数据。缓冲区的默认大小为10,000个元素。每当需要基于缓冲区的计算，但所有输入数组都小于缓冲区大小时，将在计算进行之前复制那些行为不当或类型不正确的数组。因此，调整缓冲区的大小可能会改变各种类型的ufunc计算完成的速度。用于设置此变量可见的简单接口：

方法	描述
numpy.setbufsize(size)	设置ufuncs中使用的缓冲区的大小

数据优先会放入到buffer中，然后在调用底层loop函数。numpy内部定义的ufunc函数列表，结合中文官网更香。

遗留问题：

问题1：不同的ufunc的buffer是否独立？
问题2：buffer的生命周期，是否在ufunc第一次调用就分配，然后一直存在，增长或者减少规则？

ufunc数据结构

其实际类型为 PyUFuncObject，数据结构如下

typedef void (*PyUFuncGenericFunction)
            (char **args,
             npy_intp const *dimensions,
             npy_intp const *strides,
             void *innerloopdata);

struct PyUFuncObject {
  PyObject ob_base;

  /*
   * nin: Number of inputs
   * nout: Number of outputs
   * nargs: Always nin + nout (Why is it stored?)
   */
  int nin, nout, nargs;

  /*
   * Identity for reduction, any of PyUFunc_One, PyUFunc_Zero
   * PyUFunc_MinusOne, PyUFunc_None, PyUFunc_ReorderableNone,
   * PyUFunc_IdentityValue.
   */
  int identity;

  /* Array of one-dimensional core loops */   // 一组接受不同类型输入参数的同功能函数
  PyUFuncGenericFunction *functions;
  /* Array of funcdata that gets passed into the functions */
  void **data;
  /* The number of elements in 'functions' and 'data' */
  int ntypes;

  /* Used to be unused field 'check_return' */
  int reserved1;

  /* The name of the ufunc */
  const char *name;
  /* Array of type numbers, of size ('nargs' * 'ntypes') */
  char *types;
  /* Documentation string */
  const char* doc;

  void *ptr;
  PyObject *obj;
  PyObject *userloops;

  /* generalized ufunc parameters */

  /* 0 for scalar ufunc; 1 for generalized ufunc */
  int core_enabled;
  /* number of distinct dimension names in signature */
  int core_num_dim_ix;

  /*
   * dimension indices of input/output argument k are stored in
   * core_dim_ixs[core_offsets[k]..core_offsets[k]+core_num_dims[k]-1]
   */

  /* numbers of core dimensions of each argument */
  int *core_num_dims;
  /*
   * dimension indices in a flatted form; indices
   * are in the range of [0,core_num_dim_ix)
   */
  int *core_dim_ixs;
  /*
   * positions of 1st core dimensions of each
   * argument in core_dim_ixs, equivalent to cumsum(core_num_dims)
   */
  int *core_offsets;
  /* signature string for printing purpose */
  int *core_signature;

  /*
   * A function which resolves the types and fills an array
   * with the dtypes for the inputs and outputs.
   */
  PyUFunc_TypeResolutionFunc *type_resolver;
  /*
   * A function which returns an inner loop written for
   * NumPy 1.6 and earlier ufuncs. This is for backwards
   * compatibility, and may be NULL if inner_loop_selector
   * is specified.
   */
  PyUFunc_LegacyInnerLoopSelectionFunc *legacy_inner_loop_selector;

  vectorcallfunc vectorcall;

  /* Was previously the `PyUFunc_MaskedInnerLoopSelectionFunc` */
  void *_always_null_previously_masked_innerloop_selector;

  /*
   * List of flags for each operand when ufunc is called by nditer object.
   * These flags will be used in addition to the default flags for each
   * operand set by nditer object.
   */
  npy_uint32 *op_flags;
  /*
   * List of global flags used when ufunc is called by nditer object.
   * These flags will be used in addition to the default global flags
   * set by nditer object.
   */
  npy_uint32 iter_flags;

  /* New in NPY_API_VERSION 0x0000000D and above */

  /*
   * for each core_num_dim_ix distinct dimension names,
   * the possible "frozen" size (-1 if not frozen).
   */
  npy_intp *core_dim_sizes;
  /*
   * for each distinct core dimension, a set of UFUNC_CORE_DIM* flags
   */
  npy_uint32 *core_dim_flags;

  /* Identity for reduction, when identity == PyUFunc_IdentityValue */
  PyObject *identity_value;

  /* New in NPY_API_VERSION 0x0000000F and above */

  /* New private fields related to dispatching */
  void *_dispatch_cache;
  /* A PyListObject of `(tuple of DTypes, ArrayMethod/Promoter)` */
  PyObject *_loops;
};

关于 core_enable 的作用理解：
当 core_enable == 0 代表标量ufunc，是逐元素(element-wise)操作的。
当 core_enable == 1 代表广义ufunc(generalized universal function, gufunc)，支持逐子数组(array-wise)操作，参考Generalized Universal Function。

如内积函数inner1d(a, b)，签名为 (i),(i)->().这将沿每个输入的最后一个轴应用内积，但保持其余轴不变。例如，其中 a 是形状的 (3, 5, N) 和 b 是形状的 (5, N)，这将返回 (3,5)形状的输出.底层的基础函数被调用 3 * 5 次。在签名中，我们对于每个输入指定一个core维度 (i) 和对于输出指定零core维度 () ，因为它接受两个1-D数组并返回一个标量。通过使用相同的名字 i，我们指定两个对应的维度应该具有相同的大小

对应的 PyTypeObject 对象为 PyUFunc_Type，在PyUFunc_Type 定义了5个类型方法，具体参考前文

signature和core dimensions

参考Generalized Universal Function API

Signature:
描述ufunc基本函数的输入/输出维度的字符串

Core Dimension:
函数的每个输入/输出的维数由其核维数（零核维数对应于标量输入/输出）定义。核心维度映射到输入/输出数组的最后一个维度

signature从输入参数的shape最内层维度开始匹配，指示出的维度是 core dimensions，输入输出必须严格遵守。
超出core dimensions的高维度叫loop维度。对于上文inner1d(a,b)的例子，N为输入的core dimension，外层的(3,5)即为loop维度。

签名确定如何将每个输入/输出数组的维度拆分为核心维度和循环维度：

1. 签名中的每个维度都匹配到相应传入数组的一个维度，从形状元组的末尾开始。这些是核心维度，它们必须存在于数组中，否则将引发错误。
2. 在签名中分配给同一标签的core dimension(例如 i 在 inner1d’的 (i),(i)->())的大小必须完全匹配，不执行广播。
3. 将从所有输入中删除核心维度，并将剩余维度一起广播，从而定义循环维度。
4. 每个输出的形状由loop维度加上输出的core维度决定

__ufunc_protocol__和__array_protocol__

numpy.class.__array_ufunc__(ufunc,method,*inputs,**kwargs)

当调用numpy的ufunc时(使用np.add(a, b)，a+b, a.__add__(b)等方式)，在执行具体的ufunc调用流程前，会首先检查下输入/输出参数，看他们是否定义了 __array_ufunc__属性。并且根据一定的规则，确定谁的__array_ufunc__属性优先级最高，然后执行它。
规则为：

1. 如果输入或输出参数之一具有 __array_ufunc__方法，则执行该方法而不是ufunc。
2. 如果多个参数实现 __array_ufunc__，则按顺序尝试它们：子类在超类之前，输入在输出之前，否则从左到右。
3. 返回的第一个非 NotImplemented 的方法确定结果。
4. 如果所有 __array_ufunc__ 方法都返回 NotImplemented，则会引发 TypeError
5. 如果类的有__array_ufunc__属性但它不等同于ndarray.__array_func__，这个属性被忽略

ndarray的默认 array_ufunc 实现可以如下

def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):     # 这里的inputs应该包含self
    # Cannot handle items that have __array_ufunc__ (other than our own).
    outputs = kwargs.get('out', ())
    for item in inputs + outputs:
        if (hasattr(item, '__array_ufunc__') and
                type(item).__array_ufunc__ is not ndarray.__array_ufunc__):
            return NotImplemented

    # If we didn't have to support legacy behaviour (__array_prepare__,
    # __array_wrap__, etc.), we might here convert python floats,
    # lists, etc, to arrays with
    # items = [np.asarray(item) for item in inputs]
    # and then start the right iterator for the given method.
    # However, we do have to support legacy, so call back into the ufunc.
    # Its arguments are now guaranteed not to have __array_ufunc__
    # overrides, and it will do the coercion to array for us.
    return getattr(ufunc, method)(*items, **kwargs)

numpy.class__array_function__(func, types, args, kwargs)

给用户提供了一种机制，重载numpy的python API的默认实现（通常为了适配一种新的数组对象）。同时无需修改已有的应用了numpy的工程代码

**背景：**Numpy的API非常经典，外部有很多不同的高维数组实现都参考了Numpy的API，比如GPU数组(CuPy)、稀疏数组(scipy.sparse)、并行数组(Dask array)等。甚至Tensorflow和Pytorch中的Tensor也可以囊括在内。同时当前很多实用的高阶库基于Numpy API已经实现了，比如label and indexed arrays (XArray), automatic differentiation (Autograd, Tangent), maked arrays (numpy.ma), physical unit (astropy.units, pint)等。现在考虑如何将两者结合起来？
核心思路就是，提供一种快捷的机制，能够替换掉numpy的API的实现，同时不会改变API名字不变。如此就能将CuPy用于Autograd了。

• func是一个由NumPy的公开API暴露出的任意可调用的函数，以func(*args, **kwargs)形式调用
• types来自实现 __array_function*的原始NumPy函数调用的唯一参数类型的集合
• 元组 args 和dict kwargs 直接从原始调用传递

__array_function__ 的大多数实现都将从两个检查开始：

1. 给定的函数是否知道如何重载？
2. 我们知道如何处理的类型的所有参数都是？

如果这些条件成立，__array_function__ 应该返回调用 func(*args, **kwargs) 实现的结果。 否则，它应该返回标记值 NotImplemented，表示该函数未由这些类型实现

此协议旨在为那些未被__array_ufunc__协议涵盖的ufunc (比如 np.exp)提供全面的Numpy功能.
有些偏底层函数无法采用 __array_function__机制替换。

使用自定义装饰器(下面的implements)来方便的为numpy函数注册__array_function__

HANDLED_FUNCTIONS = {}

class MyArray:
    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        # Note: this allows subclasses that don't override
        # __array_function__ to handle MyArray objects
        if not all(issubclass(t, MyArray) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)

def implements(numpy_function):
    """Register an __array_function__ implementation for MyArray objects."""
    def decorator(func):
        HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

@implements(np.concatenate)
def concatenate(arrays, axis=0, out=None):
    ...  # implementation of concatenate for MyArray objects

@implements(np.broadcast_to)
def broadcast_to(array, shape):
    ...  # implementation of broadcast_to for MyArray objects

在大多数情况下，__array_function__ 的调度规则与 __array_ufunc__ 的调度规则相匹配。尤其是：

• NumPy将从所有指定的输入中收集 __array_function__ 的实现，并按顺序调用它们：超类之前的子类，否则从左到右。 请注意，在涉及子类的某些边缘情况下，这与Python的当前行为略有不同。
• __array_function__ 的实现表明它们可以通过返回除 NotImplemented 之外的任何值来处理操作。
• 如果所有 __array_function__ 方法都返回 NotImplemented，NumPy 将引发 TypeError

与 __array_ufunc__ 的当前行为的一个偏差是，NumPy将仅对每个唯一类型的第一个参数调用 __array_function__。 这与Python调用反射方法的规则相匹配，这确保了检查重载具有可接受的性能，即使存在大量重载参数

结合自定义数组容器例子理解两个协议

数学运算

主要在_multiarray_umath 扩展中实现的，定义了一组数学基础运算，并且使用 n_ops对象存储起来。如下

struct NumericOps {
  PyObject *add;
  PyObject *subtract;
  PyObject *multiply;
  PyObject *divide;
  PyObject *remainder;
  PyObject *divmod;
  PyObject *power;
  PyObject *square;
  PyObject *reciprocal;
  PyObject *_ones_like;
  PyObject *sqrt;
  PyObject *cbrt;
  PyObject *negative;
  PyObject *positive;
  PyObject *absolute;
  PyObject *invert;
  PyObject *left_shift;
  PyObject *right_shift;
  PyObject *bitwise_and;
  PyObject *bitwise_xor;
  PyObject *bitwise_or;
  PyObject *less;
  PyObject *less_equal;
  PyObject *equal;
  PyObject *not_equal;
  PyObject *greater;
  PyObject *greater_equal;
  PyObject *floor_divide;
  PyObject *true_divide;
  PyObject *logical_or;
  PyObject *logical_and;
  PyObject *floor;
  PyObject *ceil;
  PyObject *maximum;
  PyObject *minimum;
  PyObject *rint;
  PyObject *conjugate;
  PyObject *matmul;
  PyObject *clip;
} NumericOps;

NumericOps n_ops;

其中的每个PyObject实际对应一个ufunc，即 PyUFuncObject 对象。

而 n_ops 的初始化则是在扩展载入时进行的，定义在 PyInit_multiarray_umath 方法中(按照官方教程定义Python扩展时，PyInit_<extensionname> 会在扩展载入时自动执行)

PyObject* PyInit__multiarray_umath(void) {
  PyObject *m = PyModule_Create(&moduledef);     // 当前的扩展模块

  PyObject *d = PyModule_GetDict(m);                  // 获取模块的属性字典 __dict__ 
  // ...
  initumath(m);                                     // 该方法中初始化 n_ops 对象
  // ...
}

注意在调用 initumath 方法时，会为扩展模块的 __dict__ 添加基础运算的ufunc对象(调用InitOperators(PyObject *dictionary)方法)，然后将对应的ufunc对象赋值给n_ops(调用_PyArray_SetNumericOps方法)。

先来看看 ufunc对象是如何添加到模块的属性字典中的？
首先 InitOperators 方法定义在 __umath_generated.c 文件中，该文件是在编译阶段自动生成的。numpy会在编译前判断当前系统或者硬件信息和底层依赖情况，然后选择并生成最优的底层实现。所以numpy在编译阶段的前面会添加代码生成的步骤。

以笔者自己环境中编译的 add 方法为例，介绍其如何封装成为ufunc对象，并且添加到模块的属性字典中

static PyUFuncGenericFunction add_functions[] = 
  {BOOL_add, BYTE_add, UBYTE_add, SHORT_add, USHORT_add, INT_add, UINT_add,
   LONG_add, ULONG_add, LONGLONG_add, ULONGLONG_add, HALF_add, FLOAT_add, 
   DOUBLE_add, LONGDOUBLE_add, CFLOAT_add, CDOUBLE_add, CLONGDOUBLE_add, 
   DATETIME_Mm_M_add, TIMEDELTA_mm_m_add, DATETIME_mM_M_add, NULL};
static void * add_data[] = 
  {(void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, 
   (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, 
   (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL, 
   (void *)NULL, (void *)NULL, (void *)NULL, (void *)NULL};
// 对应函数的输入输出参数，NPY_BOOl 来自一个enum类型
static char add_signatures[] = 
  {NPY_BOOL, NPY_BOOL, NPY_BOOL, 
   NPY_BYTE, NPY_BYTE, NPY_BYTE,
   NPY_UBYTE, NPY_UBYTE, NPY_UBYTE,
   NPY_SHORT, NPY_SHORT, NPY_SHORT,
   NPY_USHORT, NPY_USHORT, NPY_USHORT,
   NPY_INT, NPY_INT, NPY_INT,
   NPY_UINT, NPY_UINT, NPY_UINT,
   NPY_LONG, NPY_LONG, NPY_LONG,
   NPY_ULONG, NPY_ULONG, NPY_ULONG,
   NPY_LONGLONG, NPY_LONGLONG, NPY_LONGLONG, 
   NPY_ULONGLONG, NPY_ULONGLONG, NPY_ULONGLONG,
   NPY_HALF, NPY_HALF, NPY_HALF,
   NPY_FLOAT, NPY_FLOAT, NPY_FLOAT,
   NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
   NPY_LONGDOUBLE, NPY_LONGDOUBLE, NPY_LONGDOUBLE,
   NPY_CFLOAT, NPY_CFLOAT, NPY_CFLOAT,
   NPY_CDOUBLE, NPY_CDOUBLE, NPY_CDOUBLE,
   NPY_CLONGDOUBLE, NPY_CLONGDOUBLE, NPY_CLONGDOUBLE,
   NPY_DATETIME, NPY_TIMEDELTA, NPY_DATETIME, 
   NPY_TIMEDELTA, NPY_TIMEDELTA, NPY_TIMEDELTA, 
   NPY_TIMEDELTA, NPY_DATETIME, NPY_DATETIME,
   NPY_OBJECT, NPY_OBJECT, NPY_OBJECT};

这是add_functions的默认值，函数的实现通常在生成的代码中（如 loops_arithm_fp.dispatch.c 中）找到

static int
InitOperators(PyObject *dictionary) {
  // ...
  // 进行一系列调整
  if (NPY_CPU_HAVE(AVX2)) {
    add_functions[1] = BYTE_add_avx2;
  }
  if (NPY_CPU_HAVE(AVX2)) {
    add_functions[2] = UBYTE_add_avx2;
  }
  //...
  add_functions[21] = PyUFunc_OO_O;
  add_data[21] = (void *) PyNumber_Add;   // PyNumber_Add是个函数，应该是用于自定义对象的相加

  // 针对 Float Double CFloat CDouble 的调整
  /*
   * NPY_CPU_DISPATCH_CALL_XB宏对应如下内容
   * ((npy_cpu_have(NPY_CPU_FEATURE_AVX512F))) ? (void) (add_functions[12] = FLOAT_add_AVX512F ) : ((npy_cpu_have(NPY_CPU_FEATURE_AVX)&&npy_cpu_have(NPY_CPU_FEATURE_F16C)&&npy_cpu_have(NPY_CPU_FEATURE_AVX2))) ? (void) (add_functions[12] = FLOAT_add_AVX2 ) : ((void) 0 );
   */
  NPY_CPU_DISPATCH_CALL_XB(add_functions[12] = FLOAT_add);
  NPY_CPU_DISPATCH_CALL_XB(add_functions[13] = DOUBLE_add);
  NPY_CPU_DISPATCH_CALL_XB(add_functions[15] = CFLOAT_add);
  NPY_CPU_DISPATCH_CALL_XB(add_functions[16] = CDOUBLE_add);

  // 创建 ufunc 对象
  PyObject *f;
  f = PyUFunc_FromFuncAndDataAndSignatureAndIdentity(
        add_functions,   // 对应 functions
        add_data,          // 对应 data
        add_signatures,   // 对应 types
        22,   // 有几组signatures，也就是有几个functions，对应 ntypes
        2,    // nin
        1,    // nout
         PyUFunc_IdentityValue,   // 值为 -3, 对应 identity
         "add",      // 对应name
         DOC_NUMPY_CORE_UMATH_ADD,  // 对应 doc
         0,
         NULL,            // 对应signatures，字符串，可以对该字符串解析，影响ufunc的 core_enabled, core_num_dim_ix, core_num_dims, core_offsets, core_dim_ixs, core_dim_sizes, core_dim_flags 等值。
         identiry          // 对应 identity_value，这里仅仅是个值为 0 的python对象，不知何用
   );
  ((PyUFuncObject *)f)->type_resolver = &PyUFunc_AbsoluteTypeResolver;
  PyDict_SetItemString(dictionary, "add", f);
  Py_DECREF(f);
}

如此 add方法对应的ufunc对象构建完毕，并且加入到模块的属性字典中。

n_ops对象的作用其实不大，仅仅是缓存对应的ufunc对象，方便具体使用时快速获取

static PyObject *
array_add(PyObject *m1, PyObject *m2)
{
    PyObject *res;

    // 调用 n_ops.add 方法作用于 m1, m2
    return PyArray_GenericBinaryFunction(m1, m2, n_ops.add);
}

array_add 方法会作为ndarray对象的类型方法。通过如下方式添加

PyNumberMethods array_as_number = {
    .nb_add = array_add,
    .nb_subtract = array_subtract,
    .nb_multiply = array_multiply,
    .nb_remainder = array_remainder,
    .nb_divmod = array_divmod,
    .nb_power = (ternaryfunc)array_power,
    .nb_negative = (unaryfunc)array_negative,
    .nb_positive = (unaryfunc)array_positive,
    .nb_absolute = (unaryfunc)array_absolute,
    .nb_bool = (inquiry)_array_nonzero,
    .nb_invert = (unaryfunc)array_invert,
    .nb_lshift = array_left_shift,
    .nb_rshift = array_right_shift,
    .nb_and = array_bitwise_and,
    .nb_xor = array_bitwise_xor,
    .nb_or = array_bitwise_or,

    .nb_int = (unaryfunc)array_int,
    .nb_float = (unaryfunc)array_float,
    .nb_index = (unaryfunc)array_index,

    .nb_inplace_add = (binaryfunc)array_inplace_add,
    .nb_inplace_subtract = (binaryfunc)array_inplace_subtract,
    .nb_inplace_multiply = (binaryfunc)array_inplace_multiply,
    .nb_inplace_remainder = (binaryfunc)array_inplace_remainder,
    .nb_inplace_power = (ternaryfunc)array_inplace_power,
    .nb_inplace_lshift = (binaryfunc)array_inplace_left_shift,
    .nb_inplace_rshift = (binaryfunc)array_inplace_right_shift,
    .nb_inplace_and = (binaryfunc)array_inplace_bitwise_and,
    .nb_inplace_xor = (binaryfunc)array_inplace_bitwise_xor,
    .nb_inplace_or = (binaryfunc)array_inplace_bitwise_or,

    .nb_floor_divide = array_floor_divide,
    .nb_true_divide = array_true_divide,
    .nb_inplace_floor_divide = (binaryfunc)array_inplace_floor_divide,
    .nb_inplace_true_divide = (binaryfunc)array_inplace_true_divide,

    .nb_matrix_multiply = array_matrix_multiply,
    .nb_inplace_matrix_multiply = (binaryfunc)array_inplace_matrix_multiply,
};

PyTypeObject PyArray_Type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "numpy.ndarray",
    .tp_basicsize = sizeof(PyArrayObject_fields),
    /* methods */
    .tp_dealloc = (destructor)array_dealloc,
    .tp_repr = (reprfunc)array_repr,
    .tp_as_number = &array_as_number,
    .tp_as_sequence = &array_as_sequence,
    .tp_as_mapping = &array_as_mapping,
    .tp_str = (reprfunc)array_str,
    .tp_as_buffer = &array_as_buffer,
    .tp_flags =(Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE),

    .tp_richcompare = (richcmpfunc)array_richcompare,
    .tp_weaklistoffset = offsetof(PyArrayObject_fields, weakreflist),
    .tp_iter = (getiterfunc)array_iter,
    .tp_methods = array_methods,
    .tp_getset = array_getsetlist,
    .tp_new = (newfunc)array_new,
};

如此ndarray本身的基础运算方法定义完毕。

ufunc如何实现类型dispatch

PyUFuncObject::functions 指向一组不同类型的同功能函数。比如加法的实现：

TODO

loop类和通用类函数的调用

loop类方法以 add 方法为例
注意 array_add 方法对应了 ndarray 对象的 + 操作符。

loop类函数的调用入口是 vectorcall

ufunc_generic_vectorcall
  ufunc_generic_fastcall
    PyUFunc_GenericFunctionInternal
      execute_ufunc_loop
        generic_wrapped_legacy_loop
          DOUBLE_add_AVX512F

通用类方法以 matmul 方法为例，注意入口同样是 vectorcall

ufunc_generic_vectorcall
  ufunc_generic_fastcall
    PyUFunc_GeneralizedFunctionInternal   // 这里同 loop 类方法有所不同
      generic_wrapped_legacy_loop
        DOUBLE_matmul

reduce类函数调用

以 max 方法为例

首先在method.c 中定义了 array_max 方法，根据习惯，array_max 方法对应 ndarray 对象的成员方法。当执行如下代码时，会调用底层的 array_max 方法

a = np.array([1.0, 2.0, 3.0])
a.max()

直接调用 np.max(a) 则不会。

array_max代码：

static PyObject *
array_max(PyArrayObject *self, PyObject *args, PyObject *kwds)
{
    NPY_FORWARD_NDARRAY_METHOD("_amax");
}

/*
static PyObject *callable = NULL;
npy_cache_import("numpy.core._methods", "_amax", &callable);
return forward_ndarray_method(self, args, kwds, callable);
本质是转发调用 python 代码的 _amax 方法。注意 np.max应该是直接调用的 _amax方法
*/

在python代码中采用了特殊的方法，可以执行到一套统一的reduce执行流程。这里仅关注C代码部分

ufunc_reduce()
  PyUFunc_GenericReduction()
    PyUFunc_Reduce()
      PyUFunc_ReduceWrapper()
        reduce_loop()
          generic_wrapped_legacy_loop()
            DOUBLE_maximum_AVX512_SKX()  # 具体实现

迭代器的优化和使用

numpy ndarray迭代器C API，相关设计提案

numpy中的ndarray迭代器是根据底层memory layout来迭代数据的，可以选择 “C” order或者 “F” order，默认是以最高效(cpu cache friendly，可以交换axis顺序)的方式返回。这里默认为使用迭代器是对所有数据进行操作，对先后顺序不会有严格要求的。
同时ndarray可以选择忽略最后一维的迭代，循环内处理一组数据。同时支持连续排列数据维度的collapse。以及数据的buffer。

ufunc的调用过程就是依赖于ndarray的迭代器实现高效的循环处理的。数据可以在迭代器中缓存，可以让数据处理 type/byte-order/alignment matchment。