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 285aca4..836897e 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 @@ -74,11 +76,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 74409b6..cc21f3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Enabled support of Python 3.15 [gh-243](https://github.com/IntelPython/mkl-service/pull/243) * 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 0972a6b..b29ca79 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,20 @@ endif thread_dep = dependency('threads') cc = meson.get_compiler('c') + +atomics_args = [] +if cc.get_id() == 'msvc' + 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 + mkl_dep = dependency('MKL', method: 'cmake', modules: ['MKL::MKL'], cmake_args: [ @@ -60,7 +75,7 @@ py.extension_module( subdir: 'mkl' ) -# Cython extension +# Cython extensions py.extension_module( '_py_mkl_service', sources: ['mkl/_py_mkl_service.pyx'], @@ -71,6 +86,17 @@ py.extension_module( subdir: 'mkl' ) +py.extension_module( + '_mkl_memory', + sources: ['mkl/_mkl_memory.pyx'], + dependencies: [mkl_dep], + c_args: c_args + atomics_args, + link_args: rpath_link_args, + install: true, + subdir: 'mkl' +) + + # Python sources py.install_sources( [ @@ -82,6 +108,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' ) 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 c0eb2ae..d1ec7c2 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, @@ -121,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 new file mode 100644 index 0000000..0a152e2 --- /dev/null +++ b/mkl/_mkl_memory.pyx @@ -0,0 +1,412 @@ +# 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: +# +# * 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 +# cython: freethreading_compatible=True + +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 + + +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) + 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 _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: + 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) + + cdef void *dst = mem._memory_ptr + cdef char *src = data + + with nogil: + memcpy(dst, src, nbytes) + + return mem + + +cdef class MKLMemory: + """ + 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 + 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 + atomic_init(&self.exported_buffers, 0) + atomic_init(&self.realloc_in_progress, 0) + + cdef _cinit_malloc(self, Py_ssize_t nbytes, object alignment): + cdef int c_alignment = _check_alignment(alignment) + cdef void *p + + self._cinit_empty() + + if (nbytes > 0): + with nogil: + p = mkl_malloc(nbytes, c_alignment) + + if (p): + self._memory_ptr = p + self._nbytes = nbytes + self._alignment = c_alignment + else: + raise MemoryError( + "MKL memory allocation failed." + ) + else: + raise ValueError( + "Number of bytes of requested allocation must be positive." + ) + + 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 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, elem_size, c_alignment) + + if (p): + self._memory_ptr = p + self._nbytes = nbytes + self._alignment = c_alignment + else: + raise MemoryError( + "MKL memory allocation failed." + ) + else: + raise ValueError( + "Number of elements and element size of requested allocation " + "must be positive." + ) + + cdef _cinit_mklmemory(self, object other, object alignment): + cdef MKLMemory other_mem = other + + self._cinit_malloc(other_mem._nbytes, alignment) + with nogil: + memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes) + + def __cinit__(self, *args, **kwargs): + 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): + alignment = _extract_alignment(kwargs, 64) + self._cinit_malloc(arg, alignment) + elif isinstance(arg, MKLMemory): + alignment = _extract_alignment( + kwargs, (arg)._alignment + ) + 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] + alignment = _extract_alignment(kwargs, 64) + 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): + 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" + buffer.internal = NULL + buffer.itemsize = 1 + buffer.len = self._nbytes + buffer.ndim = 1 + buffer.obj = self + buffer.readonly = 0 + buffer.shape = &self._nbytes + buffer.strides = &buffer.itemsize + buffer.suboffsets = NULL + + atomic_fetch_add(&self.exported_buffers, 1) + + def __releasebuffer__(self, Py_buffer *buffer): + atomic_fetch_sub(&self.exported_buffers, 1) + + 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 + + # 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." + ) + 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.") + + # 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.") + + self._memory_ptr = p + self._nbytes = new_nbytes + finally: + 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): + """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"" + ) + + def __len__(self): + return self._nbytes + + def __sizeof__(self): + return self._nbytes + + def __reduce__(self): + return (_mkl_memory_from_bytes, (self.tobytes(), self._alignment)) diff --git a/mkl/_mkl_service.pxd b/mkl/_mkl_service.pxd index ed5a106..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,6 +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) + 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 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 new file mode 100644 index 0000000..3c869cd --- /dev/null +++ b/mkl/tests/test_mkl_memory.py @@ -0,0 +1,526 @@ +# 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: +# +# * 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 threading + +import pytest + +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 + 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 + assert mem.alignment == alignment + + +@pytest.mark.parametrize("alignment", [64, 128, 256]) +def test_allocation_is_actually_aligned(alignment): + assert mkl.MKLMemory(1024, alignment=alignment)._pointer % alignment == 0 + assert mkl.MKLMemory(32, 32, alignment=alignment)._pointer % alignment == 0 + source = mkl.MKLMemory(1024, alignment=alignment) + assert mkl.MKLMemory(source)._pointer % alignment == 0 + + +def test_mkl_memory_create_from_mkl_memory(): + mem1 = mkl.MKLMemory(1024) + 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" + + +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) + with pytest.raises(BufferError): + mem.realloc(2048) + del mv + + +def test_realloc_refcheck_shared(): + mem = mkl.MKLMemory(1024) + alias = mem # noqa: F841 + with pytest.raises(ValueError, match="referenced by"): + mem.realloc(2048) + del alias + + +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"): + construct(0) + with pytest.raises(ValueError, match="positive"): + construct(-1) + with pytest.raises(ValueError, match="must not exceed"): + 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(): + 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 _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, results = _concurrent_realloc_round(initial, sizes) + + assert all( + r in ("ok", "refused") for r in results + ), f"realloc raised an unexpected error: {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) + + mv = memoryview(mem) + try: + mv[0] = 1 + mv[len(mem) - 1] = 2 + finally: + 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 + alias = mem # noqa: F841 + holder_ready.set() + release_holder.wait(timeout=30) + + base_refcount = sys.getrefcount(mem) + + t = threading.Thread(target=holder) + t.start() + try: + assert holder_ready.wait(timeout=30) + shared_refcount.append(sys.getrefcount(mem)) + try: + mem.realloc(1 << 16) + outcome.append("ok") + except ValueError: + outcome.append("refused") + finally: + release_holder.set() + t.join() + + assert shared_refcount[0] > base_refcount, ( + "Holder's reference was not visible here: " + f"{base_refcount} -> {shared_refcount[0]}" + ) + assert outcome == [ + "refused" + ], f"Expected refusal while shared, got {outcome}" + assert len(mem) == 64, "Refused realloc must not change the buffer"