Add MKLMemory class to expose MKL allocated memory via Python buffer protocol - #182
Add MKLMemory class to expose MKL allocated memory via Python buffer protocol#182ndgrigorian wants to merge 12 commits into
MKLMemory class to expose MKL allocated memory via Python buffer protocol#182Conversation
aed1ff6 to
a99e626
Compare
ef09add to
ac6e9fc
Compare
b2cc6fc to
00fea26
Compare
ac6e9fc to
b39b64b
Compare
00fea26 to
c764a81
Compare
There was a problem hiding this comment.
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_memorywithMKLMemory(allocation, buffer protocol, pickling, realloc). - Add pytest coverage for allocation, buffer protocol, and pickling behavior.
- Update MKL C-API declarations/build to support
nogilcalls 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. |
exposes Python buffer protocol
as atomics are a c11+ feature, specific flags are needed to enable on window
0d4ccfb to
0c75d30
Compare
also address issues with undeclared variables and rename MKLMemory class members
0c75d30 to
6aae4fb
Compare
|
@antonwolfy |
| @@ -0,0 +1,271 @@ | |||
| # Copyright (c) 2018, Intel Corporation | |||
There was a problem hiding this comment.
Should be a year when file is created
| # Copyright (c) 2018, Intel Corporation | |
| # Copyright (c) 2026, Intel Corporation |
| subdir: 'mkl' | ||
| ) | ||
|
|
||
| py.extension_module( |
There was a problem hiding this comment.
Please populate the changelog
| if new_nbytes <= 0: | ||
| raise ValueError("New number of bytes must be positive.") | ||
|
|
||
| with nogil: |
There was a problem hiding this comment.
There might be use-after-free issue:
- Thread B (holding GIL): passes exported_buffers == 0 and the refcount check.
- Thread B enters with nogil: → releases the GIL → calls mkl_realloc.
- Thread A acquires the GIL, runs
memoryview(mem)→__getbuffer__reads the old_memory_ptrintobuffer.buf, bumps the counter. - Thread B's
mkl_reallocfrees/moves the old block, reacquires GIL, sets_memory_ptr = p. - Thread A now holds a view into freed memory → UAF.
|
|
||
| if (p): | ||
| self._memory_ptr = p | ||
| self._nbytes = num * size |
There was a problem hiding this comment.
Needs to validate there is no num > PY_SSIZE_T_MAX / size
| with nogil: | ||
| memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes) | ||
|
|
||
| def __cinit__(self, *args, **kwargs): |
There was a problem hiding this comment.
No check on name of alignment keyword.
Any typo in name, like MKLMemory(1024, alignmnet=128) (instead of alignment) will be silently ignored and fallback on 64.
There was a problem hiding this comment.
Also no type-check before the Py_ssize_t conversion
| return self._nbytes | ||
|
|
||
| @property | ||
| def size(self): |
There was a problem hiding this comment.
That duplicates nbytes.
Also that collides with the constructor's .size parameter, which means something different.
And that collides with the NumPy convention where .size is the element count and .nbytes is the byte count.
| cc = meson.get_compiler('c') | ||
| if cc.get_id() == 'msvc' | ||
| add_project_arguments( | ||
| '/experimental:c11atomics', |
There was a problem hiding this comment.
/experimental:c11atomics requires VS 2022 17.5+ and fails hard on older toolchains.
Should we document the floor?
| def __repr__(self): | ||
| return ( | ||
| f"<MKL memory allocation of {self._nbytes} bytes at " | ||
| f"{hex(<object>(<size_t>self._memory_ptr))}>" |
There was a problem hiding this comment.
It can be simplified:
| f"{hex(<object>(<size_t>self._memory_ptr))}>" | |
| f"{hex(self._pointer)}>" |
There was a problem hiding this comment.
Both forms cythonized to the same:
with_cast — hex(<object>(<size_t>x)):
__pyx_t_1 = __Pyx_PyLong_FromSize_t(((size_t)__pyx_v_x)); ...
__pyx_t_2 = __Pyx_PyNumber_Hex(__pyx_t_1); ...
without_cast — hex(<size_t>x):
__pyx_t_1 = __Pyx_PyLong_FromSize_t(((size_t)__pyx_v_x)); ...
__pyx_t_2 = __Pyx_PyNumber_Hex(__pyx_t_1); ...| @@ -0,0 +1,271 @@ | |||
| # Copyright (c) 2018, Intel Corporation | |||
There was a problem hiding this comment.
Should we update AGENTS.md, .github/copilot-instructions.md and mkl/tests/AGENTS.md with the new test file?
There was a problem hiding this comment.
The the same for _py_mkl_service files
|
|
||
| cc = meson.get_compiler('c') | ||
| if cc.get_id() == 'msvc' | ||
| add_project_arguments( |
There was a problem hiding this comment.
That is only needed for _py_mkl_service. We probably should not add it everywhere, considering it's experimental:
mkl_memory_c_args = c_args
if cc.get_id() == 'msvc'
mkl_memory_c_args += '/experimental:c11atomics'
endif
This PR proposes the introduction of
_mkl_memory.pyx, which implements anMKLMemoryclass that exposes memory allocated viamkl_mallocandmkl_callocto Python via the buffer protocolThe class uses an atomic counter incremented as
__getbuffer__and__releasebuffer__are called to track the views on the buffer to permit use ofmkl_reallocin the object (viareallocmethod). This concept was adapted from the PEP which revised the buffer protocol which proposed this kind of approach to tracking views on a bufferCloses #18