From 4d91687532844ef65d753fdd4efbb54083418bc3 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 19:56:57 -0700 Subject: [PATCH 01/12] 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 1b9ce89..95dfb83 100644 --- a/meson.build +++ b/meson.build @@ -53,7 +53,7 @@ py.extension_module( subdir: 'mkl' ) -# Cython extension +# Cython extensions py.extension_module( '_py_mkl_service', sources: ['mkl/_py_mkl_service.pyx'], @@ -64,6 +64,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 4cba157..a96afe6 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -148,6 +148,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 4defc15fb5b9784673c3e83af324a218cf864482 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 21:10:25 -0700 Subject: [PATCH 02/12] 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 95dfb83..3eda88b 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', ] ) @@ -19,6 +20,13 @@ c_args = ['-DNDEBUG'] 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 a96afe6..fc3c070 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -149,6 +149,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 e227fcfa864aee9462a8fd82769fa13fab34d850 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 21:18:53 -0700 Subject: [PATCH 03/12] 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 fc3c070..8edde28 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -150,6 +150,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 ee2636275e37d646726a928b6bbc89ac5d36e24e Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 22:10:59 -0700 Subject: [PATCH 04/12] 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 177f0d64171c2202c147737cbd063d7efa2f4662 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 23:39:23 -0700 Subject: [PATCH 05/12] 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 2807de0cfb614421c6f457dcd344cec1dde0adca Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 11 Apr 2026 23:53:17 -0700 Subject: [PATCH 06/12] 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 72e4deb725ad413bd5f47cc4eeb04c8cc7382ddf Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sun, 12 Apr 2026 00:19:16 -0700 Subject: [PATCH 07/12] 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 c764a8136410a3df6de87191e8f6552ad2833b49 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sun, 12 Apr 2026 03:07:07 -0700 Subject: [PATCH 08/12] 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 8edde28..3e07630 100644 --- a/mkl/_mkl_service.pxd +++ b/mkl/_mkl_service.pxd @@ -29,7 +29,7 @@ ctypedef unsigned long long MKL_UINT64 ctypedef int MKL_INT -cdef extern from "mkl.h": +cdef extern from "mkl.h" nogil: # MKL Function Domains Constants int MKL_DOMAIN_BLAS int MKL_DOMAIN_FFT @@ -148,10 +148,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 8cfc0450bf4bf7cc61756bac617659d47e188efe Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 09:39:44 -0700 Subject: [PATCH 09/12] fix meson.build rpath --- meson.build | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/meson.build b/meson.build index 3eda88b..c821583 100644 --- a/meson.build +++ b/meson.build @@ -77,7 +77,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 d8e3d326e9c416c89e896a5981c27b5a99b7ee31 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 09:40:02 -0700 Subject: [PATCH 10/12] 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 08897b56e871c6bf45d40b2748d92905915a29b2 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 11:23:41 -0700 Subject: [PATCH 11/12] update meson.build --- meson.build | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/meson.build b/meson.build index c821583..294f531 100644 --- a/meson.build +++ b/meson.build @@ -94,6 +94,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 0d4ccfb7b0f6aca54e87f82031f228ec86828459 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 6 Aug 2026 11:24:17 -0700 Subject: [PATCH 12/12] align realloc behavior to NumPy also address issues with undeclared variables and rename MKLMemory class members --- mkl/_mkl_memory.pyx | 107 ++++++++++++++++++++++++++--------- mkl/tests/test_mkl_memory.py | 84 +++++++++++++++++++++++++++ 2 files changed, 164 insertions(+), 27 deletions(-) diff --git a/mkl/_mkl_memory.pyx b/mkl/_mkl_memory.pyx index d4c8033..b46a8bc 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 @@ -43,6 +44,37 @@ cdef extern from "stdatomic.h" nogil: int atomic_load(atomic_int *obj) +cdef extern from *: + """ + // Check whether a MKLMemory object may be safely reallocated. + // Mirrors NumPy's PyArray_Resize_int logic. + static int _MKLMemory_MayBeShared(PyObject *op) { + #if PY_VERSION_HEX >= 0x030e00b0 + if (PyUnstable_Object_IsUniquelyReferenced(op)) { + return 0; // not shared + } + if (Py_REFCNT(op) == 2) { + return 1; // may be shared + } + return 2; // definitely shared + #else + return (Py_REFCNT(op) > 2) ? 2 : 0; + #endif + } + """ + int _MKLMemory_MayBeShared(object obj) + + +cdef int _check_alignment(Py_ssize_t alignment) except -1: + if alignment <= 0: + raise ValueError("Alignment of requested allocation must be positive.") + if alignment > 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): cdef Py_ssize_t nbytes = len(data) cdef MKLMemory mem = MKLMemory(nbytes, alignment=alignment) @@ -57,28 +89,32 @@ def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment): cdef class MKLMemory: + """MKL-backed memory object that exposes Python buffer protocol.""" cdef void *_memory_ptr - cdef Py_ssize_t nbytes - cdef Py_ssize_t alignment + cdef Py_ssize_t _nbytes + cdef Py_ssize_t _alignment cdef atomic_int exported_buffers cdef _cinit_empty(self): self._memory_ptr = NULL - self.nbytes = 0 - self.alignment = 0 + self._nbytes = 0 + self._alignment = 0 atomic_init(&self.exported_buffers, 0) cdef _cinit_malloc(self, Py_ssize_t nbytes, Py_ssize_t alignment): + cdef int c_alignment = _check_alignment(alignment) + cdef void *p + self._cinit_empty() if (nbytes > 0): with nogil: - p = mkl_malloc(nbytes, alignment) + p = mkl_malloc(nbytes, c_alignment) if (p): self._memory_ptr = p - self.nbytes = nbytes - self.alignment = alignment + self._nbytes = nbytes + self._alignment = alignment else: raise MemoryError( "MKL memory allocation failed." @@ -89,16 +125,19 @@ cdef class MKLMemory: ) cdef _cinit_calloc(self, Py_ssize_t num, Py_ssize_t size, Py_ssize_t alignment): + cdef int c_alignment = _check_alignment(alignment) + cdef void *p + self._cinit_empty() if (num > 0 and size > 0): with nogil: - p = mkl_calloc(num, size, alignment) + p = mkl_calloc(num, size, c_alignment) if (p): self._memory_ptr = p - self.nbytes = num * size - self.alignment = alignment + self._nbytes = num * size + self._alignment = alignment else: raise MemoryError( "MKL memory allocation failed." @@ -110,11 +149,11 @@ cdef class MKLMemory: ) cdef _cinit_mklmemory(self, object other, Py_ssize_t alignment): - other_mem = 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 +170,7 @@ cdef class MKLMemory: alignment = kwargs.get("alignment", 64) self._cinit_malloc(arg, alignment) elif isinstance(arg, MKLMemory): - alignment = kwargs.get("alignment", arg.alignment) + alignment = kwargs.get("alignment", (arg)._alignment) self._cinit_mklmemory(arg, alignment) else: raise TypeError( @@ -167,11 +206,11 @@ cdef class MKLMemory: buffer.format = "B" buffer.internal = NULL buffer.itemsize = 1 - buffer.len = self.nbytes + buffer.len = self._nbytes buffer.ndim = 1 buffer.obj = self buffer.readonly = 0 - buffer.shape = &self.nbytes + buffer.shape = &self._nbytes buffer.strides = &buffer.itemsize buffer.suboffsets = NULL @@ -181,12 +220,26 @@ cdef class MKLMemory: atomic_fetch_sub(&self.exported_buffers, 1) def realloc(self, Py_ssize_t new_nbytes): + cdef void *p + cdef int shared + if atomic_load(&self.exported_buffers) > 0: - raise BufferError("Cannot realloc memory while there are exported buffers.") + raise BufferError( + "Cannot realloc memory while there are exported buffers." + ) + shared = _MKLMemory_MayBeShared(self) + if shared == 1: + raise ValueError( + "Cannot realloc MKLMemory that may be referenced by another " + "object. It is possible that this is a false positive." + ) + elif shared == 2: + raise ValueError( + "Cannot realloc MKLMemory that is referenced by other objects." + ) if new_nbytes <= 0: raise ValueError("New number of bytes must be positive.") - cdef void *p with nogil: p = mkl_realloc(self._memory_ptr, new_nbytes) @@ -194,23 +247,23 @@ cdef class MKLMemory: raise MemoryError("MKL memory reallocation failed.") self._memory_ptr = p - self.nbytes = new_nbytes + self._nbytes = new_nbytes def tobytes(self): cdef char* data_ptr = 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 +271,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..3ca517b 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,84 @@ def test_pickling_with_alignment(): assert ( mem.alignment == mem_reconstructed.alignment ), "Pickling should preserve alignment" + + +def test_realloc_exported_buffer(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + with pytest.raises(BufferError): + mem.realloc(2048) + del mv + + +def test_realloc_refcheck_shared(): + mem = mkl.MKLMemory(1024) + alias = mem # noqa: F841 — extra reference + with pytest.raises(ValueError, match="referenced by"): + mem.realloc(2048) + del alias + + +def test_alignment_validation(): + with pytest.raises(ValueError, match="positive"): + mkl.MKLMemory(1024, alignment=0) + with pytest.raises(ValueError, match="positive"): + mkl.MKLMemory(1024, alignment=-1) + with pytest.raises(ValueError, match="must not exceed"): + mkl.MKLMemory(1024, alignment=2**40) + + +def test_concurrent_reads(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + for i in range(len(mem)): + mv[i] = i % 256 + del mv + + errors = [] + + def reader(): + try: + for _ in range(500): + assert len(mem) == 1024 + data = mem.tobytes() + assert len(data) == 1024 + v = memoryview(mem) + assert v[0] == 0 + v.release() + except Exception as e: + errors.append(e) + + ts = [threading.Thread(target=reader) for _ in range(4)] + for t in ts: + t.start() + for t in ts: + t.join() + assert not errors, f"Concurrent read errors: {errors}" + + +def test_concurrent_realloc_refused(): + for _ in range(50): + mem = mkl.MKLMemory(64) + barrier = threading.Barrier(2) + results = [None, None] + + def worker(idx, size): + barrier.wait() + try: + mem.realloc(size) + results[idx] = "ok" + except ValueError: + results[idx] = "refused" + + ts = [ + threading.Thread(target=worker, args=(0, 1 << 16)), + threading.Thread(target=worker, args=(1, 1 << 17)), + ] + for t in ts: + t.start() + for t in ts: + t.join() + assert results[0] == "refused" and results[1] == "refused", ( + f"Expected both threads refused, got {results}" + )