From 7bd5c95b7cb741fb2158e11c0379484ad98a96d8 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 19:56:57 -0700 Subject: [PATCH 01/16] add MKLMemory object, backed with mkl_malloc memory exposes Python buffer protocol --- meson.build | 13 +++++- mkl/__init__.py | 1 + mkl/_mkl_memory.pyx | 104 +++++++++++++++++++++++++++++++++++++++++++ mkl/_mkl_service.pxd | 2 + 4 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 mkl/_mkl_memory.pyx diff --git a/meson.build b/meson.build index 0972a6b..32ae0f2 100644 --- a/meson.build +++ b/meson.build @@ -60,7 +60,7 @@ py.extension_module( subdir: 'mkl' ) -# Cython extension +# Cython extensions py.extension_module( '_py_mkl_service', sources: ['mkl/_py_mkl_service.pyx'], @@ -71,6 +71,17 @@ py.extension_module( subdir: 'mkl' ) +py.extension_module( + '_mkl_memory', + sources: ['mkl/_mkl_memory.pyx'], + dependencies: [mkl_dep], + c_args: c_args, + install_rpath: rpath, + install: true, + subdir: 'mkl' +) + + # Python sources py.install_sources( [ diff --git a/mkl/__init__.py b/mkl/__init__.py index c0eb2ae..beadbfc 100644 --- a/mkl/__init__.py +++ b/mkl/__init__.py @@ -57,6 +57,7 @@ def __exit__(self, *args): del RTLD_for_MKL +from ._mkl_memory import MKLMemory from ._py_mkl_service import ( cbwr_get, cbwr_get_auto_branch, diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx new file mode 100644 index 0000000..bc4a87d --- /dev/null +++ b/mkl/_mkl_memory.pyx @@ -0,0 +1,104 @@ +# Copyright (c) 2018, Intel Corporation +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of Intel Corporation nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# distutils: language = c +# cython: language_level=3 + +import numbers + +from cpython cimport Py_buffer +from libc.string cimport memcpy + +from mkl._mkl_service cimport mkl_malloc, mkl_free + + +cdef class MKLMemory: + cdef void *_memory_ptr + cdef Py_ssize_t nbytes + + cdef _cinit_empty(self): + self._memory_ptr = NULL + self.nbytes = 0 + + cdef _cinit_alloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): + self._cinit_empty() + + if (nbytes > 0): + with nogil: + p = mkl_malloc(nbytes, alignment) + + if (p): + self._memory_ptr = p + self.nbytes = nbytes + else: + raise MemoryError( + "MKL memory allocation failed." + ) + else: + raise ValueError( + "Number of bytes of request allocation must be positive." + ) + + cdef _cinit_other(self, object other, Py_ssize_t alignment): + cdef MKLMemory other_mem + if isinstance(other, MKLMemory): + other_mem = other + else: + raise ValueError( + f"Argument {other} is not of type MKLMemory." + ) + self._cinit_alloc(other_mem.nbytes, alignment) + with nogil: + memcpy(self._memory_ptr, other_mem._memory_ptr, self.nbytes) + + def __cinit__(self, other, *, Py_ssize_t alignment=64): + if isinstance(other, numbers.Integral): + self._cinit_alloc(other, alignment) + else: + self._cinit_other(other, alignment) + + def __dealloc__(self): + if not (self._memory_ptr is NULL): + mkl_free(self._memory_ptr) + self._cinit_empty() + + cdef void *get_data_ptr(self): + return self._memory_ptr + + def __getbuffer__(self, Py_buffer *buffer, int flags): + buffer.buf = self._memory_ptr + buffer.format = "B" # byte + buffer.internal = NULL # see References + buffer.itemsize = 1 + buffer.len = self.nbytes + buffer.ndim = 1 + buffer.obj = self + buffer.readonly = 0 + buffer.shape = &self.nbytes + buffer.strides = &buffer.itemsize + buffer.suboffsets = NULL # for pointer arrays only + + def __releasebuffer__(self, Py_buffer *buffer): + pass diff --git a/mkl/_mkl_service.pxd b/mkl/_mkl_service.pxd index ed5a106..d839c14 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -149,6 +149,8 @@ cdef extern from "mkl.h": MKL_INT64 mkl_mem_stat(int* buf) MKL_INT64 mkl_peak_mem_usage(int mode) int mkl_set_memory_limit(int mem_type, size_t limit) + void *mkl_malloc(size_t size, int alignment) nogil + void mkl_free(void *ptr) nogil # Conditional Numerical Reproducibility int mkl_cbwr_set(int settings) From d1d7211f3e6ca74d8d4c0caa3e638ffa5d4f0dcf Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 21:10:25 -0700 Subject: [PATCH 02/16] add realloc to MKLMemory as atomics are a c11+ feature, specific flags are needed to enable on window --- meson.build | 8 ++++++++ mkl/_mkl_memory.pyx | 31 +++++++++++++++++++++++++++++-- mkl/_mkl_service.pxd | 1 + 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/meson.build b/meson.build index 32ae0f2..9308a3b 100644 --- a/meson.build +++ b/meson.build @@ -8,6 +8,7 @@ project( ).stdout().strip(), meson_version: '>=1.8.3', default_options: [ + 'c_std=c11', 'buildtype=release', ] ) @@ -25,6 +26,13 @@ endif thread_dep = dependency('threads') cc = meson.get_compiler('c') +if cc.get_id() == 'msvc' + add_project_arguments( + '/experimental:c11atomics', + language: 'c' + ) +endif + mkl_dep = dependency('MKL', method: 'cmake', modules: ['MKL::MKL'], cmake_args: [ diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index bc4a87d..c6450d4 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -31,16 +31,25 @@ import numbers from cpython cimport Py_buffer from libc.string cimport memcpy -from mkl._mkl_service cimport mkl_malloc, mkl_free +from mkl._mkl_service cimport mkl_malloc, mkl_realloc, mkl_free + +cdef extern from "stdatomic.h" nogil: + ctypedef int atomic_int "_Atomic int" + void atomic_init(atomic_int *obj, int value) + int atomic_fetch_add(atomic_int *obj, int value) + int atomic_fetch_sub(atomic_int *obj, int value) + int atomic_load(atomic_int *obj) cdef class MKLMemory: cdef void *_memory_ptr cdef Py_ssize_t nbytes + cdef atomic_int exported_buffers cdef _cinit_empty(self): self._memory_ptr = NULL self.nbytes = 0 + atomic_init(&self.exported_buffers, 0) cdef _cinit_alloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): self._cinit_empty() @@ -100,5 +109,23 @@ cdef class MKLMemory: buffer.strides = &buffer.itemsize buffer.suboffsets = NULL # for pointer arrays only + atomic_fetch_add(&self.exported_buffers, 1) + def __releasebuffer__(self, Py_buffer *buffer): - pass + atomic_fetch_sub(&self.exported_buffers, 1) + + def realloc(self, Py_ssize_t new_nbytes): + if atomic_load(&self.exported_buffers) > 0: + raise BufferError("Cannot realloc memory while there are exported buffers.") + 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 diff --git a/mkl/_mkl_service.pxd b/mkl/_mkl_service.pxd index d839c14..29854f7 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -150,6 +150,7 @@ cdef extern from "mkl.h": MKL_INT64 mkl_peak_mem_usage(int mode) int mkl_set_memory_limit(int mem_type, size_t limit) void *mkl_malloc(size_t size, int alignment) nogil + void *mkl_realloc(void *ptr, size_t size) nogil void mkl_free(void *ptr) nogil # Conditional Numerical Reproducibility From c5fdcd586117eceba8535fa5e4d8aa873be3b312 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 21:18:53 -0700 Subject: [PATCH 03/16] overload MKLMemory constructor to use mkl_calloc --- mkl/_mkl_memory.pyx | 74 +++++++++++++++++++++++++++++++++++--------- mkl/_mkl_service.pxd | 1 + 2 files changed, 61 insertions(+), 14 deletions(-) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index c6450d4..49a2365 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -31,7 +31,8 @@ import numbers from cpython cimport Py_buffer from libc.string cimport memcpy -from mkl._mkl_service cimport mkl_malloc, mkl_realloc, mkl_free +from mkl._mkl_service cimport mkl_calloc, mkl_free, mkl_malloc, mkl_realloc + cdef extern from "stdatomic.h" nogil: ctypedef int atomic_int "_Atomic int" @@ -51,7 +52,7 @@ cdef class MKLMemory: self.nbytes = 0 atomic_init(&self.exported_buffers, 0) - cdef _cinit_alloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): + cdef _cinit_malloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): self._cinit_empty() if (nbytes > 0): @@ -67,26 +68,71 @@ cdef class MKLMemory: ) else: raise ValueError( - "Number of bytes of request allocation must be positive." + "Number of bytes of requested allocation must be positive." ) - cdef _cinit_other(self, object other, Py_ssize_t alignment): - cdef MKLMemory other_mem - if isinstance(other, MKLMemory): - other_mem = other + cdef _cinit_calloc(self, Py_ssize_t num, Py_ssize_t size, Py_ssize_t alignment): + self._cinit_empty() + + if (num > 0 and size > 0): + with nogil: + p = mkl_calloc(num, size, alignment) + + if (p): + self._memory_ptr = p + self.nbytes = num * size + else: + raise MemoryError( + "MKL memory allocation failed." + ) else: raise ValueError( - f"Argument {other} is not of type MKLMemory." + "Number of elements and size of requested allocation must be " + "positive." ) - self._cinit_alloc(other_mem.nbytes, alignment) + + cdef _cinit_mklmemory(self, object other, Py_ssize_t alignment): + other_mem = other + + self._cinit_malloc(other_mem.nbytes, alignment) with nogil: memcpy(self._memory_ptr, other_mem._memory_ptr, self.nbytes) - def __cinit__(self, other, *, Py_ssize_t alignment=64): - if isinstance(other, numbers.Integral): - self._cinit_alloc(other, alignment) - else: - self._cinit_other(other, alignment) + def __cinit__(self, *args, **kwargs): + cdef Py_ssize_t alignment = kwargs.get("alignment", 64) + + n_args = len(args) + if not (0 < n_args < 3): + raise TypeError( + "MKLMemory constructor takes 1 or 2 arguments, but " + f"{n_args} were given" + ) + if n_args == 1: + arg = args[0] + if isinstance(arg, numbers.Integral): + self._cinit_malloc(arg, alignment) + elif isinstance(arg, MKLMemory): + self._cinit_mklmemory(arg, alignment) + else: + raise TypeError( + "MKLMemory single argument constructor expects an integer " + f"or MKLMemory instance, but got {type(arg)}" + ) + + elif n_args == 2: + arg0, arg1 = args[0], args[1] + if not isinstance(arg0, numbers.Integral): + raise TypeError( + "MKLMemory constructor expects first argument " + f"to be an integer, but got {type(arg0)}" + ) + if not isinstance(arg1, numbers.Integral): + raise TypeError( + "MKLMemory constructor expects second argument " + f"to be an integer, but got {type(arg1)}" + ) + + self._cinit_calloc(arg0, arg1, alignment) def __dealloc__(self): if not (self._memory_ptr is NULL): diff --git a/mkl/_mkl_service.pxd b/mkl/_mkl_service.pxd index 29854f7..c376e37 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -151,6 +151,7 @@ cdef extern from "mkl.h": int mkl_set_memory_limit(int mem_type, size_t limit) void *mkl_malloc(size_t size, int alignment) nogil void *mkl_realloc(void *ptr, size_t size) nogil + void *mkl_calloc(size_t num, size_t size, int alignment) nogil void mkl_free(void *ptr) nogil # Conditional Numerical Reproducibility From b5ea0766126029ace2d2f9541e33022b09bebc93 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 22:10:59 -0700 Subject: [PATCH 04/16] add info properties to MKLMemory --- mkl/_mkl_memory.pyx | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index 49a2365..d63bf8b 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -175,3 +175,27 @@ cdef class MKLMemory: self._memory_ptr = p self.nbytes = new_nbytes + + @property + def nbytes(self): + return self.nbytes + + @property + def size(self): + return self.nbytes + + @property + def _pointer(self): + return (self._memory_ptr) + + def __repr__(self): + return ( + f"(self._memory_ptr))}>" + ) + + def __len__(self): + return self.nbytes + + def __sizeof__(self): + return self.nbytes From 20995e7a3192007670bfd8176a493633551a5cf7 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 23:39:23 -0700 Subject: [PATCH 05/16] add pickling support for MKLMemory --- mkl/_mkl_memory.pyx | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index d63bf8b..12bb1e3 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -42,14 +42,29 @@ cdef extern from "stdatomic.h" nogil: int atomic_load(atomic_int *obj) +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) + + cdef void *dst = mem._memory_ptr + cdef char *src = data + + with nogil: + memcpy(dst, src, nbytes) + + return mem + + cdef class MKLMemory: cdef void *_memory_ptr 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 atomic_init(&self.exported_buffers, 0) cdef _cinit_malloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): @@ -62,6 +77,7 @@ cdef class MKLMemory: if (p): self._memory_ptr = p self.nbytes = nbytes + self.alignment = alignment else: raise MemoryError( "MKL memory allocation failed." @@ -81,6 +97,7 @@ cdef class MKLMemory: if (p): self._memory_ptr = p self.nbytes = num * size + self.alignment = alignment else: raise MemoryError( "MKL memory allocation failed." @@ -176,6 +193,10 @@ cdef class MKLMemory: self._memory_ptr = p self.nbytes = new_nbytes + def tobytes(self): + cdef char* data_ptr = self._memory_ptr + return data_ptr[:self.nbytes] + @property def nbytes(self): return self.nbytes @@ -184,6 +205,10 @@ cdef class MKLMemory: def size(self): return self.nbytes + @property + def alignment(self): + return self.alignment + @property def _pointer(self): return (self._memory_ptr) @@ -199,3 +224,6 @@ cdef class MKLMemory: def __sizeof__(self): return self.nbytes + + def __reduce__(self): + return (_mkl_memory_from_bytes, (self.tobytes(), self.alignment)) From f85d0b547046c39c86a348815a868c51f415b275 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 23:53:17 -0700 Subject: [PATCH 06/16] propagate alignment in MKLMemory --- mkl/_mkl_memory.pyx | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index 12bb1e3..77b6d78 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -116,7 +116,7 @@ cdef class MKLMemory: memcpy(self._memory_ptr, other_mem._memory_ptr, self.nbytes) def __cinit__(self, *args, **kwargs): - cdef Py_ssize_t alignment = kwargs.get("alignment", 64) + cdef Py_ssize_t alignment n_args = len(args) if not (0 < n_args < 3): @@ -127,8 +127,10 @@ cdef class MKLMemory: if n_args == 1: arg = args[0] if isinstance(arg, numbers.Integral): + alignment = kwargs.get("alignment", 64) self._cinit_malloc(arg, alignment) elif isinstance(arg, MKLMemory): + alignment = kwargs.get("alignment", arg.alignment) self._cinit_mklmemory(arg, alignment) else: raise TypeError( @@ -138,6 +140,7 @@ cdef class MKLMemory: elif n_args == 2: arg0, arg1 = args[0], args[1] + alignment = kwargs.get("alignment", 64) if not isinstance(arg0, numbers.Integral): raise TypeError( "MKLMemory constructor expects first argument " @@ -148,7 +151,6 @@ cdef class MKLMemory: "MKLMemory constructor expects second argument " f"to be an integer, but got {type(arg1)}" ) - self._cinit_calloc(arg0, arg1, alignment) def __dealloc__(self): From ffa398fe457c549e6241e084bbdb5788167d51ca Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sun, 12 Apr 2026 00:19:16 -0700 Subject: [PATCH 07/16] add tests for MKLMemory class --- mkl/tests/test_mkl_memory.py | 139 +++++++++++++++++++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 mkl/tests/test_mkl_memory.py diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py new file mode 100644 index 0000000..96f6f6f --- /dev/null +++ b/mkl/tests/test_mkl_memory.py @@ -0,0 +1,139 @@ +# Copyright (c) 2018, Intel Corporation +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of Intel Corporation nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import sys + +import mkl + + +def test_mkl_memory_create_malloc(): + nbytes = 1024 + mem = mkl.MKLMemory(nbytes) + assert mem.nbytes == nbytes + # default alignment is 64 bytes + assert mem.alignment == 64 + + +def test_mkl_memory_create_calloc(): + size = 32 + num = 32 + nbytes = num * size + # test creating with mkl_calloc + mem = mkl.MKLMemory(num, size) + assert mem.nbytes == nbytes + # default alignment is 64 bytes + assert mem.alignment == 64 + + +def test_mkl_memory_create_with_malloc_and_alignment(): + size = 32 + num = 32 + nbytes = num * size + alignment = 128 + mem = mkl.MKLMemory(nbytes, alignment=alignment) + assert mem.nbytes == nbytes + assert mem.alignment == alignment + + +def test_mkl_memory_create_with_calloc_and_alignment(): + size = 32 + num = 32 + nbytes = num * size + alignment = 128 + mem = mkl.MKLMemory(num, size, alignment=alignment) + assert mem.nbytes == nbytes + + +def test_mkl_memory_create_from_mkl_memory(): + mem1 = mkl.MKLMemory(1024) + mem2 = mkl.MKLMemory(mem1) + assert mem2.nbytes == mem1.nbytes + + +def test_mkl_memory_create_from_mkl_memory_with_alignment(): + mem1 = mkl.MKLMemory(1024) + alignment = 128 + mem2 = mkl.MKLMemory(mem1, alignment=alignment) + assert mem2.nbytes == mem1.nbytes + assert mem2.alignment == alignment + + +def test_mkl_memory_propagates_alignment(): + mem1 = mkl.MKLMemory(1024, alignment=128) + mem2 = mkl.MKLMemory(mem1) + assert mem2.nbytes == mem1.nbytes + assert mem2.alignment == mem1.alignment + + +def test_mkl_memory_properties(): + nbytes = 1024 + mem = mkl.MKLMemory(nbytes) + assert len(mem) == nbytes + assert type(repr(mem)) is str + assert type(bytes(mem)) is bytes + assert sys.getsizeof(mem) >= nbytes + + +def test_buffer_protocol(): + mem = mkl.MKLMemory(1024) + mv1 = memoryview(mem) + assert mv1.nbytes == mem.nbytes + mv2 = memoryview(mem) + assert mv1 == mv2 + + +def test_pickling(): + import pickle + + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + for i in range(len(mem)): + mv[i] = (i % 32) + ord("a") + + mem_reconstructed = pickle.loads(pickle.dumps(mem)) + assert type(mem) is type(mem_reconstructed), "Pickling should preserve type" + assert ( + mem.tobytes() == mem_reconstructed.tobytes() + ), "Pickling should preserve buffer content" + assert ( + mem._pointer != mem_reconstructed._pointer + ), "Pickling/unpickling should be changing pointer" + + +def test_pickling_with_alignment(): + import pickle + + mem = mkl.MKLMemory(1024, alignment=128) + mem_reconstructed = pickle.loads(pickle.dumps(mem)) + assert type(mem) is type(mem_reconstructed), "Pickling should preserve type" + assert ( + mem.tobytes() == mem_reconstructed.tobytes() + ), "Pickling should preserve buffer content" + assert ( + mem._pointer != mem_reconstructed._pointer + ), "Pickling/unpickling should be changing pointer" + assert ( + mem.alignment == mem_reconstructed.alignment + ), "Pickling should preserve alignment" From 42b59b0c7e1f9e6c402321ef2264473e76d17ed6 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sun, 12 Apr 2026 03:07:07 -0700 Subject: [PATCH 08/16] add nogil to MKL functions for freeing buffers --- mkl/_mkl_service.pxd | 10 +++++----- mkl/_py_mkl_service.pyx | 6 ++++-- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/mkl/_mkl_service.pxd b/mkl/_mkl_service.pxd index c376e37..4a2d789 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -24,7 +24,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -cdef extern from "mkl.h": +cdef extern from "mkl.h" nogil: # defer definition of integer types to mkl.h # Cython will narrow the types based on what mkl.h defines ctypedef long long MKL_INT64 @@ -149,10 +149,10 @@ cdef extern from "mkl.h": MKL_INT64 mkl_mem_stat(int* buf) MKL_INT64 mkl_peak_mem_usage(int mode) int mkl_set_memory_limit(int mem_type, size_t limit) - void *mkl_malloc(size_t size, int alignment) nogil - void *mkl_realloc(void *ptr, size_t size) nogil - void *mkl_calloc(size_t num, size_t size, int alignment) nogil - void mkl_free(void *ptr) nogil + void *mkl_malloc(size_t size, int alignment) + void *mkl_realloc(void *ptr, size_t size) + void *mkl_calloc(size_t num, size_t size, int alignment) + void mkl_free(void *ptr) # Conditional Numerical Reproducibility int mkl_cbwr_set(int settings) diff --git a/mkl/_py_mkl_service.pyx b/mkl/_py_mkl_service.pyx index 72908fe..af4ae3d 100644 --- a/mkl/_py_mkl_service.pyx +++ b/mkl/_py_mkl_service.pyx @@ -602,7 +602,8 @@ cdef inline void __free_buffers() noexcept: """ Frees unused memory allocated by the Intel(R) MKL Memory Allocator. """ - mkl.mkl_free_buffers() + with nogil: + mkl.mkl_free_buffers() return @@ -611,7 +612,8 @@ cdef inline void __thread_free_buffers() noexcept: Frees unused memory allocated by the Intel(R) MKL Memory Allocator in the current thread. """ - mkl.mkl_thread_free_buffers() + with nogil: + mkl.mkl_thread_free_buffers() return From 910d982d525c76ec537e453bc3fd140d4de4d170 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 09:39:44 -0700 Subject: [PATCH 09/16] fix meson.build rpath --- meson.build | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/meson.build b/meson.build index 9308a3b..b130247 100644 --- a/meson.build +++ b/meson.build @@ -84,7 +84,7 @@ py.extension_module( sources: ['mkl/_mkl_memory.pyx'], dependencies: [mkl_dep], c_args: c_args, - install_rpath: rpath, + link_args: rpath_link_args, install: true, subdir: 'mkl' ) From 4b88d62995838223d018b9e4ed65f8677709bd85 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 09:40:02 -0700 Subject: [PATCH 10/16] mark _mkl_memory free-threading compatible --- mkl/_mkl_memory.pyx | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index 77b6d78..d4c8033 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -1,4 +1,4 @@ -# Copyright (c) 2018, Intel Corporation +# Copyright (c) 2026, Intel Corporation # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: @@ -25,6 +25,7 @@ # distutils: language = c # cython: language_level=3 +# cython: freethreading_compatible=True import numbers @@ -163,8 +164,8 @@ cdef class MKLMemory: def __getbuffer__(self, Py_buffer *buffer, int flags): buffer.buf = self._memory_ptr - buffer.format = "B" # byte - buffer.internal = NULL # see References + buffer.format = "B" + buffer.internal = NULL buffer.itemsize = 1 buffer.len = self.nbytes buffer.ndim = 1 @@ -172,7 +173,7 @@ cdef class MKLMemory: buffer.readonly = 0 buffer.shape = &self.nbytes buffer.strides = &buffer.itemsize - buffer.suboffsets = NULL # for pointer arrays only + buffer.suboffsets = NULL atomic_fetch_add(&self.exported_buffers, 1) From 8cc7b4336efdd159d992ccd168a5a60f4d9da971 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 11:23:41 -0700 Subject: [PATCH 11/16] update meson.build --- meson.build | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/meson.build b/meson.build index b130247..407d85f 100644 --- a/meson.build +++ b/meson.build @@ -101,6 +101,9 @@ py.install_sources( ) py.install_sources( - ['mkl/tests/test_mkl_service.py'], + [ + 'mkl/tests/test_mkl_memory.py', + 'mkl/tests/test_mkl_service.py', + ], subdir: 'mkl/tests' ) From 6aae4fb28c1168e59ea6a9fabdc92da286173152 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 3 Sep 2026 14:34:09 -0700 Subject: [PATCH 12/16] align realloc behavior to NumPy also address issues with undeclared variables and rename MKLMemory class members --- mkl/_mkl_memory.pyx | 142 ++++++++++++++++++++++++++--------- mkl/tests/test_mkl_memory.py | 132 ++++++++++++++++++++++++++++++++ 2 files changed, 239 insertions(+), 35 deletions(-) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index d4c8033..ab6cfe4 100644 --- a/mkl/_mkl_memory.pyx +++ b/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 @@ -41,6 +42,41 @@ cdef extern from "stdatomic.h" nogil: int atomic_fetch_add(atomic_int *obj, int value) int atomic_fetch_sub(atomic_int *obj, int value) int atomic_load(atomic_int *obj) + void atomic_store(atomic_int *obj, int value) + bint atomic_compare_exchange_strong( + atomic_int *obj, int *expected, int desired + ) + + +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 > INT_MAX: + raise ValueError( + f"Alignment of requested allocation must not exceed {INT_MAX}." + ) + return alignment def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment): @@ -57,28 +93,35 @@ 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 + # prevents simultaneous reallocs + cdef atomic_int realloc_in_progress 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) + atomic_init(&self.realloc_in_progress, 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 +132,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 +156,11 @@ cdef class MKLMemory: ) cdef _cinit_mklmemory(self, object other, Py_ssize_t alignment): - other_mem = other + cdef MKLMemory other_mem = 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 +177,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", (arg)._alignment) self._cinit_mklmemory(arg, alignment) else: raise TypeError( @@ -167,11 +213,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,36 +227,62 @@ cdef class MKLMemory: atomic_fetch_sub(&self.exported_buffers, 1) def realloc(self, Py_ssize_t new_nbytes): - if atomic_load(&self.exported_buffers) > 0: - raise BufferError("Cannot realloc memory while there are exported buffers.") - 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) + cdef int shared + cdef int unclaimed = 0 + + # claim the exclusive right to reallocate before doing anything else + if not atomic_compare_exchange_strong( + &self.realloc_in_progress, &unclaimed, 1 + ): + raise BufferError( + "Cannot realloc memory while another thread is reallocating it." + ) + try: + if atomic_load(&self.exported_buffers) > 0: + 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.") + + with nogil: + p = mkl_realloc(self._memory_ptr, new_nbytes) - if not p: - raise MemoryError("MKL memory reallocation failed.") + if not p: + raise MemoryError("MKL memory reallocation failed.") - self._memory_ptr = p - self.nbytes = new_nbytes + self._memory_ptr = p + self._nbytes = new_nbytes + finally: + atomic_store(&self.realloc_in_progress, 0) def tobytes(self): cdef char* data_ptr = 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): @@ -218,15 +290,15 @@ cdef class MKLMemory: def __repr__(self): return ( - f"(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)) diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py index 96f6f6f..c20f3cb 100644 --- a/mkl/tests/test_mkl_memory.py +++ b/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,132 @@ 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_never_overlaps(): + initial = 64 + sizes = (1 << 16, 1 << 17) + + for _ in range(50): + mem = mkl.MKLMemory(initial) + barrier = threading.Barrier(len(sizes)) + results = [None] * len(sizes) + + def worker(idx, size, mem=mem, barrier=barrier, results=results): + barrier.wait() + try: + mem.realloc(size) + results[idx] = "ok" + except (ValueError, BufferError): + results[idx] = "refused" + + ts = [ + threading.Thread(target=worker, args=(idx, size)) + for idx, size in enumerate(sizes) + ] + for t in ts: + t.start() + for t in ts: + t.join() + + assert all( + r in ("ok", "refused") for r in results + ), f"realloc raised an unexpected error: {results}" + allowed = {initial, *sizes} + assert ( + len(mem) in allowed + ), f"Inconsistent size {len(mem)} from {results}" + assert mem.nbytes == len(mem) + assert len(mem.tobytes()) == len(mem) + + mv = memoryview(mem) + try: + mv[0] = 1 + mv[len(mem) - 1] = 2 + finally: + mv.release() + + +def test_realloc_refused_while_another_thread_holds_reference(): + mem = mkl.MKLMemory(64) + holder_ready = threading.Event() + release_holder = threading.Event() + outcome = [] + + def holder(): + # keep reference alive + alias = mem # noqa: F841 + holder_ready.set() + release_holder.wait(timeout=30) + + t = threading.Thread(target=holder) + t.start() + try: + assert holder_ready.wait(timeout=30) + try: + mem.realloc(1 << 16) + outcome.append("ok") + except ValueError: + outcome.append("refused") + finally: + release_holder.set() + t.join() + + assert outcome == [ + "refused" + ], f"Expected refusal while shared, got {outcome}" + assert len(mem) == 64, "Refused realloc must not change the buffer" From 1ed208f37013548bc69bf419dcb6e7861621fa39 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 8 Sep 2026 09:53:52 -0700 Subject: [PATCH 13/16] Fix date in test_mkl_memory.py Co-authored-by: Anton <100830759+antonwolfy@users.noreply.github.com> --- mkl/tests/test_mkl_memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py index c20f3cb..19f2930 100644 --- a/mkl/tests/test_mkl_memory.py +++ b/mkl/tests/test_mkl_memory.py @@ -1,4 +1,4 @@ -# Copyright (c) 2018, Intel Corporation +# Copyright (c) 2026, Intel Corporation # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: From 34f8ec3944d2ea384346801bb5e6ffb8b3e4e19c Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Fri, 11 Sep 2026 09:57:44 -0700 Subject: [PATCH 14/16] Apply review feedback --- .github/copilot-instructions.md | 6 +- AGENTS.md | 6 +- CHANGELOG.md | 1 + README.md | 1 + meson.build | 15 +- mkl/AGENTS.md | 11 ++ mkl/__init__.py | 1 + mkl/_mkl_memory.pyx | 186 +++++++++++++++---- mkl/tests/AGENTS.md | 8 + mkl/tests/test_mkl_memory.py | 313 +++++++++++++++++++++++++++++--- 10 files changed, 471 insertions(+), 77 deletions(-) diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 435e801..4c4abd5 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -20,7 +20,7 @@ Higher-precedence file overrides; lower must not restate overridden guidance. ## Contribution expectations - Keep diffs minimal; prefer atomic single-purpose commits. - Preserve public API signatures in `mkl/__init__.py` unless change is explicitly requested. -- For user-visible behavior changes: update tests in `mkl/tests/test_mkl_service.py`. +- For user-visible behavior changes: update tests in `mkl/tests/test_mkl_service.py`, or `mkl/tests/test_mkl_memory.py` for `MKLMemory`. - For bug fixes: add or extend regression tests in the same change. - Do not generate code without corresponding test updates when behavior changes. - Run `pre-commit run --all-files` when `.pre-commit-config.yaml` is present. @@ -37,8 +37,8 @@ Higher-precedence file overrides; lower must not restate overridden guidance. - Build/config: `pyproject.toml`, `meson.build` - Recipe/deps: `conda-recipe/meta.yaml`, `conda-recipe/conda_build_config.yaml` - CI: `.github/workflows/*.{yml,yaml}` -- API contracts: `mkl/__init__.py`, `mkl/_py_mkl_service.pyx` -- Tests: `mkl/tests/test_mkl_service.py` +- API contracts: `mkl/__init__.py`, `mkl/_py_mkl_service.pyx`, `mkl/_mkl_memory.pyx` +- Tests: `mkl/tests/test_mkl_service.py`, `mkl/tests/test_mkl_memory.py` ## MKL-specific constraints - Linux runtime init path may require `RTLD_GLOBAL` preloading (`mkl/_mklinitmodule.c`). diff --git a/AGENTS.md b/AGENTS.md index 55f57d7..60a524d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,6 +7,7 @@ Entry point for agent context in this repo. - Threading control (set/get number of threads, domain-specific threading) - Version information (MKL version, build info) - Memory management (peak memory usage, memory statistics) +- Aligned memory allocation (`MKLMemory`, a buffer-protocol object backed by `mkl_malloc`) - Conditional Numerical Reproducibility (CNR) - Timing functions (get CPU/wall clock time) - Miscellaneous utilities (MKL_VERBOSE control, etc.) @@ -16,6 +17,7 @@ Originally part of Intel® Distribution for Python*, now a standalone package av ## Key components - **Python interface:** `mkl/__init__.py` — public API surface - **Cython wrapper:** `mkl/_py_mkl_service.pyx` — wraps MKL support functions +- **Cython allocator:** `mkl/_mkl_memory.pyx` — `MKLMemory`, wraps `mkl_malloc`/`mkl_calloc`/`mkl_realloc`/`mkl_free` - **C init module:** `mkl/_mklinitmodule.c` — Linux-side MKL runtime preloading / initialization - **Helper:** `mkl/_init_helper.py` — Windows venv DLL loading helper - **Build system:** meson-python + Cython @@ -73,11 +75,11 @@ mkl.get_version_string() # MKL version info - **API stability:** Preserve existing function signatures (widely used in ecosystem) - **Threading:** Changes to threading control must be thread-safe - **CNR:** Conditional Numerical Reproducibility flags require careful documentation -- **Testing:** Add tests to `mkl/tests/test_mkl_service.py` +- **Testing:** Add tests to `mkl/tests/test_mkl_service.py`, or `mkl/tests/test_mkl_memory.py` for `MKLMemory` - **Docs:** MKL support functions documented in [Intel oneMKL Developer Reference](https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2025-2/support-functions.html) ## Code structure -- **Cython layer:** `_py_mkl_service.pyx` + `_mkl_service.pxd` (C declarations) +- **Cython layer:** `_py_mkl_service.pyx` and `_mkl_memory.pyx` + `_mkl_service.pxd` (C declarations) - **C init:** `_mklinitmodule.c` handles Linux preloading (`dlopen(..., RTLD_GLOBAL)`) for MKL runtime - **Windows loading helper:** `_init_helper.py` handles DLL path setup in Windows venv - **Python wrapper:** `__init__.py` imports `_py_mkl_service` (generated from `.pyx`) diff --git a/CHANGELOG.md b/CHANGELOG.md index d267ff7..6172fb4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added * Added support for free-threaded (GIL-disabled) CPython builds: the Cython extension is compiled with `freethreading_compatible=True` and `_mklinit` declares `Py_MOD_GIL_NOT_USED`, so importing `mkl` no longer re-enables the GIL [gh-213](https://github.com/IntelPython/mkl-service/pull/213) * Added support for new build option `ilp64` to initialize MKL with the ILP64 interface, which also resolves some build warnings [gh-184](https://github.com/IntelPython/mkl-service/pull/184) +* Exposed `mkl_malloc` and related MKL calls to Python via `MKLMemory` class which supports the Python buffer protocol [gh-182](https://github.com/IntelPython/mkl-service/pull/182) ### Changed * Raised the minimum build-time `Cython` requirement to `3.1.0`, the first release providing the `freethreading_compatible` directive [gh-213](https://github.com/IntelPython/mkl-service/pull/213) diff --git a/README.md b/README.md index 3eb6e97..cb51f4c 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ For more information about the usage of support functions see [Developer Referen ## Building A C compiler and Intel(R) oneAPI Math Kernel Library (oneMKL) are required to build mkl-service from source. +The compiler must support C11 atomics (i.e., for Windows, Visual Studio 2022 17.5 or newer). Executing ```sh diff --git a/meson.build b/meson.build index 407d85f..b29ca79 100644 --- a/meson.build +++ b/meson.build @@ -26,10 +26,17 @@ endif thread_dep = dependency('threads') cc = meson.get_compiler('c') + +atomics_args = [] if cc.get_id() == 'msvc' - add_project_arguments( - '/experimental:c11atomics', - language: 'c' + atomics_args += '/experimental:c11atomics' +endif + +# checked to fail early if missing header +if not cc.has_header('stdatomic.h', args: atomics_args) + error( + 'mkl-service requires a C compiler supporting C11 atomics', + '(i.e., for Windows, Visual Studio 2022 17.5 or newer).' ) endif @@ -83,7 +90,7 @@ py.extension_module( '_mkl_memory', sources: ['mkl/_mkl_memory.pyx'], dependencies: [mkl_dep], - c_args: c_args, + c_args: c_args + atomics_args, link_args: rpath_link_args, install: true, subdir: 'mkl' diff --git a/mkl/AGENTS.md b/mkl/AGENTS.md index 00dc3a6..8f832d7 100644 --- a/mkl/AGENTS.md +++ b/mkl/AGENTS.md @@ -5,6 +5,7 @@ Core Python/Cython implementation: MKL support function wrappers and runtime con ## Structure - `__init__.py` — public API, RTLD_GLOBAL context manager, module initialization - `_py_mkl_service.pyx` — Cython wrappers for MKL support functions +- `_mkl_memory.pyx` — `MKLMemory`, a buffer-protocol object over MKL's allocator - `_mkl_service.pxd` — Cython declarations (C function signatures) - `_mklinitmodule.c` — C extension for Linux-side MKL runtime preloading/init - `_init_helper.py` — Windows loading helper (DLL path setup in venv) @@ -26,6 +27,13 @@ Core Python/Cython implementation: MKL support function wrappers and runtime con - `peak_mem_usage(memtype)` — peak memory usage stats - `mem_stat()` — memory allocation statistics +### Memory allocation +- `MKLMemory(nbytes, alignment=64)` — aligned allocation via `mkl_malloc` +- `MKLMemory(num, elem_size, alignment=64)` — zeroed allocation via `mkl_calloc` +- `MKLMemory(other, alignment=other.alignment)` — copy of another allocation +- `realloc(new_nbytes, refcheck=True)` — resize in place via `mkl_realloc` +- `nbytes` / `__len__`, `alignment`, `tobytes()`, buffer protocol, pickling + ### CNR (Conditional Numerical Reproducibility) - `set_num_threads_local(n)` — thread-local thread count - CNR mode control functions @@ -39,11 +47,14 @@ Core Python/Cython implementation: MKL support function wrappers and runtime con - **API stability:** Preserve function signatures (widely used in ecosystem) - **MKL dependency:** Assumes MKL is available at runtime (conda: mkl package). Do **not** list `mkl` in `pyproject.toml` `[project].dependencies` — its PyPI wheel lacks `.dist-info`, which breaks `pip check`; on conda-forge there is no pip-visible `mkl` distribution. - **RTLD_GLOBAL preload path:** Linux preload is handled in `_mklinitmodule.c`; Windows DLL setup is in `_init_helper.py` +- **`MKLMemory` mutation:** `realloc` moves the underlying block, so it must refuse while a buffer is exported, while another thread is resizing, or (unless `refcheck=False`) while the object looks referenced elsewhere. The GIL must not be released across those checks and the pointer store, mirroring NumPy's `PyArray_Resize`. The reference-count check stays NumPy's: `PyUnstable_Object_IsUniquelyReferenced` from 3.14, `Py_REFCNT > 2` before it, keyed on `PY_VERSION_HEX` and not on `Py_GIL_DISABLED`. It is a check against dangling references, not against other threads — on a free-threaded build before 3.14 it cannot be either, and resizing an allocation another thread can reach is the caller's responsibility, as it is for `numpy.ndarray.resize`. ## Cython details - `_py_mkl_service.pyx` → generates `_py_mkl_service` extension module +- `_mkl_memory.pyx` → generates `_mkl_memory` extension module - `.pxd` file declares external C functions from MKL headers - Cython build requires MKL headers (`mkl-devel`) +- `_mkl_memory.pyx` uses C11 atomics (``); `meson.build` scopes MSVC's `/experimental:c11atomics` to that one target ## C init module - `_mklinitmodule.c` → `_mklinit` extension diff --git a/mkl/__init__.py b/mkl/__init__.py index beadbfc..d1ec7c2 100644 --- a/mkl/__init__.py +++ b/mkl/__init__.py @@ -122,6 +122,7 @@ def __exit__(self, *args): "mem_stat", "peak_mem_usage", "set_memory_limit", + "MKLMemory", "cbwr_set", "cbwr_get", "cbwr_get_auto_branch", diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index ab6cfe4..0a152e2 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -36,6 +36,10 @@ from libc.string cimport memcpy from mkl._mkl_service cimport mkl_calloc, mkl_free, mkl_malloc, mkl_realloc +cdef extern from "Python.h": + const Py_ssize_t PY_SSIZE_T_MAX + + cdef extern from "stdatomic.h" nogil: ctypedef int atomic_int "_Atomic int" void atomic_init(atomic_int *obj, int value) @@ -50,8 +54,8 @@ cdef extern from "stdatomic.h" nogil: cdef extern from *: """ - // Check whether a MKLMemory object may be safely reallocated. - // Mirrors NumPy's PyArray_Resize_int logic. + // 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)) { @@ -69,10 +73,29 @@ cdef extern from *: int _MKLMemory_MayBeShared(object obj) -cdef int _check_alignment(Py_ssize_t alignment) except -1: +cdef _extract_alignment(dict kwargs, object default): + """ + Return the ``alignment`` keyword, or `default` when it was not given. + """ + for name in kwargs: + if name != "alignment": + raise TypeError( + "MKLMemory constructor got an unexpected keyword argument " + f"'{name}'" + ) + + return kwargs.get("alignment", default) + + +cdef int _check_alignment(object alignment) except -1: + if not isinstance(alignment, numbers.Integral): + raise TypeError( + "Alignment of requested allocation must be an integer, but got " + f"{type(alignment)}" + ) if alignment <= 0: raise ValueError("Alignment of requested allocation must be positive.") - if alignment > INT_MAX: + if alignment > INT_MAX: raise ValueError( f"Alignment of requested allocation must not exceed {INT_MAX}." ) @@ -93,7 +116,36 @@ def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment): cdef class MKLMemory: - """MKL-backed memory object that exposes Python buffer protocol.""" + """ + MKLMemory(nbytes, alignment=64) + MKLMemory(num, elem_size, alignment=64) + MKLMemory(other, alignment=other.alignment) + + An object representing an aligned allocation made by oneMKL's allocator, + exposed through the Python buffer protocol. + + The first form allocates ``nbytes`` uninitialized bytes with + ``mkl_malloc``, the second ``num * elem_size`` zeroed bytes with + ``mkl_calloc``, and the third a copy of the content of another + :class:`MKLMemory`. + + Args: + nbytes (int): + number of bytes to allocate. + Expected to be positive. + num (int): + number of elements to allocate. + Expected to be positive. + elem_size (int): + size of a single element in bytes. + Expected to be positive. + other (:class:`MKLMemory`): + allocation whose size and content the new allocation takes. + alignment (Optional[int]): + address alignment of the allocation in bytes. Expected to be + positive and to not exceed ``INT_MAX``. Defaults to the alignment + of ``other`` in the copy form, and to `64` otherwise. + """ cdef void *_memory_ptr cdef Py_ssize_t _nbytes cdef Py_ssize_t _alignment @@ -108,7 +160,7 @@ cdef class MKLMemory: atomic_init(&self.exported_buffers, 0) atomic_init(&self.realloc_in_progress, 0) - cdef _cinit_malloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): + cdef _cinit_malloc(self, Py_ssize_t nbytes, object alignment): cdef int c_alignment = _check_alignment(alignment) cdef void *p @@ -121,7 +173,7 @@ cdef class MKLMemory: if (p): self._memory_ptr = p self._nbytes = nbytes - self._alignment = alignment + self._alignment = c_alignment else: raise MemoryError( "MKL memory allocation failed." @@ -131,31 +183,41 @@ cdef class MKLMemory: "Number of bytes of requested allocation must be positive." ) - cdef _cinit_calloc(self, Py_ssize_t num, Py_ssize_t size, Py_ssize_t alignment): + cdef _cinit_calloc( + self, Py_ssize_t num, Py_ssize_t elem_size, object alignment + ): cdef int c_alignment = _check_alignment(alignment) + cdef Py_ssize_t nbytes cdef void *p self._cinit_empty() - if (num > 0 and size > 0): + if (num > 0 and elem_size > 0): + if num > PY_SSIZE_T_MAX // elem_size: + raise ValueError( + "Total size of requested allocation must not exceed " + f"{PY_SSIZE_T_MAX} bytes." + ) + nbytes = num * elem_size + with nogil: - p = mkl_calloc(num, size, c_alignment) + p = mkl_calloc(num, elem_size, c_alignment) if (p): self._memory_ptr = p - self._nbytes = num * size - self._alignment = alignment + self._nbytes = nbytes + self._alignment = c_alignment else: raise MemoryError( "MKL memory allocation failed." ) else: raise ValueError( - "Number of elements and size of requested allocation must be " - "positive." + "Number of elements and element size of requested allocation " + "must be positive." ) - cdef _cinit_mklmemory(self, object other, Py_ssize_t alignment): + cdef _cinit_mklmemory(self, object other, object alignment): cdef MKLMemory other_mem = other self._cinit_malloc(other_mem._nbytes, alignment) @@ -163,8 +225,6 @@ cdef class MKLMemory: memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes) def __cinit__(self, *args, **kwargs): - cdef Py_ssize_t alignment - n_args = len(args) if not (0 < n_args < 3): raise TypeError( @@ -174,10 +234,12 @@ cdef class MKLMemory: if n_args == 1: arg = args[0] if isinstance(arg, numbers.Integral): - alignment = kwargs.get("alignment", 64) + alignment = _extract_alignment(kwargs, 64) self._cinit_malloc(arg, alignment) elif isinstance(arg, MKLMemory): - alignment = kwargs.get("alignment", (arg)._alignment) + alignment = _extract_alignment( + kwargs, (arg)._alignment + ) self._cinit_mklmemory(arg, alignment) else: raise TypeError( @@ -187,7 +249,7 @@ cdef class MKLMemory: elif n_args == 2: arg0, arg1 = args[0], args[1] - alignment = kwargs.get("alignment", 64) + alignment = _extract_alignment(kwargs, 64) if not isinstance(arg0, numbers.Integral): raise TypeError( "MKLMemory constructor expects first argument " @@ -226,7 +288,41 @@ cdef class MKLMemory: def __releasebuffer__(self, Py_buffer *buffer): atomic_fetch_sub(&self.exported_buffers, 1) - def realloc(self, Py_ssize_t new_nbytes): + def realloc(self, Py_ssize_t new_nbytes, *, bint refcheck=True): + """ + realloc(new_nbytes, refcheck=True) + + Resizes this allocation in place, keeping the content that fits. + + Args: + new_nbytes (int): + new size of the allocation in bytes. + Expected to be positive. + refcheck (Optional[bool]): + whether to refuse the resize when this object appears to be + referenced from elsewhere. + Default: `True`. + + Resizing moves the underlying memory, so any other reference to this + object would be left pointing at freed memory. The check for such + references is a heuristic based on the reference count and can refuse a + resize that would have been safe, especially in the case of a reference + reachable from more than one thread. + + Passing ``refcheck=False`` skips that check, and it is the caller's + responsibility to ensure that nothing else refers to this object and + that no other thread can reach it until the call returns. + + Neither the check nor its absence is a substitute for locking. Under the + GIL, and on free-threaded builds from Python 3.14 where the object can + be asked whether it is uniquely referenced, nothing else can reach the + object between the check and the resize. On a free-threaded build before + 3.14 there is neither, and a reference the caller holds cannot be told + apart from one another thread holds: resizing an allocation another + thread can reach may leave that thread reading freed memory whatever + ``refcheck`` is set to, so arrange for exclusive access. The same + applies to :meth:`numpy.ndarray.resize`. + """ cdef void *p cdef int shared cdef int unclaimed = 0 @@ -243,22 +339,29 @@ cdef class MKLMemory: 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 refcheck: + 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. If you are sure that this MKLMemory is " + "uniquely referenced, pass refcheck=False." + ) + elif shared == 2: + raise ValueError( + "Cannot realloc MKLMemory that is referenced by other " + "objects. Pass refcheck=False to realloc anyway, at the " + "risk of leaving those references pointing at freed " + "memory." + ) if new_nbytes <= 0: raise ValueError("New number of bytes must be positive.") - with nogil: - p = mkl_realloc(self._memory_ptr, new_nbytes) + # do not release the GIL here, as that can allow another thread to + # read the or export a buffer with the old pointer before + # mkl_realloc frees it + p = mkl_realloc(self._memory_ptr, new_nbytes) if not p: raise MemoryError("MKL memory reallocation failed.") @@ -269,29 +372,34 @@ cdef class MKLMemory: atomic_store(&self.realloc_in_progress, 0) def tobytes(self): + """ + Constructs bytes object populated with copy of this allocation. + """ cdef char* data_ptr = self._memory_ptr return data_ptr[:self._nbytes] @property def nbytes(self): - return self._nbytes - - @property - def size(self): + """Extent of this allocation in bytes.""" return self._nbytes @property def alignment(self): + """Address alignment of this allocation in bytes, as requested.""" return self._alignment @property def _pointer(self): + """ + Pointer to the start of this allocation + represented as Python integer. + """ return (self._memory_ptr) def __repr__(self): return ( f"(self._memory_ptr))}>" + f"{hex(self._pointer)}>" ) def __len__(self): diff --git a/mkl/tests/AGENTS.md b/mkl/tests/AGENTS.md index 021d695..955e7f5 100644 --- a/mkl/tests/AGENTS.md +++ b/mkl/tests/AGENTS.md @@ -4,6 +4,7 @@ Unit tests for MKL runtime control API. ## Test files - **test_mkl_service.py** — API functionality, threading control, version info +- **test_mkl_memory.py** — `MKLMemory` allocation, buffer protocol, `realloc`, concurrency ## Test coverage - Threading: `set_num_threads`, `get_max_threads`, domain-specific threading @@ -11,6 +12,10 @@ Unit tests for MKL runtime control API. - Memory: `peak_mem_usage`, `mem_stat` (if supported by MKL build) - CNR: Conditional Numerical Reproducibility flags - Edge cases currently covered: thread-local settings and API round-trips +- `MKLMemory` construction: all three forms, argument count/type errors, non-positive sizes, `num * elem_size` overflow, alignment bounds and types, unexpected keywords +- `MKLMemory` buffers: buffer protocol, `tobytes`, pickle round-trip, actual address alignment +- `MKLMemory.realloc`: grow/shrink with data preservation, alignment preserved across a resize, refusal while a buffer is exported or the object looks shared, `refcheck=False`, non-positive sizes +- `MKLMemory` concurrency: concurrent reads, overlapping `realloc` calls, readers racing a `realloc` ## Running tests ```bash @@ -24,5 +29,8 @@ pytest mkl/tests/ ## Adding tests - New API functions → add to `test_mkl_service.py` with validation +- `MKLMemory` behavior → add to `test_mkl_memory.py` - Threading behavior → test thread count changes take effect - Use `mkl.get_version()` to check MKL availability before tests +- Concurrency tests must be checked for vacuity: a `realloc` refused by every thread satisfies loose assertions without ever reaching `mkl_realloc` +- Tests must pass on free-threaded builds, where `realloc`'s reference-count check does not guard against other threads before 3.14: a test that races a resize against live readers must be gated on `REALLOC_RACE_IS_CONTAINED`, or it reads freed memory there instead of testing a guard diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py index 19f2930..3c869cd 100644 --- a/mkl/tests/test_mkl_memory.py +++ b/mkl/tests/test_mkl_memory.py @@ -30,6 +30,11 @@ import mkl +# on free-threaded Python prior to 3.14, only the caller can ensure exclusive +# access during realloc +_GIL_ENABLED = getattr(sys, "_is_gil_enabled", lambda: True)() +REALLOC_RACE_IS_CONTAINED = _GIL_ENABLED or sys.version_info >= (3, 14) + def test_mkl_memory_create_malloc(): nbytes = 1024 @@ -67,6 +72,15 @@ def test_mkl_memory_create_with_calloc_and_alignment(): alignment = 128 mem = mkl.MKLMemory(num, size, alignment=alignment) assert mem.nbytes == nbytes + assert mem.alignment == alignment + + +@pytest.mark.parametrize("alignment", [64, 128, 256]) +def test_allocation_is_actually_aligned(alignment): + assert mkl.MKLMemory(1024, alignment=alignment)._pointer % alignment == 0 + assert mkl.MKLMemory(32, 32, alignment=alignment)._pointer % alignment == 0 + source = mkl.MKLMemory(1024, alignment=alignment) + assert mkl.MKLMemory(source)._pointer % alignment == 0 def test_mkl_memory_create_from_mkl_memory(): @@ -142,6 +156,51 @@ def test_pickling_with_alignment(): ), "Pickling should preserve alignment" +def test_realloc_grow_and_shrink_preserves_data(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + for i in range(len(mem)): + mv[i] = i % 256 + mv.release() + original = mem.tobytes() + + grown = 1 << 20 + mem.realloc(grown) + assert mem.nbytes == grown + assert len(mem) == grown + assert len(mem.tobytes()) == grown + # growing keeps every byte that was there + assert mem.tobytes()[:1024] == original + + mem.realloc(256) + assert mem.nbytes == 256 + assert len(mem) == 256 + # shrinking keeps the surviving prefix + assert mem.tobytes() == original[:256] + + # and the resized buffer is still writable through the buffer protocol + mv = memoryview(mem) + try: + mv[0] = 7 + mv[len(mem) - 1] = 9 + finally: + mv.release() + assert mem.tobytes()[0] == 7 + assert mem.tobytes()[-1] == 9 + + +@pytest.mark.parametrize("alignment", [64, 128, 4096]) +def test_realloc_preserves_alignment(alignment): + # test that alignment is preserved by realloc, which is undocumented in MKL + # but holds experimentally + mem = mkl.MKLMemory(1024, alignment=alignment) + assert mem._pointer % alignment == 0 + for nbytes in (1 << 20, 256): + mem.realloc(nbytes) + assert mem.alignment == alignment + assert mem._pointer % alignment == 0 + + def test_realloc_exported_buffer(): mem = mkl.MKLMemory(1024) mv = memoryview(mem) @@ -152,19 +211,152 @@ def test_realloc_exported_buffer(): def test_realloc_refcheck_shared(): mem = mkl.MKLMemory(1024) - alias = mem # noqa: F841 — extra reference + alias = mem # noqa: F841 with pytest.raises(ValueError, match="referenced by"): mem.realloc(2048) del alias -def test_alignment_validation(): +def test_realloc_refcheck_false_allows_shared(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + for i in range(len(mem)): + mv[i] = i % 256 + del mv + + alias = mem # noqa: F841 + with pytest.raises(ValueError, match="refcheck=False"): + mem.realloc(2048) + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + assert len(mem) == 2048 + # the leading bytes must have survived the move + assert mem.tobytes()[:256] == bytes(range(256)) + del alias + + +def test_realloc_refcheck_false_still_refuses_exported_buffer(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + try: + with pytest.raises(BufferError): + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 1024 + finally: + mv.release() + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_realloc_validates_size(): + mem = mkl.MKLMemory(1024) + with pytest.raises(ValueError, match="positive"): + mem.realloc(0, refcheck=False) + with pytest.raises(ValueError, match="positive"): + mem.realloc(-1, refcheck=False) + assert mem.nbytes == 1024 + + +def test_realloc_refcheck_is_keyword_only(): + mem = mkl.MKLMemory(1024) + with pytest.raises(TypeError): + mem.realloc(2048, False) + assert mem.nbytes == 1024 + + +def test_constructor_argument_count(): + with pytest.raises(TypeError, match="takes 1 or 2 arguments"): + mkl.MKLMemory() + with pytest.raises(TypeError, match="takes 1 or 2 arguments"): + mkl.MKLMemory(32, 32, 32) + + +@pytest.mark.parametrize("arg", ["1024", 1024.0, None, 1024j, [1024], {}]) +def test_constructor_single_argument_type(arg): + with pytest.raises(TypeError, match="expects an integer or MKLMemory"): + mkl.MKLMemory(arg) + + +@pytest.mark.parametrize("arg", ["32", 32.0, None, 32j, [32]]) +def test_constructor_two_argument_types(arg): + with pytest.raises(TypeError, match="first argument"): + mkl.MKLMemory(arg, 32) + with pytest.raises(TypeError, match="second argument"): + mkl.MKLMemory(32, arg) + + +@pytest.mark.parametrize("nbytes", [0, -1]) +def test_malloc_rejects_non_positive_size(nbytes): + with pytest.raises(ValueError, match="must be positive"): + mkl.MKLMemory(nbytes) + + +@pytest.mark.parametrize( + "num,elem_size", [(0, 32), (32, 0), (0, 0), (-1, 32), (32, -1), (-1, -1)] +) +def test_calloc_rejects_non_positive_size(num, elem_size): + with pytest.raises(ValueError, match="must be positive"): + mkl.MKLMemory(num, elem_size) + + +def test_calloc_total_size_overflow_validation(): + with pytest.raises(ValueError, match="must not exceed"): + mkl.MKLMemory(2**32, 2**32) + with pytest.raises(ValueError, match="must not exceed"): + mkl.MKLMemory(sys.maxsize, 2) + with pytest.raises(ValueError, match="must not exceed"): + mkl.MKLMemory(2, sys.maxsize) + with pytest.raises(ValueError, match="must not exceed"): + mkl.MKLMemory(sys.maxsize // 2 + 1, 2) + + +@pytest.mark.parametrize( + "construct", + [ + lambda alignment: mkl.MKLMemory(1024, alignment=alignment), + lambda alignment: mkl.MKLMemory(32, 32, alignment=alignment), + lambda alignment: mkl.MKLMemory(mkl.MKLMemory(64), alignment=alignment), + ], + ids=["malloc", "calloc", "copy"], +) +def test_alignment_validation(construct): with pytest.raises(ValueError, match="positive"): - mkl.MKLMemory(1024, alignment=0) + construct(0) with pytest.raises(ValueError, match="positive"): - mkl.MKLMemory(1024, alignment=-1) + construct(-1) with pytest.raises(ValueError, match="must not exceed"): - mkl.MKLMemory(1024, alignment=2**40) + construct(2**40) + with pytest.raises(ValueError, match="must not exceed"): + construct(2**100) + + +@pytest.mark.parametrize("alignment", ["64", 64.0, None, 64j, [64]]) +def test_alignment_type_validation(alignment): + with pytest.raises(TypeError, match="must be an integer"): + mkl.MKLMemory(1024, alignment=alignment) + with pytest.raises(TypeError, match="must be an integer"): + mkl.MKLMemory(32, 32, alignment=alignment) + with pytest.raises(TypeError, match="must be an integer"): + mkl.MKLMemory(mkl.MKLMemory(64), alignment=alignment) + + +def test_unexpected_keyword_argument(): + keyword = "align" + match = f"unexpected keyword argument '{keyword}'" + with pytest.raises(TypeError, match=match): + mkl.MKLMemory(1024, **{keyword: 128}) + with pytest.raises(TypeError, match=match): + mkl.MKLMemory(32, 32, **{keyword: 128}) + with pytest.raises(TypeError, match=match): + mkl.MKLMemory(mkl.MKLMemory(64, alignment=128), **{keyword: 256}) + + +def test_alignment_keyword_still_accepted(): + assert mkl.MKLMemory(1024, alignment=128).alignment == 128 + assert mkl.MKLMemory(32, 32, alignment=128).alignment == 128 + source = mkl.MKLMemory(64, alignment=128) + assert mkl.MKLMemory(source).alignment == 128 + assert mkl.MKLMemory(source, alignment=256).alignment == 256 def test_concurrent_reads(): @@ -196,39 +388,43 @@ def reader(): assert not errors, f"Concurrent read errors: {errors}" +def _concurrent_realloc_round(initial, sizes): + mem = mkl.MKLMemory(initial) + barrier = threading.Barrier(len(sizes)) + results = [None] * len(sizes) + + def worker(idx, size): + barrier.wait() + try: + mem.realloc(size, refcheck=False) + results[idx] = "ok" + except BufferError: + results[idx] = "refused" + + ts = [ + threading.Thread(target=worker, args=(idx, size)) + for idx, size in enumerate(sizes) + ] + for t in ts: + t.start() + for t in ts: + t.join() + + return mem, results + + def test_concurrent_realloc_never_overlaps(): initial = 64 sizes = (1 << 16, 1 << 17) for _ in range(50): - mem = mkl.MKLMemory(initial) - barrier = threading.Barrier(len(sizes)) - results = [None] * len(sizes) - - def worker(idx, size, mem=mem, barrier=barrier, results=results): - barrier.wait() - try: - mem.realloc(size) - results[idx] = "ok" - except (ValueError, BufferError): - results[idx] = "refused" - - ts = [ - threading.Thread(target=worker, args=(idx, size)) - for idx, size in enumerate(sizes) - ] - for t in ts: - t.start() - for t in ts: - t.join() + mem, results = _concurrent_realloc_round(initial, sizes) assert all( r in ("ok", "refused") for r in results ), f"realloc raised an unexpected error: {results}" - allowed = {initial, *sizes} - assert ( - len(mem) in allowed - ), f"Inconsistent size {len(mem)} from {results}" + assert "ok" in results, f"no realloc completed: {results}" + assert len(mem) in sizes, f"Inconsistent size {len(mem)} from {results}" assert mem.nbytes == len(mem) assert len(mem.tobytes()) == len(mem) @@ -240,11 +436,63 @@ def worker(idx, size, mem=mem, barrier=barrier, results=results): mv.release() +@pytest.mark.skipif( + not REALLOC_RACE_IS_CONTAINED, + reason=( + "before 3.14 a free-threaded build cannot establish unique ownership, " + "so keeping readers off a resized allocation is the caller's job" + ), +) +def test_concurrent_realloc_and_reads(): + mem = mkl.MKLMemory(64) + stop = threading.Event() + errors = [] + + def reader(): + try: + while not stop.is_set(): + mv = memoryview(mem) + try: + n = mv.nbytes + assert n > 0 + # touch both ends of whatever block was handed out + mv[0] = 1 + mv[n - 1] = 2 + finally: + mv.release() + assert len(mem.tobytes()) == mem.nbytes + except Exception as e: # pragma: no cover - only on failure + errors.append(e) + + def reallocer(): + try: + for i in range(200): + try: + mem.realloc(1 << 12 if i % 2 == 0 else 1 << 13) + except (ValueError, BufferError): + pass + except Exception as e: # pragma: no cover - only on failure + errors.append(e) + finally: + stop.set() + + ts = [threading.Thread(target=reader) for _ in range(3)] + ts.append(threading.Thread(target=reallocer)) + for t in ts: + t.start() + for t in ts: + t.join() + + assert not errors, f"Concurrent realloc/read errors: {errors}" + assert mem.nbytes == len(mem) + + def test_realloc_refused_while_another_thread_holds_reference(): mem = mkl.MKLMemory(64) holder_ready = threading.Event() release_holder = threading.Event() outcome = [] + shared_refcount = [] def holder(): # keep reference alive @@ -252,10 +500,13 @@ def holder(): holder_ready.set() release_holder.wait(timeout=30) + base_refcount = sys.getrefcount(mem) + t = threading.Thread(target=holder) t.start() try: assert holder_ready.wait(timeout=30) + shared_refcount.append(sys.getrefcount(mem)) try: mem.realloc(1 << 16) outcome.append("ok") @@ -265,6 +516,10 @@ def holder(): release_holder.set() t.join() + assert shared_refcount[0] > base_refcount, ( + "Holder's reference was not visible here: " + f"{base_refcount} -> {shared_refcount[0]}" + ) assert outcome == [ "refused" ], f"Expected refusal while shared, got {outcome}" From a851c93ea6b6d2d5cf7d417006551b6bfb333bf5 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 22 Sep 2026 14:54:45 -0700 Subject: [PATCH 15/16] Apply review comments --- mkl/AGENTS.md | 8 +- mkl/_mkl_memory.pyx | 86 ++++--- mkl/tests/AGENTS.md | 17 +- mkl/tests/test_mkl_memory.py | 450 +++++++++++++++++++++++++---------- 4 files changed, 403 insertions(+), 158 deletions(-) diff --git a/mkl/AGENTS.md b/mkl/AGENTS.md index 8f832d7..3510ace 100644 --- a/mkl/AGENTS.md +++ b/mkl/AGENTS.md @@ -28,7 +28,7 @@ Core Python/Cython implementation: MKL support function wrappers and runtime con - `mem_stat()` — memory allocation statistics ### Memory allocation -- `MKLMemory(nbytes, alignment=64)` — aligned allocation via `mkl_malloc` +- `MKLMemory(nbytes, alignment=64)` — aligned allocation via `mkl_malloc`; `alignment` must be a power of two - `MKLMemory(num, elem_size, alignment=64)` — zeroed allocation via `mkl_calloc` - `MKLMemory(other, alignment=other.alignment)` — copy of another allocation - `realloc(new_nbytes, refcheck=True)` — resize in place via `mkl_realloc` @@ -48,6 +48,12 @@ Core Python/Cython implementation: MKL support function wrappers and runtime con - **MKL dependency:** Assumes MKL is available at runtime (conda: mkl package). Do **not** list `mkl` in `pyproject.toml` `[project].dependencies` — its PyPI wheel lacks `.dist-info`, which breaks `pip check`; on conda-forge there is no pip-visible `mkl` distribution. - **RTLD_GLOBAL preload path:** Linux preload is handled in `_mklinitmodule.c`; Windows DLL setup is in `_init_helper.py` - **`MKLMemory` mutation:** `realloc` moves the underlying block, so it must refuse while a buffer is exported, while another thread is resizing, or (unless `refcheck=False`) while the object looks referenced elsewhere. The GIL must not be released across those checks and the pointer store, mirroring NumPy's `PyArray_Resize`. The reference-count check stays NumPy's: `PyUnstable_Object_IsUniquelyReferenced` from 3.14, `Py_REFCNT > 2` before it, keyed on `PY_VERSION_HEX` and not on `Py_GIL_DISABLED`. It is a check against dangling references, not against other threads — on a free-threaded build before 3.14 it cannot be either, and resizing an allocation another thread can reach is the caller's responsibility, as it is for `numpy.ndarray.resize`. +- **`MKLMemory` alignment:** `mkl_malloc`/`mkl_calloc` honor only power-of-two alignments and silently fall back to their own (64 bytes, measured) for anything else, so `_check_alignment` rejects non-powers of two — otherwise `.alignment` would report a value the allocation does not have. Powers of two are delivered exactly, up to at least 1 GiB. +- **`MKLMemory` pickling:** `__reduce__` must rebuild `type(self)`, not `MKLMemory`, and carry the instance `__dict__` so a subclass survives a round trip. `_mkl_memory_from_bytes` takes the class as an optional third argument — optional so that older pickles still load, and omitted for `MKLMemory` itself so that its pickles stay loadable by older versions — and must reject anything that is not a `MKLMemory` subclass, since every pickle names that function. +- **`MKLMemory` buffer export:** `__getbuffer__` hands the view to `PyBuffer_FillInfo`, which describes a flat block of unsigned bytes and answers `flags` — `format` only under `PyBUF_FORMAT`, `shape` under `PyBUF_ND`, `strides` under `PyBUF_STRIDES` — instead of filling in fields the consumer did not request. It also takes the reference on the exporter, so `__releasebuffer__` must stay a bare decrement of `exported_buffers`. +- **Claim before fill:** the `atomic_fetch_add(&self.exported_buffers, 1)` comes *before* the fill, with the claim given back in an `except` clause if the fill raises (`PyBuffer_FillInfo` is declared `except -1`). Claiming afterwards leaves a window in which a concurrent `realloc` frees the block the view was already handed, and the consumer keeps that view — every array over the allocation holds it for as long as the array lives, so the cost is a durably dangling array rather than one bad read. Reproduced with the window widened by a 5 ms sleep on 3.13t: the resize went through and ASan reported `heap-use-after-free` in `array_tobytes`; with the claim first the same resize is refused. No test can observe the ordering, so it has to be kept on purpose. It narrows rather than closes the race — a `realloc` already past its own count check can still free under a fill — which only mutual exclusion would fix. +- **Backing a NumPy array:** `np.asarray(mem)` and `np.frombuffer(mem, dtype=...)` keep a `memoryview` as `.base` and hold the export for the array's whole lifetime, so `realloc` is refused with `BufferError` until the array goes away. `np.ndarray(shape, buffer=mem)` releases the `Py_buffer` and keeps only an object reference, so only the reference check stands in the way and `refcheck=False` leaves the array dangling — `bytearray` behaves the same there, so it is NumPy's property, not this object's. The array cannot resize the allocation either: it does not own its data, which `PyArray_Resize` refuses ahead of its own reference check, so neither `ndarray.resize(..., refcheck=False)` nor a C caller invoking `PyArray_Resize` directly gets past it (both measured). What does drop the export a live array depends on is `arr.base.release()`, which is caller error the same way it is for any exporter. +- **Reading another `MKLMemory`'s block:** code that reads someone else's allocation with the GIL released must claim a buffer on it (`atomic_fetch_add(&other.exported_buffers, 1)` in a `try`/`finally`, as the copy constructor does) *before* reading its size, so that a concurrent `realloc` is refused rather than freeing the block mid-read or shrinking it under a size that was already read. ## Cython details - `_py_mkl_service.pyx` → generates `_py_mkl_service` extension module diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index 0a152e2..530baff 100644 --- a/mkl/_mkl_memory.pyx +++ b/mkl/_mkl_memory.pyx @@ -30,6 +30,7 @@ import numbers from cpython cimport Py_buffer +from cpython.buffer cimport PyBuffer_FillInfo from libc.limits cimport INT_MAX from libc.string cimport memcpy @@ -88,6 +89,8 @@ cdef _extract_alignment(dict kwargs, object default): cdef int _check_alignment(object alignment) except -1: + cdef int c_alignment + if not isinstance(alignment, numbers.Integral): raise TypeError( "Alignment of requested allocation must be an integer, but got " @@ -99,16 +102,31 @@ cdef int _check_alignment(object alignment) except -1: raise ValueError( f"Alignment of requested allocation must not exceed {INT_MAX}." ) - return alignment + c_alignment = alignment + if c_alignment & (c_alignment - 1): + raise ValueError( + "Alignment of requested allocation must be a power of two, but got " + f"{c_alignment}." + ) + + return c_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) - cdef void *dst = mem._memory_ptr +def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment, cls=None): + cdef Py_ssize_t nbytes = len(data) + cdef MKLMemory mem + cdef void *dst cdef char *src = data + if cls is None: + cls = MKLMemory + elif not (isinstance(cls, type) and issubclass(cls, MKLMemory)): + raise TypeError(f"{cls} is not a subclass of MKLMemory") + + mem = cls(nbytes, alignment=alignment) + dst = mem._memory_ptr + with nogil: memcpy(dst, src, nbytes) @@ -142,9 +160,9 @@ cdef class MKLMemory: other (:class:`MKLMemory`): allocation whose size and content the new allocation takes. alignment (Optional[int]): - address alignment of the allocation in bytes. Expected to be - positive and to not exceed ``INT_MAX``. Defaults to the alignment - of ``other`` in the copy form, and to `64` otherwise. + address alignment of the allocation in bytes. Expected to be a + power of two and to not exceed ``INT_MAX``. Defaults to the + alignment of ``other`` in the copy form, and to `64` otherwise. """ cdef void *_memory_ptr cdef Py_ssize_t _nbytes @@ -220,9 +238,13 @@ cdef class MKLMemory: cdef _cinit_mklmemory(self, object other, object alignment): cdef MKLMemory other_mem = other - self._cinit_malloc(other_mem._nbytes, alignment) - with nogil: - memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes) + atomic_fetch_add(&other_mem.exported_buffers, 1) + try: + self._cinit_malloc(other_mem._nbytes, alignment) + with nogil: + memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes) + finally: + atomic_fetch_sub(&other_mem.exported_buffers, 1) def __cinit__(self, *args, **kwargs): n_args = len(args) @@ -271,19 +293,14 @@ cdef class MKLMemory: return self._memory_ptr def __getbuffer__(self, Py_buffer *buffer, int flags): - buffer.buf = self._memory_ptr - buffer.format = "B" - buffer.internal = NULL - buffer.itemsize = 1 - buffer.len = self._nbytes - buffer.ndim = 1 - buffer.obj = self - buffer.readonly = 0 - buffer.shape = &self._nbytes - buffer.strides = &buffer.itemsize - buffer.suboffsets = NULL - atomic_fetch_add(&self.exported_buffers, 1) + try: + PyBuffer_FillInfo( + buffer, self, self._memory_ptr, self._nbytes, 0, flags + ) + except BaseException: + atomic_fetch_sub(&self.exported_buffers, 1) + raise def __releasebuffer__(self, Py_buffer *buffer): atomic_fetch_sub(&self.exported_buffers, 1) @@ -320,13 +337,15 @@ cdef class MKLMemory: 3.14 there is neither, and a reference the caller holds cannot be told apart from one another thread holds: resizing an allocation another thread can reach may leave that thread reading freed memory whatever - ``refcheck`` is set to, so arrange for exclusive access. The same - applies to :meth:`numpy.ndarray.resize`. + ``refcheck`` is set to, so arrange for exclusive access. """ cdef void *p cdef int shared cdef int unclaimed = 0 + if new_nbytes <= 0: + raise ValueError("New number of bytes must be positive.") + # claim the exclusive right to reallocate before doing anything else if not atomic_compare_exchange_strong( &self.realloc_in_progress, &unclaimed, 1 @@ -355,11 +374,8 @@ cdef class MKLMemory: "risk of leaving those references pointing at freed " "memory." ) - if new_nbytes <= 0: - raise ValueError("New number of bytes must be positive.") - # do not release the GIL here, as that can allow another thread to - # read the or export a buffer with the old pointer before + # read from or export a buffer with the old pointer before # mkl_realloc frees it p = mkl_realloc(self._memory_ptr, new_nbytes) @@ -406,7 +422,15 @@ cdef class MKLMemory: return self._nbytes def __sizeof__(self): - return self._nbytes + return object.__sizeof__(self) + self._nbytes def __reduce__(self): - return (_mkl_memory_from_bytes, (self.tobytes(), self._alignment)) + cdef type cls = type(self) + + # a subclass should come back as itself + if cls is MKLMemory: + args = (self.tobytes(), self._alignment) + else: + args = (self.tobytes(), self._alignment, cls) + + return (_mkl_memory_from_bytes, args, getattr(self, "__dict__", None)) diff --git a/mkl/tests/AGENTS.md b/mkl/tests/AGENTS.md index 955e7f5..a1d7fe9 100644 --- a/mkl/tests/AGENTS.md +++ b/mkl/tests/AGENTS.md @@ -12,10 +12,12 @@ Unit tests for MKL runtime control API. - Memory: `peak_mem_usage`, `mem_stat` (if supported by MKL build) - CNR: Conditional Numerical Reproducibility flags - Edge cases currently covered: thread-local settings and API round-trips -- `MKLMemory` construction: all three forms, argument count/type errors, non-positive sizes, `num * elem_size` overflow, alignment bounds and types, unexpected keywords -- `MKLMemory` buffers: buffer protocol, `tobytes`, pickle round-trip, actual address alignment -- `MKLMemory.realloc`: grow/shrink with data preservation, alignment preserved across a resize, refusal while a buffer is exported or the object looks shared, `refcheck=False`, non-positive sizes -- `MKLMemory` concurrency: concurrent reads, overlapping `realloc` calls, readers racing a `realloc` +- `MKLMemory` construction: all three forms, argument count/type errors, non-positive sizes, `num * elem_size` overflow, alignment bounds and types, non-power-of-two alignments refused in all three forms, unexpected keywords, `mkl_calloc` actually zeroing, and the copy form copying the content into an allocation of its own +- `MKLMemory` buffers: two simultaneous views alias one block and each counts as an export of its own, the exported view's own fields (exporter, format, itemsize, ndim, shape, strides, suboffsets, writability, contiguity), no reference left behind per export/release cycle, `tobytes`, pickle round-trip, actual address alignment — every accepted alignment is checked against the delivered pointer, so a value MKL would ignore cannot pass unnoticed +- `MKLMemory` pickling: a subclass comes back as itself with its attributes, and the reconstructor refuses a class that is not a `MKLMemory` subclass +- `MKLMemory.realloc`: grow/shrink with data preservation, alignment preserved across a resize, refusal while a buffer is exported or the object looks shared, `refcheck=False`, non-positive sizes, and every refusal being a no-op — pointer, size, alignment and content unchanged, whatever the reason +- `MKLMemory` copy construction: the source's buffer is claimed for the duration, which is observed by resizing the source from an alignment object's `__int__` (called inside the copy), and released on both the success and the failure path +- `MKLMemory` concurrency: concurrent reads, two threads resizing at once (the CAS latch's losing side is only reachable on a free-threaded build — `realloc` holds the GIL otherwise), and readers hammering the object across resizes taken while they are parked, asserting the resizes were not quietly refused ## Running tests ```bash @@ -32,5 +34,8 @@ pytest mkl/tests/ - `MKLMemory` behavior → add to `test_mkl_memory.py` - Threading behavior → test thread count changes take effect - Use `mkl.get_version()` to check MKL availability before tests -- Concurrency tests must be checked for vacuity: a `realloc` refused by every thread satisfies loose assertions without ever reaching `mkl_realloc` -- Tests must pass on free-threaded builds, where `realloc`'s reference-count check does not guard against other threads before 3.14: a test that races a resize against live readers must be gated on `REALLOC_RACE_IS_CONTAINED`, or it reads freed memory there instead of testing a guard +- Concurrency tests must be checked for vacuity by counting outcomes, not by reading the code: a `realloc` refused by every thread satisfies loose assertions without ever reaching `mkl_realloc`. The predecessor of `test_concurrent_reads_across_reallocs` swallowed `(ValueError, BufferError)` around a `refcheck=True` resize and completed 0 of 200 resizes on every build where it ran — an object reachable from both the test frame and a closure cell has a reference count the check calls shared — and it still passed with `__releasebuffer__` gutted to a no-op +- Prefer a deterministic test over threads where the window can be entered on purpose: a callback from an argument the implementation converts inside the window (see `_AlignmentProbe`) tests the same guard without depending on the scheduler +- Tests must pass on free-threaded builds, and a test must not resize an allocation that other threads can reach — not on any version. `refcheck` is not a guard against other threads, and `tobytes` re-reads the pointer with nothing claimed, so a reallocer that retries until it slips between two reads is a use-after-free, not a test (ASan confirms it on 3.13t *and* 3.14t, in `tobytes`; a GIL build completes 200/200 resizes clean, which is why this looks fine locally). Park the readers on a `threading.Barrier` instead, resize while they are parked, and assert the resizes happened +- `numpy` is not a dependency of this package and CI does not install it for the test job, so any test that needs it must call `pytest.importorskip("numpy")` in the body — not import it at module level, which would break collection. Note that `pytest.importorskip` skips on `ModuleNotFoundError` only, so a probe that blocks the module by raising plain `ImportError` reports failures rather than skips and says nothing about the real behavior +- Content comparison alone does not prove a copy: a fresh allocation can be handed recycled heap memory that already holds the pattern, so `assert copy.tobytes() == source.tobytes()` has been observed to pass with the `memcpy` removed. Write the pattern a byte at a time so no freed `bytes` copy of it is left on the heap, and assert `_pointer` differs diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py index 3c869cd..781084b 100644 --- a/mkl/tests/test_mkl_memory.py +++ b/mkl/tests/test_mkl_memory.py @@ -23,6 +23,7 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import numbers import sys import threading @@ -30,11 +31,6 @@ import mkl -# on free-threaded Python prior to 3.14, only the caller can ensure exclusive -# access during realloc -_GIL_ENABLED = getattr(sys, "_is_gil_enabled", lambda: True)() -REALLOC_RACE_IS_CONTAINED = _GIL_ENABLED or sys.version_info >= (3, 14) - def test_mkl_memory_create_malloc(): nbytes = 1024 @@ -53,6 +49,8 @@ def test_mkl_memory_create_calloc(): assert mem.nbytes == nbytes # default alignment is 64 bytes assert mem.alignment == 64 + # mkl_calloc hands back zeroed memory + assert mem.tobytes() == bytes(nbytes) def test_mkl_memory_create_with_malloc_and_alignment(): @@ -75,18 +73,23 @@ def test_mkl_memory_create_with_calloc_and_alignment(): assert mem.alignment == alignment -@pytest.mark.parametrize("alignment", [64, 128, 256]) -def test_allocation_is_actually_aligned(alignment): - assert mkl.MKLMemory(1024, alignment=alignment)._pointer % alignment == 0 - assert mkl.MKLMemory(32, 32, alignment=alignment)._pointer % alignment == 0 - source = mkl.MKLMemory(1024, alignment=alignment) - assert mkl.MKLMemory(source)._pointer % alignment == 0 - - def test_mkl_memory_create_from_mkl_memory(): mem1 = mkl.MKLMemory(1024) + mv = memoryview(mem1) + for i in range(len(mem1)): + mv[i] = (i * 5 + 1) % 256 + mv.release() + mem2 = mkl.MKLMemory(mem1) assert mem2.nbytes == mem1.nbytes + assert mem2.tobytes() == mem1.tobytes() + assert mem2._pointer != mem1._pointer + + mv = memoryview(mem2) + mv[0] = mem1.tobytes()[0] ^ 0xFF + mv.release() + assert mem2.tobytes()[0] != mem1.tobytes()[0] + assert mem1.tobytes()[0] == (0 * 5 + 1) % 256 def test_mkl_memory_create_from_mkl_memory_with_alignment(): @@ -104,21 +107,142 @@ def test_mkl_memory_propagates_alignment(): assert mem2.alignment == mem1.alignment -def test_mkl_memory_properties(): - nbytes = 1024 - mem = mkl.MKLMemory(nbytes) - assert len(mem) == nbytes - assert type(repr(mem)) is str - assert type(bytes(mem)) is bytes - assert sys.getsizeof(mem) >= nbytes +class _AlignmentProbe: + def __init__(self, value, callback): + self._value = value + self._callback = callback + self._fired = False + + def __le__(self, other): + return self._value <= other + + def __gt__(self, other): + return self._value > other + + def _convert(self): + if not self._fired: + self._fired = True + self._callback() + return self._value + + def __int__(self): + return self._convert() + + def __index__(self): + return self._convert() + + +numbers.Integral.register(_AlignmentProbe) + + +def test_realloc_refused_while_copy_reads_source(): + source = mkl.MKLMemory(1024) + mv = memoryview(source) + for i in range(len(source)): + mv[i] = (i * 7 + 3) % 256 + mv.release() + pattern = source.tobytes() + + outcome = [] + + def probe(): + try: + source.realloc(256, refcheck=False) + outcome.append("resized") + except BufferError: + outcome.append("refused") + + copy = mkl.MKLMemory(source, alignment=_AlignmentProbe(64, probe)) + + assert outcome == ["refused"], f"resize was not refused: {outcome}" + assert source.nbytes == 1024 + assert copy.nbytes == 1024 + assert copy.alignment == 64 + assert copy.tobytes() == pattern + + +def test_copy_constructor_releases_source_claim(): + mem = mkl.MKLMemory(1024) + copy = mkl.MKLMemory(mem) + assert copy.nbytes == mem.nbytes + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_copy_constructor_releases_source_claim_on_failure(): + mem = mkl.MKLMemory(1024) + with pytest.raises(ValueError, match="Alignment"): + mkl.MKLMemory(mem, alignment=-1) + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_sizeof_accounts_for_object_too(): + small, large = 1024, 1 << 20 + mem, big = mkl.MKLMemory(small), mkl.MKLMemory(large) + + assert big.__sizeof__() - mem.__sizeof__() == large - small + overhead = mem.__sizeof__() - small + assert overhead > 0 + assert big.__sizeof__() - large == overhead + + assert sys.getsizeof(mem) >= mem.__sizeof__() + + mem.realloc(large) + assert mem.__sizeof__() == big.__sizeof__() def test_buffer_protocol(): mem = mkl.MKLMemory(1024) mv1 = memoryview(mem) - assert mv1.nbytes == mem.nbytes mv2 = memoryview(mem) - assert mv1 == mv2 + try: + assert mv1.nbytes == mem.nbytes + mv1[0] = 7 + assert mv2[0] == 7 + mv2[1] = 9 + assert mv1[1] == 9 + + with pytest.raises(BufferError, match="exported buffers"): + mem.realloc(2048, refcheck=False) + mv1.release() + with pytest.raises(BufferError, match="exported buffers"): + mem.realloc(2048, refcheck=False) + finally: + mv2.release() + + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_exported_buffer_describes_a_flat_writable_block(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + try: + assert mv.obj is mem + assert mv.format == "B" + assert mv.itemsize == 1 + assert mv.ndim == 1 + assert mv.shape == (mem.nbytes,) + assert mv.strides == (1,) + assert mv.suboffsets == () + assert not mv.readonly + assert mv.c_contiguous and mv.f_contiguous + mv[0] = 7 + assert mem.tobytes()[0] == 7 + finally: + mv.release() + + +def test_each_export_gives_back_its_reference(): + mem = mkl.MKLMemory(64) + before = sys.getrefcount(mem) + for _ in range(200): + memoryview(mem).release() + + assert sys.getrefcount(mem) == before + mem.realloc(128, refcheck=False) + assert mem.nbytes == 128 def test_pickling(): @@ -156,6 +280,47 @@ def test_pickling_with_alignment(): ), "Pickling should preserve alignment" +class _MKLMemorySubclass(mkl.MKLMemory): + pass + + +def test_pickling_preserves_subclass(): + import pickle + + mem = _MKLMemorySubclass(1024, alignment=128) + mv = memoryview(mem) + mv[:] = bytes((i * 3 + 1) % 256 for i in range(len(mem))) + mv.release() + + reconstructed = pickle.loads(pickle.dumps(mem)) + assert type(reconstructed) is _MKLMemorySubclass + assert reconstructed.nbytes == mem.nbytes + assert reconstructed.alignment == 128 + assert reconstructed.tobytes() == mem.tobytes() + + +def test_pickling_preserves_subclass_attributes(): + import pickle + + mem = _MKLMemorySubclass(256) + mem.label = "kept" + reconstructed = pickle.loads(pickle.dumps(mem)) + assert reconstructed.label == "kept" + + +def test_reconstruct_rejects_foreign_class(): + from mkl._mkl_memory import _mkl_memory_from_bytes + + with pytest.raises(TypeError, match="not a subclass of MKLMemory"): + _mkl_memory_from_bytes(b"abcd", 64, bytearray) + with pytest.raises(TypeError, match="not a subclass of MKLMemory"): + _mkl_memory_from_bytes(b"abcd", 64, "not a class") + + mem = _mkl_memory_from_bytes(b"abcd", 64) + assert type(mem) is mkl.MKLMemory + assert mem.tobytes() == b"abcd" + + def test_realloc_grow_and_shrink_preserves_data(): mem = mkl.MKLMemory(1024) mv = memoryview(mem) @@ -169,7 +334,6 @@ def test_realloc_grow_and_shrink_preserves_data(): assert mem.nbytes == grown assert len(mem) == grown assert len(mem.tobytes()) == grown - # growing keeps every byte that was there assert mem.tobytes()[:1024] == original mem.realloc(256) @@ -178,7 +342,6 @@ def test_realloc_grow_and_shrink_preserves_data(): # shrinking keeps the surviving prefix assert mem.tobytes() == original[:256] - # and the resized buffer is still writable through the buffer protocol mv = memoryview(mem) try: mv[0] = 7 @@ -191,8 +354,7 @@ def test_realloc_grow_and_shrink_preserves_data(): @pytest.mark.parametrize("alignment", [64, 128, 4096]) def test_realloc_preserves_alignment(alignment): - # test that alignment is preserved by realloc, which is undocumented in MKL - # but holds experimentally + # test that alignment is preserved by realloc mem = mkl.MKLMemory(1024, alignment=alignment) assert mem._pointer % alignment == 0 for nbytes in (1 << 20, 256): @@ -201,14 +363,6 @@ def test_realloc_preserves_alignment(alignment): assert mem._pointer % alignment == 0 -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 @@ -230,7 +384,7 @@ def test_realloc_refcheck_false_allows_shared(): mem.realloc(2048, refcheck=False) assert mem.nbytes == 2048 assert len(mem) == 2048 - # the leading bytes must have survived the move + # the leading bytes must be preserved assert mem.tobytes()[:256] == bytes(range(256)) del alias @@ -248,13 +402,26 @@ def test_realloc_refcheck_false_still_refuses_exported_buffer(): assert mem.nbytes == 2048 -def test_realloc_validates_size(): +@pytest.mark.parametrize("new_nbytes", [0, -1]) +def test_realloc_validates_size_before_state(new_nbytes): + match = "New number of bytes must be positive" mem = mkl.MKLMemory(1024) - with pytest.raises(ValueError, match="positive"): - mem.realloc(0, refcheck=False) - with pytest.raises(ValueError, match="positive"): - mem.realloc(-1, refcheck=False) + mv = memoryview(mem) + try: + with pytest.raises(ValueError, match=match): + mem.realloc(new_nbytes) + with pytest.raises(ValueError, match=match): + mem.realloc(new_nbytes, refcheck=False) + finally: + mv.release() + + held = mem + with pytest.raises(ValueError, match=match): + held.realloc(new_nbytes) + assert mem.nbytes == 1024 + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 def test_realloc_refcheck_is_keyword_only(): @@ -264,6 +431,49 @@ def test_realloc_refcheck_is_keyword_only(): assert mem.nbytes == 1024 +def test_failed_realloc_leaves_the_allocation_untouched(): + mem = mkl.MKLMemory(1024, alignment=128) + mv = memoryview(mem) + mv[:] = bytes((i * 11 + 5) % 256 for i in range(len(mem))) + mv.release() + + pointer, nbytes = mem._pointer, mem.nbytes + alignment, pattern = mem.alignment, mem.tobytes() + + def assert_untouched(): + assert mem._pointer == pointer + assert mem.nbytes == nbytes + assert len(mem) == nbytes + assert mem.alignment == alignment + assert mem.tobytes() == pattern + + with pytest.raises(ValueError, match="must be positive"): + mem.realloc(0, refcheck=False) + assert_untouched() + + mv = memoryview(mem) + try: + with pytest.raises(BufferError, match="exported buffers"): + mem.realloc(2048, refcheck=False) + finally: + mv.release() + assert_untouched() + + held = mem # noqa: F841 + with pytest.raises(ValueError, match="Cannot realloc MKLMemory"): + mem.realloc(2048) + assert_untouched() + + with pytest.raises(TypeError): + mem.realloc(2048, False) + assert_untouched() + + mem.realloc(4096, refcheck=False) + assert mem.nbytes == 4096 + assert mem.alignment == alignment + assert mem.tobytes()[:nbytes] == pattern + + def test_constructor_argument_count(): with pytest.raises(TypeError, match="takes 1 or 2 arguments"): mkl.MKLMemory() @@ -330,6 +540,34 @@ def test_alignment_validation(construct): construct(2**100) +@pytest.mark.parametrize("alignment", [3, 5, 12, 24, 96, 100, 129, 1000]) +@pytest.mark.parametrize( + "construct", + [ + lambda alignment: mkl.MKLMemory(1024, alignment=alignment), + lambda alignment: mkl.MKLMemory(32, 32, alignment=alignment), + lambda alignment: mkl.MKLMemory(mkl.MKLMemory(64), alignment=alignment), + ], + ids=["malloc", "calloc", "copy"], +) +def test_alignment_must_be_a_power_of_two(construct, alignment): + with pytest.raises(ValueError, match="must be a power of two"): + construct(alignment) + + +@pytest.mark.parametrize( + "alignment", [1, 2, 4, 8, 16, 32, 64, 128, 256, 4096, 1 << 20] +) +def test_powers_of_two_are_honored(alignment): + for mem in ( + mkl.MKLMemory(1024, alignment=alignment), + mkl.MKLMemory(32, 32, alignment=alignment), + mkl.MKLMemory(mkl.MKLMemory(64), alignment=alignment), + ): + assert mem.alignment == alignment + assert mem._pointer % alignment == 0 + + @pytest.mark.parametrize("alignment", ["64", 64.0, None, 64j, [64]]) def test_alignment_type_validation(alignment): with pytest.raises(TypeError, match="must be an integer"): @@ -351,14 +589,6 @@ def test_unexpected_keyword_argument(): mkl.MKLMemory(mkl.MKLMemory(64, alignment=128), **{keyword: 256}) -def test_alignment_keyword_still_accepted(): - assert mkl.MKLMemory(1024, alignment=128).alignment == 128 - assert mkl.MKLMemory(32, 32, alignment=128).alignment == 128 - source = mkl.MKLMemory(64, alignment=128) - assert mkl.MKLMemory(source).alignment == 128 - assert mkl.MKLMemory(source, alignment=256).alignment == 256 - - def test_concurrent_reads(): mem = mkl.MKLMemory(1024) mv = memoryview(mem) @@ -413,7 +643,7 @@ def worker(idx, size): return mem, results -def test_concurrent_realloc_never_overlaps(): +def test_concurrent_realloc_leaves_a_consistent_allocation(): initial = 64 sizes = (1 << 16, 1 << 17) @@ -436,91 +666,71 @@ def test_concurrent_realloc_never_overlaps(): mv.release() -@pytest.mark.skipif( - not REALLOC_RACE_IS_CONTAINED, - reason=( - "before 3.14 a free-threaded build cannot establish unique ownership, " - "so keeping readers off a resized allocation is the caller's job" - ), -) -def test_concurrent_realloc_and_reads(): - mem = mkl.MKLMemory(64) - stop = threading.Event() +def test_concurrent_reads_across_reallocs(): + sizes = [1 << 12, 1 << 13, 1 << 12, 1 << 14, 1 << 12] + n_readers = 3 + reads_per_round = 50 + mem = mkl.MKLMemory(sizes[0]) errors = [] + resizes = 0 + + def fill(value): + mv = memoryview(mem) + try: + mv[:] = bytes([value]) * len(mv) + finally: + mv.release() + + quiesce = threading.Barrier(n_readers + 1, timeout=60) + fill(0xA5) def reader(): try: - while not stop.is_set(): - mv = memoryview(mem) - try: - n = mv.nbytes - assert n > 0 - # touch both ends of whatever block was handed out - mv[0] = 1 - mv[n - 1] = 2 - finally: - mv.release() - assert len(mem.tobytes()) == mem.nbytes + for _ in sizes: + for _ in range(reads_per_round): + mv = memoryview(mem) + try: + n = mv.nbytes + assert n == mem.nbytes + data = bytes(mv) + finally: + mv.release() + assert data == data[:1] * n, "view spans two blocks" + + copy = mem.tobytes() + assert len(copy) == n + assert copy == data + quiesce.wait() # no reader is inside `mem` past this point + quiesce.wait() # the resize is done + except threading.BrokenBarrierError: # pragma: no cover - on failure + pass except Exception as e: # pragma: no cover - only on failure errors.append(e) + quiesce.abort() - def reallocer(): + def resizer(): + nonlocal resizes try: - for i in range(200): - try: - mem.realloc(1 << 12 if i % 2 == 0 else 1 << 13) - except (ValueError, BufferError): - pass + for round_ in range(len(sizes)): + quiesce.wait() + if round_ + 1 < len(sizes): + mem.realloc(sizes[round_ + 1], refcheck=False) + resizes += 1 + fill(round_ + 1) + quiesce.wait() + except threading.BrokenBarrierError: # pragma: no cover - on failure + pass except Exception as e: # pragma: no cover - only on failure errors.append(e) - finally: - stop.set() + quiesce.abort() - ts = [threading.Thread(target=reader) for _ in range(3)] - ts.append(threading.Thread(target=reallocer)) + ts = [threading.Thread(target=reader) for _ in range(n_readers)] + ts.append(threading.Thread(target=resizer)) for t in ts: t.start() for t in ts: t.join() assert not errors, f"Concurrent realloc/read errors: {errors}" - assert mem.nbytes == len(mem) - - -def test_realloc_refused_while_another_thread_holds_reference(): - mem = mkl.MKLMemory(64) - holder_ready = threading.Event() - release_holder = threading.Event() - outcome = [] - shared_refcount = [] - - def holder(): - # keep reference alive - alias = mem # noqa: F841 - holder_ready.set() - release_holder.wait(timeout=30) - - base_refcount = sys.getrefcount(mem) - - t = threading.Thread(target=holder) - t.start() - try: - assert holder_ready.wait(timeout=30) - shared_refcount.append(sys.getrefcount(mem)) - try: - mem.realloc(1 << 16) - outcome.append("ok") - except ValueError: - outcome.append("refused") - finally: - release_holder.set() - t.join() - - assert shared_refcount[0] > base_refcount, ( - "Holder's reference was not visible here: " - f"{base_refcount} -> {shared_refcount[0]}" - ) - assert outcome == [ - "refused" - ], f"Expected refusal while shared, got {outcome}" - assert len(mem) == 64, "Refused realloc must not change the buffer" + assert resizes == len(sizes) - 1 + assert mem.nbytes == sizes[-1] == len(mem) From 5898e7a195a1b42c711f663f28662bdaa9102d0d Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 22 Sep 2026 15:05:53 -0700 Subject: [PATCH 16/16] fix pre-commit --- mkl/tests/test_mkl_memory.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py index 781084b..e38bf91 100644 --- a/mkl/tests/test_mkl_memory.py +++ b/mkl/tests/test_mkl_memory.py @@ -309,6 +309,7 @@ def test_pickling_preserves_subclass_attributes(): def test_reconstruct_rejects_foreign_class(): + # pylint: disable-next=no-name-in-module from mkl._mkl_memory import _mkl_memory_from_bytes with pytest.raises(TypeError, match="not a subclass of MKLMemory"):