From 726676a21ac4661e81ecdb4a236324f5740500ae Mon Sep 17 00:00:00 2001 From: Nathan Goldbaum Date: Fri, 14 Apr 2023 13:20:47 -0600 Subject: [PATCH] ENH: Refactor special zero-filling to be managed by the DType --- numpy/core/include/numpy/_dtype_api.h | 11 +++++- numpy/core/src/multiarray/common.c | 18 ---------- numpy/core/src/multiarray/common.h | 3 -- numpy/core/src/multiarray/ctors.c | 48 ++++++++++++++++++--------- numpy/core/src/multiarray/dtypemeta.c | 25 ++++++++++++++ numpy/core/src/multiarray/dtypemeta.h | 12 ++++++- numpy/core/src/multiarray/getset.c | 19 +++++------ 7 files changed, 87 insertions(+), 49 deletions(-) diff --git a/numpy/core/include/numpy/_dtype_api.h b/numpy/core/include/numpy/_dtype_api.h index 9a2e508903fa..36a8524a5718 100644 --- a/numpy/core/include/numpy/_dtype_api.h +++ b/numpy/core/include/numpy/_dtype_api.h @@ -5,7 +5,7 @@ #ifndef NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ #define NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ -#define __EXPERIMENTAL_DTYPE_API_VERSION 9 +#define __EXPERIMENTAL_DTYPE_API_VERSION 10 struct PyArrayMethodObject_tag; @@ -199,6 +199,14 @@ typedef int (get_reduction_initial_function)( PyArrayMethod_Context *context, npy_bool reduction_is_empty, char *initial); +/** + * A function to fill an array with an initial value. + * + * @param arr The array to be filled. + * + * @returns -1 or 0 indicating error and arr being successfully filled. + */ +typedef int (fill_initial_function)(PyArrayObject *arr); /* * The following functions are only used be the wrapping array method defined @@ -319,6 +327,7 @@ typedef int (get_traverse_loop_function)( #define NPY_DT_setitem 7 #define NPY_DT_getitem 8 #define NPY_DT_get_clear_loop 9 +#define NPY_DT_fill_zero_value 10 // These PyArray_ArrFunc slots will be deprecated and replaced eventually // getitem and setitem can be defined as a performance optimization; diff --git a/numpy/core/src/multiarray/common.c b/numpy/core/src/multiarray/common.c index da8d23a262aa..001d299c7f11 100644 --- a/numpy/core/src/multiarray/common.c +++ b/numpy/core/src/multiarray/common.c @@ -128,24 +128,6 @@ PyArray_DTypeFromObject(PyObject *obj, int maxdims, PyArray_Descr **out_dtype) } -NPY_NO_EXPORT int -_zerofill(PyArrayObject *ret) -{ - if (PyDataType_REFCHK(PyArray_DESCR(ret))) { - PyObject *zero = PyLong_FromLong(0); - PyArray_FillObjectArray(ret, zero); - Py_DECREF(zero); - if (PyErr_Occurred()) { - return -1; - } - } - else { - npy_intp n = PyArray_NBYTES(ret); - memset(PyArray_DATA(ret), 0, n); - } - return 0; -} - NPY_NO_EXPORT npy_bool _IsWriteable(PyArrayObject *ap) { diff --git a/numpy/core/src/multiarray/common.h b/numpy/core/src/multiarray/common.h index 4e067b22c1d5..127c6250da80 100644 --- a/numpy/core/src/multiarray/common.h +++ b/numpy/core/src/multiarray/common.h @@ -55,9 +55,6 @@ PyArray_DTypeFromObject(PyObject *obj, int maxdims, NPY_NO_EXPORT PyArray_Descr * _array_find_python_scalar_type(PyObject *op); -NPY_NO_EXPORT int -_zerofill(PyArrayObject *ret); - NPY_NO_EXPORT npy_bool _IsWriteable(PyArrayObject *ap); diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c index a6335c78372b..49296bd4d45f 100644 --- a/numpy/core/src/multiarray/ctors.c +++ b/numpy/core/src/multiarray/ctors.c @@ -783,6 +783,10 @@ PyArray_NewFromDescr_int( } + /* dtype-specific function to fill array with zero values, may be NULL */ + fill_initial_function *fill_zero_value = + NPY_DT_SLOTS(NPY_DTYPE(descr))->fill_zero_value; + if (data == NULL) { /* Store the handler in case the default is modified */ fa->mem_handler = PyDataMem_GetHandler(); @@ -801,14 +805,29 @@ PyArray_NewFromDescr_int( fa->strides[i] = 0; } } - /* - * It is bad to have uninitialized OBJECT pointers - * which could also be sub-fields of a VOID array - */ - if (zeroed || PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) { + + if ( + /* + * If NPY_NEEDS_INIT is set, always zero out the buffer, even if + * zeroed isn't set. + */ + (PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) || + /* + * If fill_zero_value isn't set, the dtype definitely has no + * special zero value, so get a zero-filled buffer using + * calloc. Otherwise, get the buffer with malloc and allow + * fill_zero_value to zero out the array with the appropriate + * special zero value for the dtype. VOID arrays have + * fill_zero_value set but may or may not contain objects, so + * always allocate with calloc if zeroed is set. + */ + (zeroed && + ((fill_zero_value == NULL) || (descr->type_num == NPY_VOID)))) { + /* allocate array buffer with calloc */ data = PyDataMem_UserNEW_ZEROED(nbytes, 1, fa->mem_handler); } else { + /* allocate array buffer with malloc */ data = PyDataMem_UserNEW(nbytes, fa->mem_handler); } if (data == NULL) { @@ -845,6 +864,15 @@ PyArray_NewFromDescr_int( } } + /* + * If the array needs special dtype-specific zero-filling logic, do that + */ + if (zeroed && (fill_zero_value != NULL)) { + if (fill_zero_value((PyArrayObject *)fa) < 0) { + goto fail; + } + } + /* * call the __array_finalize__ method if a subtype was requested. * If obj is NULL use Py_None for the Python callback. @@ -3017,17 +3045,7 @@ PyArray_Zeros(int nd, npy_intp const *dims, PyArray_Descr *type, int is_f_order) return NULL; } - /* handle objects */ - if (PyDataType_REFCHK(PyArray_DESCR(ret))) { - if (_zerofill(ret) < 0) { - Py_DECREF(ret); - return NULL; - } - } - - return (PyObject *)ret; - } /*NUMPY_API diff --git a/numpy/core/src/multiarray/dtypemeta.c b/numpy/core/src/multiarray/dtypemeta.c index f268ba2cb2dd..b896be3cb3d0 100644 --- a/numpy/core/src/multiarray/dtypemeta.c +++ b/numpy/core/src/multiarray/dtypemeta.c @@ -698,6 +698,28 @@ object_common_dtype( return cls; } +int +object_fill_zero_value(PyArrayObject *arr) +{ + PyObject *zero = PyLong_FromLong(0); + PyArray_FillObjectArray(arr, zero); + Py_DECREF(zero); + if (PyErr_Occurred()) { + return -1; + } + return 0; +} + +int +void_fill_zero_value(PyArrayObject *arr) +{ + if (PyDataType_REFCHK(PyArray_DESCR(arr))) { + if (object_fill_zero_value(arr) < 0) { + return -1; + } + } + return 0; +} /** * This function takes a PyArray_Descr and replaces its base class with @@ -855,6 +877,7 @@ dtypemeta_wrap_legacy_descriptor(PyArray_Descr *descr) dt_slots->common_dtype = default_builtin_common_dtype; dt_slots->common_instance = NULL; dt_slots->ensure_canonical = ensure_native_byteorder; + dt_slots->fill_zero_value = NULL; if (PyTypeNum_ISSIGNED(dtype_class->type_num)) { /* Convert our scalars (raise on too large unsigned and NaN, etc.) */ @@ -866,6 +889,7 @@ dtypemeta_wrap_legacy_descriptor(PyArray_Descr *descr) } else if (descr->type_num == NPY_OBJECT) { dt_slots->common_dtype = object_common_dtype; + dt_slots->fill_zero_value = object_fill_zero_value; } else if (PyTypeNum_ISDATETIME(descr->type_num)) { /* Datetimes are flexible, but were not considered previously */ @@ -887,6 +911,7 @@ dtypemeta_wrap_legacy_descriptor(PyArray_Descr *descr) void_discover_descr_from_pyobject); dt_slots->common_instance = void_common_instance; dt_slots->ensure_canonical = void_ensure_canonical; + dt_slots->fill_zero_value = void_fill_zero_value; } else { dt_slots->default_descr = string_and_unicode_default_descr; diff --git a/numpy/core/src/multiarray/dtypemeta.h b/numpy/core/src/multiarray/dtypemeta.h index 3b4dbad24d21..897d2a13384d 100644 --- a/numpy/core/src/multiarray/dtypemeta.h +++ b/numpy/core/src/multiarray/dtypemeta.h @@ -41,6 +41,16 @@ typedef struct { * Python objects. */ get_traverse_loop_function *get_clear_loop; + /* + Either NULL or a function that fills an array with zero values + appropriate for the dtype. If this function pointer is NULL, the initial + value is assumed to be zero. If this function is defined, the array + buffer is allocated with malloc and then this function is called to fill + the buffer. As a performance optimization, the array buffer is allocated + using calloc if this function is NULL, so it is best to avoid using this + function if zero-filling the array buffer makes sense for the dtype. + */ + fill_initial_function *fill_zero_value; /* * The casting implementation (ArrayMethod) to convert between two * instances of this DType, stored explicitly for fast access: @@ -63,7 +73,7 @@ typedef struct { // This must be updated if new slots before within_dtype_castingimpl // are added -#define NPY_NUM_DTYPE_SLOTS 9 +#define NPY_NUM_DTYPE_SLOTS 10 #define NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS 22 #define NPY_DT_MAX_ARRFUNCS_SLOT \ NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c index ab35548ed728..d019acbb5b23 100644 --- a/numpy/core/src/multiarray/getset.c +++ b/numpy/core/src/multiarray/getset.c @@ -797,20 +797,17 @@ array_imag_get(PyArrayObject *self, void *NPY_UNUSED(ignored)) } else { Py_INCREF(PyArray_DESCR(self)); - ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self), - PyArray_DESCR(self), - PyArray_NDIM(self), - PyArray_DIMS(self), - NULL, NULL, - PyArray_ISFORTRAN(self), - (PyObject *)self); + ret = (PyArrayObject *)PyArray_NewFromDescr_int( + Py_TYPE(self), + PyArray_DESCR(self), + PyArray_NDIM(self), + PyArray_DIMS(self), + NULL, NULL, + PyArray_ISFORTRAN(self), + (PyObject *)self, NULL, 1, 0); if (ret == NULL) { return NULL; } - if (_zerofill(ret) < 0) { - Py_DECREF(ret); - return NULL; - } PyArray_CLEARFLAGS(ret, NPY_ARRAY_WRITEABLE); } return (PyObject *) ret;