Add MKLMemory class to expose MKL allocated memory via Python buffer protocol

[ndgrigorian]

This PR proposes the introduction of _mkl_memory.pyx, which implements an MKLMemory class that exposes memory allocated via mkl_malloc and mkl_calloc to Python via the buffer protocol

The class uses an atomic counter incremented as __getbuffer__ and __releasebuffer__ are called to track the views on the buffer to permit use of mkl_realloc in the object (via realloc method). This concept was adapted from the PEP which revised the buffer protocol which proposed this kind of approach to tracking views on a buffer

Closes #18

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a new MKLMemory Cython extension type backed by MKL’s allocator, exposes it from the top-level mkl package, and introduces tests/build changes to support C11 atomics and nogil MKL calls.

Changes:

• Introduce mkl._mkl_memory with MKLMemory (allocation, buffer protocol, pickling, realloc).
• Add pytest coverage for allocation, buffer protocol, and pickling behavior.
• Update MKL C-API declarations/build to support nogil calls and C11 atomics (plus MSVC flag).

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
mkl/tests/test_mkl_memory.py	Adds tests for MKLMemory creation, buffer protocol, and pickling behavior.
mkl/_py_mkl_service.pyx	Releases the GIL around MKL buffer-free calls.
mkl/_mkl_service.pxd	Marks MKL externs as nogil and adds malloc/calloc/realloc/free declarations.
mkl/_mkl_memory.pyx	Adds the new MKLMemory Cython extension implementing allocation + buffer protocol + pickling.
mkl/init.py	Exposes MKLMemory at the package top level.
meson.build	Enables C11, adds MSVC atomics flag, and builds the new _mkl_memory extension.

[ndgrigorian]

@antonwolfy
Would be nice to get this reviewed and see if it can make it in the next release

[vlad-perevezentsev]

with nogil allows another thread to call realloc while memcpy is reading from the old buffer.
Could you temporarily increment exported_buffers during the copy similar to how it is done in __getbuffer__?

Like this I think

atomic_fetch_add(&other_mem.exported_buffers, 1)
try:
    with nogil:
         ...
finally:
    atomic_fetch_sub(&other_mem.exported_buffers, 1)

[vlad-perevezentsev]

It looks like oneMKL does not support non-power-of-two alignment values.

In [1]: import mkl

In [2]: A = 100

In [3]: ptrs = [mkl.MKLMemory(1024, alignment=A)._pointer for _ in range(200)]

In [4]: print(sum(p % A == 0 for p in ptrs), "/", len(ptrs))
0 / 200

In [5]: A = 64

In [6]: ptrs = [mkl.MKLMemory(1024, alignment=A)._pointer for _ in range(200)]

In [7]: print(sum(p % A == 0 for p in ptrs), "/", len(ptrs))
200 / 200

I think we should reject such values, document this restriction and extend tests

[vlad-perevezentsev]

__reduce__ always reconstructs MKLMemory so subclasses are lost after pickling.

In [8]: class Sub(mkl.MKLMemory): pass

In [9]: s = Sub(256, alignment=128)

In [10]: type(pickle.loads(pickle.dumps(s)))

In [11]: type(pickle.loads(pickle.dumps(s)))
Out[11]: mkl._mkl_memory.MKLMemory # not Sub

Could we preserve type(self) here, for example with cdef type cls = type(self) and pass it to _mkl_memory_from_bytes?

def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment, cls=None):
    cdef Py_ssize_t nbytes = len(data)
    cdef MKLMemory mem

    if cls is None:
        cls = MKLMemory
    elif not (isinstance(cls, type) and issubclass(cls, MKLMemory)):
        raise TypeError(f"{cls} is not a subclass of MKLMemory")

    mem = cls(nbytes, alignment=alignment)

[vlad-perevezentsev]

Do we need to consider the Python object overhead here (object.__ sizeof __(self))?

[vlad-perevezentsev]

new_nbytes is validated after the CAS checks so realloc(0) may raise BufferError instead of ValueError.

In [13]: mem = mkl.MKLMemory(1024)

In [14]: mv = memoryview(mem)

In [15]: mem.realloc(0)
---------------------------------------------------------------------------
BufferError                               Traceback (most recent call last)
Cell In[15], line 1
----> 1 mem.realloc(0)

File mkl/_mkl_memory.pyx:339, in mkl._mkl_memory.MKLMemory.realloc()
--> 339 'Could not get source, probably due dynamically evaluated source code.'

BufferError: Cannot realloc memory while there are exported buffers.

It would be better to validate the size first (above if not atomic_compare_exchange_strong)

[vlad-perevezentsev]

The test checks only nbytes. Could we also verify that the content is copied and that the two objects use different allocations?

[vlad-perevezentsev]

assert mem2.tobytes() == mem1.tobytes()
assert mem2._pointer != mem1._pointer

[vlad-perevezentsev]

It would be useful to add a test that a failed realloc leaves the original pointer, size and contents unchanged

[vlad-perevezentsev]

Suggested change

	# read the or export a buffer with the old pointer before
	# read from or export a buffer with the old pointer before