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..3510ace 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`; `alignment` must be a power of two +- `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,20 @@ 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`. +- **`MKLMemory` alignment:** `mkl_malloc`/`mkl_calloc` honor only power-of-two alignments and silently fall back to their own (64 bytes, measured) for anything else, so `_check_alignment` rejects non-powers of two — otherwise `.alignment` would report a value the allocation does not have. Powers of two are delivered exactly, up to at least 1 GiB. +- **`MKLMemory` pickling:** `__reduce__` must rebuild `type(self)`, not `MKLMemory`, and carry the instance `__dict__` so a subclass survives a round trip. `_mkl_memory_from_bytes` takes the class as an optional third argument — optional so that older pickles still load, and omitted for `MKLMemory` itself so that its pickles stay loadable by older versions — and must reject anything that is not a `MKLMemory` subclass, since every pickle names that function. +- **`MKLMemory` buffer export:** `__getbuffer__` hands the view to `PyBuffer_FillInfo`, which describes a flat block of unsigned bytes and answers `flags` — `format` only under `PyBUF_FORMAT`, `shape` under `PyBUF_ND`, `strides` under `PyBUF_STRIDES` — instead of filling in fields the consumer did not request. It also takes the reference on the exporter, so `__releasebuffer__` must stay a bare decrement of `exported_buffers`. +- **Claim before fill:** the `atomic_fetch_add(&self.exported_buffers, 1)` comes *before* the fill, with the claim given back in an `except` clause if the fill raises (`PyBuffer_FillInfo` is declared `except -1`). Claiming afterwards leaves a window in which a concurrent `realloc` frees the block the view was already handed, and the consumer keeps that view — every array over the allocation holds it for as long as the array lives, so the cost is a durably dangling array rather than one bad read. Reproduced with the window widened by a 5 ms sleep on 3.13t: the resize went through and ASan reported `heap-use-after-free` in `array_tobytes`; with the claim first the same resize is refused. No test can observe the ordering, so it has to be kept on purpose. It narrows rather than closes the race — a `realloc` already past its own count check can still free under a fill — which only mutual exclusion would fix. +- **Backing a NumPy array:** `np.asarray(mem)` and `np.frombuffer(mem, dtype=...)` keep a `memoryview` as `.base` and hold the export for the array's whole lifetime, so `realloc` is refused with `BufferError` until the array goes away. `np.ndarray(shape, buffer=mem)` releases the `Py_buffer` and keeps only an object reference, so only the reference check stands in the way and `refcheck=False` leaves the array dangling — `bytearray` behaves the same there, so it is NumPy's property, not this object's. The array cannot resize the allocation either: it does not own its data, which `PyArray_Resize` refuses ahead of its own reference check, so neither `ndarray.resize(..., refcheck=False)` nor a C caller invoking `PyArray_Resize` directly gets past it (both measured). What does drop the export a live array depends on is `arr.base.release()`, which is caller error the same way it is for any exporter. +- **Reading another `MKLMemory`'s block:** code that reads someone else's allocation with the GIL released must claim a buffer on it (`atomic_fetch_add(&other.exported_buffers, 1)` in a `try`/`finally`, as the copy constructor does) *before* reading its size, so that a concurrent `realloc` is refused rather than freeing the block mid-read or shrinking it under a size that was already read. ## 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..530baff --- /dev/null +++ b/mkl/_mkl_memory.pyx @@ -0,0 +1,436 @@ +# 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 cpython.buffer cimport PyBuffer_FillInfo +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: + cdef int c_alignment + + 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}." + ) + + c_alignment = alignment + if c_alignment & (c_alignment - 1): + raise ValueError( + "Alignment of requested allocation must be a power of two, but got " + f"{c_alignment}." + ) + + return c_alignment + + +def _mkl_memory_from_bytes(bytes data, Py_ssize_t alignment, cls=None): + cdef Py_ssize_t nbytes = len(data) + cdef MKLMemory mem + cdef void *dst + cdef char *src = data + + 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) + dst = mem._memory_ptr + + 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 a + power of two 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 + + atomic_fetch_add(&other_mem.exported_buffers, 1) + try: + self._cinit_malloc(other_mem._nbytes, alignment) + with nogil: + memcpy(self._memory_ptr, other_mem._memory_ptr, self._nbytes) + finally: + atomic_fetch_sub(&other_mem.exported_buffers, 1) + + 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): + atomic_fetch_add(&self.exported_buffers, 1) + try: + PyBuffer_FillInfo( + buffer, self, self._memory_ptr, self._nbytes, 0, flags + ) + except BaseException: + atomic_fetch_sub(&self.exported_buffers, 1) + raise + + 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. + """ + cdef void *p + cdef int shared + cdef int unclaimed = 0 + + if new_nbytes <= 0: + raise ValueError("New number of bytes must be positive.") + + # 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." + ) + # do not release the GIL here, as that can allow another thread to + # read from 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 object.__sizeof__(self) + self._nbytes + + def __reduce__(self): + cdef type cls = type(self) + + # a subclass should come back as itself + if cls is MKLMemory: + args = (self.tobytes(), self._alignment) + else: + args = (self.tobytes(), self._alignment, cls) + + return (_mkl_memory_from_bytes, args, getattr(self, "__dict__", None)) 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..a1d7fe9 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,12 @@ 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, non-power-of-two alignments refused in all three forms, unexpected keywords, `mkl_calloc` actually zeroing, and the copy form copying the content into an allocation of its own +- `MKLMemory` buffers: two simultaneous views alias one block and each counts as an export of its own, the exported view's own fields (exporter, format, itemsize, ndim, shape, strides, suboffsets, writability, contiguity), no reference left behind per export/release cycle, `tobytes`, pickle round-trip, actual address alignment — every accepted alignment is checked against the delivered pointer, so a value MKL would ignore cannot pass unnoticed +- `MKLMemory` pickling: a subclass comes back as itself with its attributes, and the reconstructor refuses a class that is not a `MKLMemory` subclass +- `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, and every refusal being a no-op — pointer, size, alignment and content unchanged, whatever the reason +- `MKLMemory` copy construction: the source's buffer is claimed for the duration, which is observed by resizing the source from an alignment object's `__int__` (called inside the copy), and released on both the success and the failure path +- `MKLMemory` concurrency: concurrent reads, two threads resizing at once (the CAS latch's losing side is only reachable on a free-threaded build — `realloc` holds the GIL otherwise), and readers hammering the object across resizes taken while they are parked, asserting the resizes were not quietly refused ## Running tests ```bash @@ -24,5 +31,11 @@ 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 by counting outcomes, not by reading the code: a `realloc` refused by every thread satisfies loose assertions without ever reaching `mkl_realloc`. The predecessor of `test_concurrent_reads_across_reallocs` swallowed `(ValueError, BufferError)` around a `refcheck=True` resize and completed 0 of 200 resizes on every build where it ran — an object reachable from both the test frame and a closure cell has a reference count the check calls shared — and it still passed with `__releasebuffer__` gutted to a no-op +- Prefer a deterministic test over threads where the window can be entered on purpose: a callback from an argument the implementation converts inside the window (see `_AlignmentProbe`) tests the same guard without depending on the scheduler +- Tests must pass on free-threaded builds, and a test must not resize an allocation that other threads can reach — not on any version. `refcheck` is not a guard against other threads, and `tobytes` re-reads the pointer with nothing claimed, so a reallocer that retries until it slips between two reads is a use-after-free, not a test (ASan confirms it on 3.13t *and* 3.14t, in `tobytes`; a GIL build completes 200/200 resizes clean, which is why this looks fine locally). Park the readers on a `threading.Barrier` instead, resize while they are parked, and assert the resizes happened +- `numpy` is not a dependency of this package and CI does not install it for the test job, so any test that needs it must call `pytest.importorskip("numpy")` in the body — not import it at module level, which would break collection. Note that `pytest.importorskip` skips on `ModuleNotFoundError` only, so a probe that blocks the module by raising plain `ImportError` reports failures rather than skips and says nothing about the real behavior +- Content comparison alone does not prove a copy: a fresh allocation can be handed recycled heap memory that already holds the pattern, so `assert copy.tobytes() == source.tobytes()` has been observed to pass with the `memcpy` removed. Write the pattern a byte at a time so no freed `bytes` copy of it is left on the heap, and assert `_pointer` differs diff --git a/mkl/tests/test_mkl_memory.py b/mkl/tests/test_mkl_memory.py new file mode 100644 index 0000000..e38bf91 --- /dev/null +++ b/mkl/tests/test_mkl_memory.py @@ -0,0 +1,737 @@ +# 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 numbers +import sys +import threading + +import pytest + +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 + # mkl_calloc hands back zeroed memory + assert mem.tobytes() == bytes(nbytes) + + +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 + + +def test_mkl_memory_create_from_mkl_memory(): + mem1 = mkl.MKLMemory(1024) + mv = memoryview(mem1) + for i in range(len(mem1)): + mv[i] = (i * 5 + 1) % 256 + mv.release() + + mem2 = mkl.MKLMemory(mem1) + assert mem2.nbytes == mem1.nbytes + assert mem2.tobytes() == mem1.tobytes() + assert mem2._pointer != mem1._pointer + + mv = memoryview(mem2) + mv[0] = mem1.tobytes()[0] ^ 0xFF + mv.release() + assert mem2.tobytes()[0] != mem1.tobytes()[0] + assert mem1.tobytes()[0] == (0 * 5 + 1) % 256 + + +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 + + +class _AlignmentProbe: + def __init__(self, value, callback): + self._value = value + self._callback = callback + self._fired = False + + def __le__(self, other): + return self._value <= other + + def __gt__(self, other): + return self._value > other + + def _convert(self): + if not self._fired: + self._fired = True + self._callback() + return self._value + + def __int__(self): + return self._convert() + + def __index__(self): + return self._convert() + + +numbers.Integral.register(_AlignmentProbe) + + +def test_realloc_refused_while_copy_reads_source(): + source = mkl.MKLMemory(1024) + mv = memoryview(source) + for i in range(len(source)): + mv[i] = (i * 7 + 3) % 256 + mv.release() + pattern = source.tobytes() + + outcome = [] + + def probe(): + try: + source.realloc(256, refcheck=False) + outcome.append("resized") + except BufferError: + outcome.append("refused") + + copy = mkl.MKLMemory(source, alignment=_AlignmentProbe(64, probe)) + + assert outcome == ["refused"], f"resize was not refused: {outcome}" + assert source.nbytes == 1024 + assert copy.nbytes == 1024 + assert copy.alignment == 64 + assert copy.tobytes() == pattern + + +def test_copy_constructor_releases_source_claim(): + mem = mkl.MKLMemory(1024) + copy = mkl.MKLMemory(mem) + assert copy.nbytes == mem.nbytes + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_copy_constructor_releases_source_claim_on_failure(): + mem = mkl.MKLMemory(1024) + with pytest.raises(ValueError, match="Alignment"): + mkl.MKLMemory(mem, alignment=-1) + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_sizeof_accounts_for_object_too(): + small, large = 1024, 1 << 20 + mem, big = mkl.MKLMemory(small), mkl.MKLMemory(large) + + assert big.__sizeof__() - mem.__sizeof__() == large - small + overhead = mem.__sizeof__() - small + assert overhead > 0 + assert big.__sizeof__() - large == overhead + + assert sys.getsizeof(mem) >= mem.__sizeof__() + + mem.realloc(large) + assert mem.__sizeof__() == big.__sizeof__() + + +def test_buffer_protocol(): + mem = mkl.MKLMemory(1024) + mv1 = memoryview(mem) + mv2 = memoryview(mem) + try: + assert mv1.nbytes == mem.nbytes + mv1[0] = 7 + assert mv2[0] == 7 + mv2[1] = 9 + assert mv1[1] == 9 + + with pytest.raises(BufferError, match="exported buffers"): + mem.realloc(2048, refcheck=False) + mv1.release() + with pytest.raises(BufferError, match="exported buffers"): + mem.realloc(2048, refcheck=False) + finally: + mv2.release() + + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +def test_exported_buffer_describes_a_flat_writable_block(): + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + try: + assert mv.obj is mem + assert mv.format == "B" + assert mv.itemsize == 1 + assert mv.ndim == 1 + assert mv.shape == (mem.nbytes,) + assert mv.strides == (1,) + assert mv.suboffsets == () + assert not mv.readonly + assert mv.c_contiguous and mv.f_contiguous + mv[0] = 7 + assert mem.tobytes()[0] == 7 + finally: + mv.release() + + +def test_each_export_gives_back_its_reference(): + mem = mkl.MKLMemory(64) + before = sys.getrefcount(mem) + for _ in range(200): + memoryview(mem).release() + + assert sys.getrefcount(mem) == before + mem.realloc(128, refcheck=False) + assert mem.nbytes == 128 + + +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" + + +class _MKLMemorySubclass(mkl.MKLMemory): + pass + + +def test_pickling_preserves_subclass(): + import pickle + + mem = _MKLMemorySubclass(1024, alignment=128) + mv = memoryview(mem) + mv[:] = bytes((i * 3 + 1) % 256 for i in range(len(mem))) + mv.release() + + reconstructed = pickle.loads(pickle.dumps(mem)) + assert type(reconstructed) is _MKLMemorySubclass + assert reconstructed.nbytes == mem.nbytes + assert reconstructed.alignment == 128 + assert reconstructed.tobytes() == mem.tobytes() + + +def test_pickling_preserves_subclass_attributes(): + import pickle + + mem = _MKLMemorySubclass(256) + mem.label = "kept" + reconstructed = pickle.loads(pickle.dumps(mem)) + assert reconstructed.label == "kept" + + +def test_reconstruct_rejects_foreign_class(): + # pylint: disable-next=no-name-in-module + from mkl._mkl_memory import _mkl_memory_from_bytes + + with pytest.raises(TypeError, match="not a subclass of MKLMemory"): + _mkl_memory_from_bytes(b"abcd", 64, bytearray) + with pytest.raises(TypeError, match="not a subclass of MKLMemory"): + _mkl_memory_from_bytes(b"abcd", 64, "not a class") + + mem = _mkl_memory_from_bytes(b"abcd", 64) + assert type(mem) is mkl.MKLMemory + assert mem.tobytes() == b"abcd" + + +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 + 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] + + 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 + 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_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 be preserved + 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 + + +@pytest.mark.parametrize("new_nbytes", [0, -1]) +def test_realloc_validates_size_before_state(new_nbytes): + match = "New number of bytes must be positive" + mem = mkl.MKLMemory(1024) + mv = memoryview(mem) + try: + with pytest.raises(ValueError, match=match): + mem.realloc(new_nbytes) + with pytest.raises(ValueError, match=match): + mem.realloc(new_nbytes, refcheck=False) + finally: + mv.release() + + held = mem + with pytest.raises(ValueError, match=match): + held.realloc(new_nbytes) + + assert mem.nbytes == 1024 + mem.realloc(2048, refcheck=False) + assert mem.nbytes == 2048 + + +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_failed_realloc_leaves_the_allocation_untouched(): + mem = mkl.MKLMemory(1024, alignment=128) + mv = memoryview(mem) + mv[:] = bytes((i * 11 + 5) % 256 for i in range(len(mem))) + mv.release() + + pointer, nbytes = mem._pointer, mem.nbytes + alignment, pattern = mem.alignment, mem.tobytes() + + def assert_untouched(): + assert mem._pointer == pointer + assert mem.nbytes == nbytes + assert len(mem) == nbytes + assert mem.alignment == alignment + assert mem.tobytes() == pattern + + with pytest.raises(ValueError, match="must be positive"): + mem.realloc(0, refcheck=False) + assert_untouched() + + mv = memoryview(mem) + try: + with pytest.raises(BufferError, match="exported buffers"): + mem.realloc(2048, refcheck=False) + finally: + mv.release() + assert_untouched() + + held = mem # noqa: F841 + with pytest.raises(ValueError, match="Cannot realloc MKLMemory"): + mem.realloc(2048) + assert_untouched() + + with pytest.raises(TypeError): + mem.realloc(2048, False) + assert_untouched() + + mem.realloc(4096, refcheck=False) + assert mem.nbytes == 4096 + assert mem.alignment == alignment + assert mem.tobytes()[:nbytes] == pattern + + +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", [3, 5, 12, 24, 96, 100, 129, 1000]) +@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_must_be_a_power_of_two(construct, alignment): + with pytest.raises(ValueError, match="must be a power of two"): + construct(alignment) + + +@pytest.mark.parametrize( + "alignment", [1, 2, 4, 8, 16, 32, 64, 128, 256, 4096, 1 << 20] +) +def test_powers_of_two_are_honored(alignment): + for mem in ( + mkl.MKLMemory(1024, alignment=alignment), + mkl.MKLMemory(32, 32, alignment=alignment), + mkl.MKLMemory(mkl.MKLMemory(64), alignment=alignment), + ): + assert mem.alignment == alignment + assert mem._pointer % alignment == 0 + + +@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_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_leaves_a_consistent_allocation(): + 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() + + +def test_concurrent_reads_across_reallocs(): + sizes = [1 << 12, 1 << 13, 1 << 12, 1 << 14, 1 << 12] + n_readers = 3 + reads_per_round = 50 + mem = mkl.MKLMemory(sizes[0]) + errors = [] + resizes = 0 + + def fill(value): + mv = memoryview(mem) + try: + mv[:] = bytes([value]) * len(mv) + finally: + mv.release() + + quiesce = threading.Barrier(n_readers + 1, timeout=60) + fill(0xA5) + + def reader(): + try: + for _ in sizes: + for _ in range(reads_per_round): + mv = memoryview(mem) + try: + n = mv.nbytes + assert n == mem.nbytes + data = bytes(mv) + finally: + mv.release() + assert data == data[:1] * n, "view spans two blocks" + + copy = mem.tobytes() + assert len(copy) == n + assert copy == data + quiesce.wait() # no reader is inside `mem` past this point + quiesce.wait() # the resize is done + except threading.BrokenBarrierError: # pragma: no cover - on failure + pass + except Exception as e: # pragma: no cover - only on failure + errors.append(e) + quiesce.abort() + + def resizer(): + nonlocal resizes + try: + for round_ in range(len(sizes)): + quiesce.wait() + if round_ + 1 < len(sizes): + mem.realloc(sizes[round_ + 1], refcheck=False) + resizes += 1 + fill(round_ + 1) + quiesce.wait() + except threading.BrokenBarrierError: # pragma: no cover - on failure + pass + except Exception as e: # pragma: no cover - only on failure + errors.append(e) + quiesce.abort() + + ts = [threading.Thread(target=reader) for _ in range(n_readers)] + ts.append(threading.Thread(target=resizer)) + for t in ts: + t.start() + for t in ts: + t.join() + + assert not errors, f"Concurrent realloc/read errors: {errors}" + assert resizes == len(sizes) - 1 + assert mem.nbytes == sizes[-1] == len(mem)