Merge pull request #16523 from seberg/buffer-speed

MAINT: Improve buffer speed

2 files changed, 44 insertions, 52 deletions

diff --git a/numpy/core/src/multiarray/buffer.c b/numpy/core/src/multiarray/buffer.c
index 232176011..8b482dc03 100644
--- a/numpy/core/src/multiarray/buffer.c
+++ b/numpy/core/src/multiarray/buffer.c
@@ -64,7 +64,7 @@ _append_char(_tmp_string_t *s, char c)
         char *p;
         size_t to_alloc = (s->allocated == 0) ? INIT_SIZE : (2 * s->allocated);
 
-        p = realloc(s->s, to_alloc);
+        p = PyObject_Realloc(s->s, to_alloc);
         if (p == NULL) {
             PyErr_SetString(PyExc_MemoryError, "memory allocation failed");
             return -1;
@@ -135,12 +135,25 @@ fail:
  * AND, the descr element size is a multiple of the alignment,
  * AND, the array data is positioned to alignment granularity.
  */
-static int
+static NPY_INLINE int
 _is_natively_aligned_at(PyArray_Descr *descr,
                         PyArrayObject *arr, Py_ssize_t offset)
 {
     int k;
 
+    if (NPY_LIKELY(descr == PyArray_DESCR(arr))) {
+        /*
+         * If the descriptor is the arrays descriptor we can assume the
+         * array's alignment is correct.
+         */
+        assert(offset == 0);
+        if (PyArray_ISALIGNED(arr)) {
+            assert(descr->elsize % descr->alignment == 0);
+            return 1;
+        }
+        return 0;
+    }
+
     if ((Py_ssize_t)(PyArray_DATA(arr)) % descr->alignment != 0) {
         return 0;
     }
@@ -297,8 +310,6 @@ _buffer_format_string(PyArray_Descr *descr, _tmp_string_t *str,
                 descr->type_num == NPY_ULONGLONG);
         }
 
-        *offset += descr->elsize;
-
         if (PyArray_IsScalar(obj, Generic)) {
             /* scalars are always natively aligned */
             is_natively_aligned = 1;
@@ -308,6 +319,8 @@ _buffer_format_string(PyArray_Descr *descr, _tmp_string_t *str,
                                               (PyArrayObject*)obj, *offset);
         }
 
+        *offset += descr->elsize;
+
         if (descr->byteorder == '=' && is_natively_aligned) {
             /* Prefer native types, to cater for Cython */
             is_standard_size = 0;
@@ -445,49 +458,22 @@ static PyObject *_buffer_info_cache = NULL;
 static _buffer_info_t*
 _buffer_info_new(PyObject *obj)
 {
+    /*
+     * Note that the buffer info is cached as PyLongObjects making them appear
+     * like unreachable lost memory to valgrind.
+     */
     _buffer_info_t *info;
     _tmp_string_t fmt = {NULL, 0, 0};
     int k;
     PyArray_Descr *descr = NULL;
     int err = 0;
 
-    /*
-     * Note that the buffer info is cached as pyints making them appear like
-     * unreachable lost memory to valgrind.
-     */
-    info = malloc(sizeof(_buffer_info_t));
-    if (info == NULL) {
-        PyErr_NoMemory();
-        goto fail;
-    }
-
-    if (PyArray_IsScalar(obj, Datetime) || PyArray_IsScalar(obj, Timedelta)) {
-        /*
-         * Special case datetime64 scalars to remain backward compatible.
-         * This will change in a future version.
-         * Note arrays of datetime64 and structured arrays with datetime64
-         * fields will not hit this code path and are currently unsupported
-         * in _buffer_format_string.
-         */
-        if (_append_char(&fmt, 'B') < 0) {
-            goto fail;
-        }
-        if (_append_char(&fmt, '\0') < 0) {
-            goto fail;
-        }
-        info->ndim = 1;
-        info->shape = malloc(sizeof(Py_ssize_t) * 2);
-        if (info->shape == NULL) {
+    if (PyArray_IsScalar(obj, Void)) {
+        info = PyObject_Malloc(sizeof(_buffer_info_t));
+        if (info == NULL) {
             PyErr_NoMemory();
             goto fail;
         }
-        info->strides = info->shape + info->ndim;
-        info->shape[0] = 8;
-        info->strides[0] = 1;
-        info->format = fmt.s;
-        return info;
-    }
-    else if (PyArray_IsScalar(obj, Generic)) {
         descr = PyArray_DescrFromScalar(obj);
         if (descr == NULL) {
             goto fail;
@@ -497,8 +483,16 @@ _buffer_info_new(PyObject *obj)
         info->strides = NULL;
     }
     else {
+        assert(PyArray_Check(obj));
         PyArrayObject * arr = (PyArrayObject *)obj;
         descr = PyArray_DESCR(arr);
+
+        info = PyObject_Malloc(sizeof(_buffer_info_t) +
+                               sizeof(Py_ssize_t) * PyArray_NDIM(arr) * 2);
+        if (info == NULL) {
+            PyErr_NoMemory();
+            goto fail;
+        }
         /* Fill in shape and strides */
         info->ndim = PyArray_NDIM(arr);
 
@@ -507,11 +501,8 @@ _buffer_info_new(PyObject *obj)
             info->strides = NULL;
         }
         else {
-            info->shape = malloc(sizeof(Py_ssize_t) * PyArray_NDIM(arr) * 2 + 1);
-            if (info->shape == NULL) {
-                PyErr_NoMemory();
-                goto fail;
-            }
+            info->shape = (npy_intp *)((char *)info + sizeof(_buffer_info_t));
+            assert((size_t)info->shape % sizeof(npy_intp) == 0);
             info->strides = info->shape + PyArray_NDIM(arr);
             for (k = 0; k < PyArray_NDIM(arr); ++k) {
                 info->shape[k] = PyArray_DIMS(arr)[k];
@@ -525,11 +516,9 @@ _buffer_info_new(PyObject *obj)
     err = _buffer_format_string(descr, &fmt, obj, NULL, NULL);
     Py_DECREF(descr);
     if (err != 0) {
-        free(info->shape);
         goto fail;
     }
     if (_append_char(&fmt, '\0') < 0) {
-        free(info->shape);
         goto fail;
     }
     info->format = fmt.s;
@@ -537,8 +526,8 @@ _buffer_info_new(PyObject *obj)
     return info;
 
 fail:
-    free(fmt.s);
-    free(info);
+    PyObject_Free(fmt.s);
+    PyObject_Free(info);
     return NULL;
 }
 
@@ -569,12 +558,9 @@ static void
 _buffer_info_free(_buffer_info_t *info)
 {
     if (info->format) {
-        free(info->format);
-    }
-    if (info->shape) {
-        free(info->shape);
+        PyObject_Free(info->format);
     }
-    free(info);
+    PyObject_Free(info);
 }
 
 /* Get buffer info from the global dictionary */
diff --git a/numpy/core/tests/test_scalarbuffer.py b/numpy/core/tests/test_scalarbuffer.py
index b1c1bbbb1..574c56864 100644
--- a/numpy/core/tests/test_scalarbuffer.py
+++ b/numpy/core/tests/test_scalarbuffer.py
@@ -2,6 +2,7 @@
 Test scalar buffer interface adheres to PEP 3118
 """
 import numpy as np
+from numpy.core._rational_tests import rational
 import pytest
 
 from numpy.testing import assert_, assert_equal, assert_raises
@@ -117,3 +118,8 @@ class TestScalarPEP3118:
         code_points = np.frombuffer(v, dtype='i4')
 
         assert_equal(code_points, [ord(c) for c in s])
+
+    def test_user_scalar_fails_buffer(self):
+        r = rational(1)
+        with assert_raises(TypeError):
+            memoryview(r)