align realloc behavior to NumPy

also address issues with undeclared variables and rename MKLMemory class members

Author: ndgrigorian

mkl/_mkl_memory.pyx

@@ -30,6 +30,7 @@
 import numbers
 
 from cpython cimport Py_buffer
+from libc.limits cimport INT_MAX
 from libc.string cimport memcpy
 
 from mkl._mkl_service cimport mkl_calloc, mkl_free, mkl_malloc, mkl_realloc
@@ -43,6 +44,37 @@ cdef extern from "stdatomic.h" nogil:
     int atomic_load(atomic_int *obj)
 
 
+cdef extern from *:
+    """
+    // Check whether a MKLMemory object may be safely reallocated.
+    // Mirrors NumPy's PyArray_Resize_int logic.
+    static int _MKLMemory_MayBeShared(PyObject *op) {
+    #if PY_VERSION_HEX >= 0x030e00b0
+        if (PyUnstable_Object_IsUniquelyReferenced(op)) {
+            return 0; // not shared
+        }
+        if (Py_REFCNT(op) == 2) {
+            return 1;  // may be shared
+        }
+        return 2;  // definitely shared
+    #else
+        return (Py_REFCNT(op) > 2) ? 2 : 0;
+    #endif
+    }
+    """
+    int _MKLMemory_MayBeShared(object obj)
+
+
+cdef int _check_alignment(Py_ssize_t alignment) except -1:
+    if alignment <= 0:
+        raise ValueError("Alignment of requested allocation must be positive.")
+    if alignment > <Py_ssize_t>INT_MAX:
+        raise ValueError(
+            f"Alignment of requested allocation must not exceed {INT_MAX}."
+        )
+    return <int>alignment
+
+
 def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment):
     cdef Py_ssize_t nbytes = len(data)
     cdef MKLMemory mem = MKLMemory(nbytes, alignment=alignment)
@@ -57,28 +89,32 @@ def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment):
 
 
 cdef class MKLMemory:
+    """MKL-backed memory object that exposes Python buffer protocol."""
     cdef void *_memory_ptr
-    cdef Py_ssize_t nbytes
-    cdef Py_ssize_t alignment
+    cdef Py_ssize_t _nbytes
+    cdef Py_ssize_t _alignment
     cdef atomic_int exported_buffers
 
     cdef _cinit_empty(self):
         self._memory_ptr = NULL
-        self.nbytes = 0
-        self.alignment = 0
+        self._nbytes = 0
+        self._alignment = 0
         atomic_init(&self.exported_buffers, 0)
 
     cdef _cinit_malloc(self, Py_ssize_t nbytes, Py_ssize_t alignment):
+        cdef int c_alignment = _check_alignment(alignment)
+        cdef void *p
+
         self._cinit_empty()
 
         if (nbytes > 0):
             with nogil:
-                p = mkl_malloc(nbytes, alignment)
+                p = mkl_malloc(nbytes, c_alignment)
 
             if (p):
                 self._memory_ptr = p
-                self.nbytes = nbytes
-                self.alignment = alignment
+                self._nbytes = nbytes
+                self._alignment = alignment
             else:
                 raise MemoryError(
                     "MKL memory allocation failed."
@@ -89,16 +125,19 @@ cdef class MKLMemory:
             )
 
     cdef _cinit_calloc(self, Py_ssize_t num, Py_ssize_t size, Py_ssize_t alignment):
+        cdef int c_alignment = _check_alignment(alignment)
+        cdef void *p
+
         self._cinit_empty()
 
         if (num > 0 and size > 0):
             with nogil:
-                p = mkl_calloc(num, size, alignment)
+                p = mkl_calloc(num, size, c_alignment)
 
             if (p):
                 self._memory_ptr = p
-                self.nbytes = num * size
-                self.alignment = alignment
+                self._nbytes = num * size
+                self._alignment = alignment
             else:
                 raise MemoryError(
                     "MKL memory allocation failed."
@@ -110,11 +149,11 @@ cdef class MKLMemory:
             )
 
     cdef _cinit_mklmemory(self, object other, Py_ssize_t alignment):
-        other_mem = <MKLMemory> other
+        cdef MKLMemory other_mem = <MKLMemory> other
 
-        self._cinit_malloc(other_mem.nbytes, alignment)
+        self._cinit_malloc(other_mem._nbytes, alignment)
         with nogil:
-            memcpy(self._memory_ptr, other_mem._memory_ptr, self.nbytes)
+            memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes)
 
     def __cinit__(self, *args, **kwargs):
         cdef Py_ssize_t alignment
@@ -131,7 +170,7 @@ cdef class MKLMemory:
                 alignment = kwargs.get("alignment", 64)
                 self._cinit_malloc(arg, alignment)
             elif isinstance(arg, MKLMemory):
-                alignment = kwargs.get("alignment", arg.alignment)
+                alignment = kwargs.get("alignment", (<MKLMemory>arg)._alignment)
                 self._cinit_mklmemory(arg, alignment)
             else:
                 raise TypeError(
@@ -167,11 +206,11 @@ cdef class MKLMemory:
         buffer.format = "B"
         buffer.internal = NULL
         buffer.itemsize = 1
-        buffer.len = self.nbytes
+        buffer.len = self._nbytes
         buffer.ndim = 1
         buffer.obj = self
         buffer.readonly = 0
-        buffer.shape = &self.nbytes
+        buffer.shape = &self._nbytes
         buffer.strides = &buffer.itemsize
         buffer.suboffsets = NULL
 
@@ -181,52 +220,66 @@ cdef class MKLMemory:
         atomic_fetch_sub(&self.exported_buffers, 1)
 
     def realloc(self, Py_ssize_t new_nbytes):
+        cdef void *p
+        cdef int shared
+
         if atomic_load(&self.exported_buffers) > 0:
-            raise BufferError("Cannot realloc memory while there are exported buffers.")
+            raise BufferError(
+                "Cannot realloc memory while there are exported buffers."
+            )
+        shared = _MKLMemory_MayBeShared(self)
+        if shared == 1:
+            raise ValueError(
+                "Cannot realloc MKLMemory that may be referenced by another "
+                "object. It is possible that this is a false positive."
+            )
+        elif shared == 2:
+            raise ValueError(
+                "Cannot realloc MKLMemory that is referenced by other objects."
+            )
         if new_nbytes <= 0:
             raise ValueError("New number of bytes must be positive.")
 
-        cdef void *p
         with nogil:
             p = mkl_realloc(self._memory_ptr, new_nbytes)
 
         if not p:
             raise MemoryError("MKL memory reallocation failed.")
 
         self._memory_ptr = p
-        self.nbytes = new_nbytes
+        self._nbytes = new_nbytes
 
     def tobytes(self):
         cdef char* data_ptr = <char*>self._memory_ptr
-        return data_ptr[:self.nbytes]
+        return data_ptr[:self._nbytes]
 
     @property
     def nbytes(self):
-        return self.nbytes
+        return self._nbytes
 
     @property
     def size(self):
-        return self.nbytes
+        return self._nbytes
 
     @property
     def alignment(self):
-        return self.alignment
+        return self._alignment
 
     @property
     def _pointer(self):
         return <size_t>(self._memory_ptr)
 
     def __repr__(self):
         return (
-            f"<MKL memory allocation of {self.nbytes} bytes at "
+            f"<MKL memory allocation of {self._nbytes} bytes at "
             f"{hex(<object>(<size_t>self._memory_ptr))}>"
         )
 
     def __len__(self):
-        return self.nbytes
+        return self._nbytes
 
     def __sizeof__(self):
-        return self.nbytes
+        return self._nbytes
 
     def __reduce__(self):
-        return (_mkl_memory_from_bytes, (self.tobytes(), self.alignment))
+        return (_mkl_memory_from_bytes, (self.tobytes(), self._alignment))

mkl/tests/test_mkl_memory.py

@@ -24,6 +24,9 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import sys
+import threading
+
+import pytest
 
 import mkl
 
@@ -137,3 +140,84 @@ def test_pickling_with_alignment():
     assert (
         mem.alignment == mem_reconstructed.alignment
     ), "Pickling should preserve alignment"
+
+
+def test_realloc_exported_buffer():
+    mem = mkl.MKLMemory(1024)
+    mv = memoryview(mem)
+    with pytest.raises(BufferError):
+        mem.realloc(2048)
+    del mv
+
+
+def test_realloc_refcheck_shared():
+    mem = mkl.MKLMemory(1024)
+    alias = mem  # noqa: F841 — extra reference
+    with pytest.raises(ValueError, match="referenced by"):
+        mem.realloc(2048)
+    del alias
+
+
+def test_alignment_validation():
+    with pytest.raises(ValueError, match="positive"):
+        mkl.MKLMemory(1024, alignment=0)
+    with pytest.raises(ValueError, match="positive"):
+        mkl.MKLMemory(1024, alignment=-1)
+    with pytest.raises(ValueError, match="must not exceed"):
+        mkl.MKLMemory(1024, alignment=2**40)
+
+
+def test_concurrent_reads():
+    mem = mkl.MKLMemory(1024)
+    mv = memoryview(mem)
+    for i in range(len(mem)):
+        mv[i] = i % 256
+    del mv
+
+    errors = []
+
+    def reader():
+        try:
+            for _ in range(500):
+                assert len(mem) == 1024
+                data = mem.tobytes()
+                assert len(data) == 1024
+                v = memoryview(mem)
+                assert v[0] == 0
+                v.release()
+        except Exception as e:
+            errors.append(e)
+
+    ts = [threading.Thread(target=reader) for _ in range(4)]
+    for t in ts:
+        t.start()
+    for t in ts:
+        t.join()
+    assert not errors, f"Concurrent read errors: {errors}"
+
+
+def test_concurrent_realloc_refused():
+    for _ in range(50):
+        mem = mkl.MKLMemory(64)
+        barrier = threading.Barrier(2)
+        results = [None, None]
+
+        def worker(idx, size):
+            barrier.wait()
+            try:
+                mem.realloc(size)
+                results[idx] = "ok"
+            except ValueError:
+                results[idx] = "refused"
+
+        ts = [
+            threading.Thread(target=worker, args=(0, 1 << 16)),
+            threading.Thread(target=worker, args=(1, 1 << 17)),
+        ]
+        for t in ts:
+            t.start()
+        for t in ts:
+            t.join()
+        assert results[0] == "refused" and results[1] == "refused", (
+            f"Expected both threads refused, got {results}"
+        )