From 7bd5c95b7cb741fb2158e11c0379484ad98a96d8 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 19:56:57 -0700 Subject: [PATCH 01/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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}"