强制 NumPy ndarray 在 Cython 中拥有其内存

Question

根据这个回答，我想通过Cython的NumPy包装器使用Python C API函数PyArray_ENABLEFLAGS，但发现这个函数没有被暴露出来。

我尝试手动暴露这个函数（下面是一个最小的示例，展示了失败的情况）

from libc.stdlib cimport malloc
import numpy as np
cimport numpy as np

np.import_array()

ctypedef np.int32_t DTYPE_t

cdef extern from "numpy/ndarraytypes.h":
    void PyArray_ENABLEFLAGS(np.PyArrayObject *arr, int flags)

def test():
    cdef int N = 1000

    cdef DTYPE_t *data = <DTYPE_t *>malloc(N * sizeof(DTYPE_t))
    cdef np.ndarray[DTYPE_t, ndim=1] arr = np.PyArray_SimpleNewFromData(1, &N, np.NPY_INT32, data)
    PyArray_ENABLEFLAGS(arr, np.NPY_ARRAY_OWNDATA)

结果出现了编译错误：

Error compiling Cython file:
------------------------------------------------------------
...
def test():
    cdef int N = 1000

    cdef DTYPE_t *data = <DTYPE_t *>malloc(N * sizeof(DTYPE_t))
    cdef np.ndarray[DTYPE_t, ndim=1] arr = np.PyArray_SimpleNewFromData(1, &N, np.NPY_INT32, data)
    PyArray_ENABLEFLAGS(arr, np.NPY_ARRAY_OWNDATA)
                          ^
------------------------------------------------------------

/tmp/test.pyx:19:27: Cannot convert Python object to 'PyArrayObject *'

我的问题：在这种情况下，我的做法对吗？如果对，那我哪里做错了？如果不对，我该如何在Cython中强制NumPy拥有内存，而不需要使用C扩展模块呢？

Answer 1

最新的Cython版本让你可以用很简单的语法来处理一些事情，虽然相比于更底层的解决方案，它的开销稍微大一点。

numpy_array = np.asarray(<np.int32_t[:10, :10]> my_pointer)

https://cython.readthedocs.io/en/latest/src/userguide/memoryviews.html#coercion-to-numpy

这一步并没有转移所有权。

值得注意的是，通过这个调用会生成一个Cython数组，具体是通过array_cwrapper来实现的。

这个过程生成了一个cython.array，而且没有分配新的内存。这个cython.array默认使用的是stdlib.h中的malloc和free，所以你应该使用默认的malloc，而不是任何特别的CPython/Numpy分配器。

只有在这个cython.array设置了所有权时，free才会被调用，默认情况下只有在它分配了数据时才会这样。对于我们的情况，我们可以手动设置它，方法是：

my_cyarr.free_data = True

---

所以如果要返回一个一维数组，操作就简单得多：

from cython.view cimport array as cvarray

# ...
    cdef cvarray cvarr = <np.int32_t[:N]> data
    cvarr.free_data = True
    return np.asarray(cvarr)

Answer 2

@Stefan的解决方案在大多数情况下都能用，但有点脆弱。Numpy使用PyDataMem_NEW/PyDataMem_FREE来管理内存，这些调用实际上是和常见的malloc/free结合使用的，还有一些内存追踪的功能（我不太清楚Stefan的解决方案对内存追踪有什么影响，至少看起来没有崩溃）。

还有一些比较复杂的情况，比如在numpy库中调用的free和在cython代码中调用的malloc可能不是用同一个内存分配器（例如链接了不同的运行时，具体可以参考这个github问题或这个SO帖子）。

正确的工具来传递或管理数据的所有权是PyArray_SetBaseObject。

首先，我们需要一个python对象，负责释放内存。我这里使用了一个自定义的cdef类（主要是为了记录和演示），但显然还有其他选择：

%%cython
from libc.stdlib cimport free

cdef class MemoryNanny:
    cdef void* ptr # set to NULL by "constructor"
    def __dealloc__(self):
        print("freeing ptr=", <unsigned long long>(self.ptr)) #just for debugging
        free(self.ptr)
        
    @staticmethod
    cdef create(void* ptr):
        cdef MemoryNanny result = MemoryNanny()
        result.ptr = ptr
        print("nanny for ptr=", <unsigned long long>(result.ptr)) #just for debugging
        return result

 ...

现在，我们使用一个MemoryNanny对象作为内存的守护者，一旦父numpy数组被销毁，这个内存就会被释放。代码有点复杂，因为PyArray_SetBaseObject会“偷走”引用，而这在Cython中并不会自动处理：

%%cython
...
from cpython.object cimport PyObject
from cpython.ref cimport Py_INCREF

cimport numpy as np

#needed to initialize PyArray_API in order to be able to use it
np.import_array()


cdef extern from "numpy/arrayobject.h":
    # a little bit awkward: the reference to obj will be stolen
    # using PyObject*  to signal that Cython cannot handle it automatically
    int PyArray_SetBaseObject(np.ndarray arr, PyObject *obj) except -1 # -1 means there was an error
          
cdef array_from_ptr(void * ptr, np.npy_intp N, int np_type):
    cdef np.ndarray arr = np.PyArray_SimpleNewFromData(1, &N, np_type, ptr)
    nanny = MemoryNanny.create(ptr)
    Py_INCREF(nanny) # a reference will get stolen, so prepare nanny
    PyArray_SetBaseObject(arr, <PyObject*>nanny) 
    return arr
...

下面是一个示例，展示如何调用这个功能：

%%cython
...
from libc.stdlib cimport malloc
def create():
    cdef double *ptr=<double*>malloc(sizeof(double)*8);
    ptr[0]=42.0
    return array_from_ptr(ptr, 8, np.NPY_FLOAT64)

可以按如下方式使用：

>>> m =  create()
nanny for ptr= 94339864945184
>>> m.flags
...
OWNDATA : False
...
>>> m[0]
42.0
>>> del m
freeing ptr= 94339864945184

结果/输出如预期。

注意：生成的数组并不真正拥有数据（也就是说，标志返回OWNDATA : False），因为内存是由内存守护者拥有的，但结果是一样的：一旦数组被删除，内存就会被释放（因为没有人再持有对守护者的引用）。

---

MemoryNanny不一定要保护一个原始的C指针。它可以是其他任何东西，比如一个std::vector：

%%cython -+
from libcpp.vector cimport vector
cdef class VectorNanny:
    #automatically default initialized/destructed by Cython:
    cdef vector[double] vec 
    @staticmethod
    cdef create(vector[double]& vec):
        cdef VectorNanny result = VectorNanny()
        result.vec.swap(vec) # swap and not copy
        return result
   
# for testing:
def create_vector(int N):
    cdef vector[double] vec;
    vec.resize(N, 2.0)
    return VectorNanny.create(vec)

以下测试显示，守护者是有效的：

nanny=create_vector(10**8) # top shows additional 800MB memory are used
del nanny                  # top shows, this additional memory is no longer used.

Answer 3

你在接口定义上有一些小错误。下面的代码对我来说是有效的：

from libc.stdlib cimport malloc
import numpy as np
cimport numpy as np

np.import_array()

ctypedef np.int32_t DTYPE_t

cdef extern from "numpy/arrayobject.h":
    void PyArray_ENABLEFLAGS(np.ndarray arr, int flags)

cdef data_to_numpy_array_with_spec(void * ptr, np.npy_intp N, int t):
    cdef np.ndarray[DTYPE_t, ndim=1] arr = np.PyArray_SimpleNewFromData(1, &N, t, ptr)
    PyArray_ENABLEFLAGS(arr, np.NPY_OWNDATA)
    return arr

def test():
    N = 1000

    cdef DTYPE_t *data = <DTYPE_t *>malloc(N * sizeof(DTYPE_t))
    arr = data_to_numpy_array_with_spec(data, N, np.NPY_INT32)
    return arr

这是我的 setup.py 文件：

from distutils.core import setup, Extension
from Cython.Distutils import build_ext
ext_modules = [Extension("_owndata", ["owndata.pyx"])]
setup(cmdclass={'build_ext': build_ext}, ext_modules=ext_modules)

用 python setup.py build_ext --inplace 来构建。然后检查一下数据是否真的被拥有：

import _owndata
arr = _owndata.test()
print arr.flags

你应该能看到 OWNDATA : True。

而且 没错，这确实是处理这个问题的正确方法，因为 numpy.pxd 也是用同样的方式把其他函数导出到 Cython。