Skip to content

Commit 5f82a77

Browse files
committed
add MKLMemory object, backed with mkl_malloc memory
exposes Python buffer protocol
1 parent 3fb7aff commit 5f82a77

4 files changed

Lines changed: 119 additions & 1 deletion

File tree

meson.build

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ py.extension_module(
4343
subdir: 'mkl'
4444
)
4545

46-
# Cython extension
46+
# Cython extensions
4747
py.extension_module(
4848
'_py_mkl_service',
4949
sources: ['mkl/_py_mkl_service.pyx'],
@@ -54,6 +54,17 @@ py.extension_module(
5454
subdir: 'mkl'
5555
)
5656

57+
py.extension_module(
58+
'_mkl_memory',
59+
sources: ['mkl/_mkl_memory.pyx'],
60+
dependencies: [mkl_dep],
61+
c_args: c_args,
62+
install_rpath: rpath,
63+
install: true,
64+
subdir: 'mkl'
65+
)
66+
67+
5768
# Python sources
5869
py.install_sources(
5970
[

mkl/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ def __exit__(self, *args):
5757
del RTLD_for_MKL
5858
del sys
5959

60+
from ._mkl_memory import MKLMemory
6061
from ._py_mkl_service import *
6162
from ._version import __version__
6263

mkl/_mkl_memory.pyx

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Copyright (c) 2018, Intel Corporation
2+
#
3+
# Redistribution and use in source and binary forms, with or without
4+
# modification, are permitted provided that the following conditions are met:
5+
#
6+
# * Redistributions of source code must retain the above copyright notice,
7+
# this list of conditions and the following disclaimer.
8+
# * Redistributions in binary form must reproduce the above copyright
9+
# notice, this list of conditions and the following disclaimer in the
10+
# documentation and/or other materials provided with the distribution.
11+
# * Neither the name of Intel Corporation nor the names of its contributors
12+
# may be used to endorse or promote products derived from this software
13+
# without specific prior written permission.
14+
#
15+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
16+
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
18+
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
19+
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
20+
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
21+
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
22+
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
23+
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
24+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
25+
26+
# distutils: language = c
27+
# cython: language_level=3
28+
29+
import numbers
30+
31+
from cpython cimport Py_buffer
32+
from libc.string cimport memcpy
33+
34+
from mkl._mkl_service cimport mkl_malloc, mkl_free
35+
36+
37+
cdef class MKLMemory:
38+
cdef void *_memory_ptr
39+
cdef Py_ssize_t nbytes
40+
41+
cdef _cinit_empty(self):
42+
self._memory_ptr = NULL
43+
self.nbytes = 0
44+
45+
cdef _cinit_alloc(self, Py_ssize_t nbytes, Py_ssize_t alignment):
46+
self._cinit_empty()
47+
48+
if (nbytes > 0):
49+
with nogil:
50+
p = mkl_malloc(nbytes, alignment)
51+
52+
if (p):
53+
self._memory_ptr = p
54+
self.nbytes = nbytes
55+
else:
56+
raise MemoryError(
57+
"MKL memory allocation failed."
58+
)
59+
else:
60+
raise ValueError(
61+
"Number of bytes of request allocation must be positive."
62+
)
63+
64+
cdef _cinit_other(self, object other, Py_ssize_t alignment):
65+
cdef MKLMemory other_mem
66+
if isinstance(other, MKLMemory):
67+
other_mem = <MKLMemory> other
68+
else:
69+
raise ValueError(
70+
f"Argument {other} is not of type MKLMemory."
71+
)
72+
self._cinit_alloc(other_mem.nbytes, alignment)
73+
with nogil:
74+
memcpy(self._memory_ptr, other_mem._memory_ptr, self.nbytes)
75+
76+
def __cinit__(self, other, *, Py_ssize_t alignment=64):
77+
if isinstance(other, numbers.Integral):
78+
self._cinit_alloc(other, alignment)
79+
else:
80+
self._cinit_other(other, alignment)
81+
82+
def __dealloc__(self):
83+
if not (self._memory_ptr is NULL):
84+
mkl_free(self._memory_ptr)
85+
self._cinit_empty()
86+
87+
cdef void *get_data_ptr(self):
88+
return self._memory_ptr
89+
90+
def __getbuffer__(self, Py_buffer *buffer, int flags):
91+
buffer.buf = <void *>self._memory_ptr
92+
buffer.format = "B" # byte
93+
buffer.internal = NULL # see References
94+
buffer.itemsize = 1
95+
buffer.len = self.nbytes
96+
buffer.ndim = 1
97+
buffer.obj = self
98+
buffer.readonly = 0
99+
buffer.shape = &self.nbytes
100+
buffer.strides = &buffer.itemsize
101+
buffer.suboffsets = NULL # for pointer arrays only
102+
103+
def __releasebuffer__(self, Py_buffer *buffer):
104+
pass

mkl/_mkl_service.pxd

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,8 @@ cdef extern from "mkl.h":
151151
MKL_INT64 mkl_mem_stat(int* buf)
152152
MKL_INT64 mkl_peak_mem_usage(int mode)
153153
int mkl_set_memory_limit(int mem_type, size_t limit)
154+
void *mkl_malloc(size_t size, int alignment) nogil
155+
void mkl_free(void *ptr) nogil
154156

155157
# Conditional Numerical Reproducibility
156158
int mkl_cbwr_set(int settings)

0 commit comments

Comments
 (0)