From 7e14227d78d058df348ee60748c2c552a6fa5f51 Mon Sep 17 00:00:00 2001 From: Lisandro Dalcin Date: Thu, 5 May 2016 18:23:17 +0300 Subject: [PATCH] More work on the memory object and MPI memory allocation --- src/MPI/MPI.pyx | 5 +-- src/MPI/asbuffer.pxi | 78 ++++++++++++++++++++++++++++++++++++-------- src/MPI/stdlib.pxi | 5 +++ src/pycompat.h | 18 ++++++++++ test/test_mem.py | 60 ++++++++++++++++++++++++++++++++++ 5 files changed, 150 insertions(+), 16 deletions(-) diff --git a/src/MPI/MPI.pyx b/src/MPI/MPI.pyx index 7fc508841..06480bdf3 100644 --- a/src/MPI/MPI.pyx +++ b/src/MPI/MPI.pyx @@ -96,13 +96,14 @@ def Alloc_mem(Aint size, Info info=INFO_NULL): CHKERR( MPI_Alloc_mem(size, cinfo, &base) ) return tomemory(base, size) -def Free_mem(memory): +def Free_mem(mem): """ Free memory allocated with `Alloc_mem()` """ cdef void *base = NULL - asmemory(memory, &base, NULL) + cdef memory m = asmemory(mem, &base, NULL) CHKERR( MPI_Free_mem(base) ) + m.release() # Initialization and Exit # ----------------------- diff --git a/src/MPI/asbuffer.pxi b/src/MPI/asbuffer.pxi index 6f28755c2..9d61ff5e9 100644 --- a/src/MPI/asbuffer.pxi +++ b/src/MPI/asbuffer.pxi @@ -1,5 +1,15 @@ #------------------------------------------------------------------------------ +cdef extern from "Python.h": + int PyIndex_Check(object) + int PySlice_Check(object) + int PySlice_GetIndicesEx(object, Py_ssize_t, + Py_ssize_t *, Py_ssize_t *, + Py_ssize_t *, Py_ssize_t *) except -1 + Py_ssize_t PyNumber_AsSsize_t(object, object) except -1 + +#------------------------------------------------------------------------------ + # Python 3 buffer interface (PEP 3118) cdef extern from "Python.h": enum: PY3 "(PY_MAJOR_VERSION>=3)" @@ -142,6 +152,10 @@ cdef class memory: cdef Py_buffer view + def __cinit__(self): + PyBuffer_FillInfo(&self.view, NULL, + NULL, 0, 0, PyBUF_SIMPLE) + def __dealloc__(self): PyBuffer_Release(&self.view) @@ -162,6 +176,18 @@ cdef class memory: def __get__(self): return self.view.readonly + # convenience methods + + def tobytes(self): + """Return the data in the buffer as a byte string""" + return PyBytes_FromStringAndSize(self.view.buf, self.view.len) + + def release(self): + """Release the underlying buffer exposed by the memory object""" + PyBuffer_Release(&self.view) + PyBuffer_FillInfo(&self.view, NULL, + NULL, 0, 0, PyBUF_SIMPLE) + # buffer interface (PEP 3118) def __getbuffer__(self, Py_buffer *view, int flags): @@ -208,17 +234,35 @@ cdef class memory: if i < 0: i += self.view.len if i < 0 or i >= self.view.len: raise IndexError("index out of range") - cdef unsigned char *buf = self.view.buf + cdef unsigned char *buf = self.view.buf return buf[i] - def __setitem__(self, Py_ssize_t i, unsigned char v): + def __setitem__(self, object item, object value): if self.view.readonly: raise TypeError("memory buffer is read-only") - if i < 0: i += self.view.len - if i < 0 or i >= self.view.len: - raise IndexError("index out of range") cdef unsigned char *buf = self.view.buf - buf[i] = v + cdef Py_ssize_t start, stop, step, length + cdef memory inmem + if PyIndex_Check(item): + start = PyNumber_AsSsize_t(item, IndexError) + if start < 0: start += self.view.len + if start < 0 or start >= self.view.len: + raise IndexError("index out of range") + buf[start] = value + elif PySlice_Check(item): + PySlice_GetIndicesEx(item, self.view.len, + &start, &stop, &step, &length) + if step != 1: + raise IndexError("slice with step not supported") + if PyIndex_Check(value): + memset(buf+start, value, length) + else: + inmem = getbuffer(value, 1, 0) + if inmem.view.len != length: + raise ValueError("slice length does not match buffer") + memcpy(buf+start, inmem.view.buf, length) + else: + raise TypeError("indices must be integers or slices") #------------------------------------------------------------------------------ @@ -270,23 +314,29 @@ cdef inline object getformat(memory buf): cdef inline memory getbuffer_r(object ob, void **base, MPI_Aint *size): cdef memory buf = getbuffer(ob, 1, 0) - if base != NULL: base[0] = buf.view.buf - if size != NULL: size[0] = buf.view.len + if base != NULL: base[0] = buf.view.buf + if size != NULL: size[0] = buf.view.len return buf cdef inline memory getbuffer_w(object ob, void **base, MPI_Aint *size): cdef memory buf = getbuffer(ob, 0, 0) - if base != NULL: base[0] = buf.view.buf - if size != NULL: size[0] = buf.view.len + if base != NULL: base[0] = buf.view.buf + if size != NULL: size[0] = buf.view.len return buf #------------------------------------------------------------------------------ -cdef inline object asmemory(object ob, void **base, MPI_Aint *size): - cdef memory buf = getbuffer_w(ob, base, size) - return buf +cdef inline memory asmemory(object ob, void **base, MPI_Aint *size): + cdef memory mem + if isinstance(ob, memory): + mem = ob + else: + mem = getbuffer(ob, 1, 0) + if base != NULL: base[0] = mem.view.buf + if size != NULL: size[0] = mem.view.len + return mem -cdef inline object tomemory(void *base, MPI_Aint size): +cdef inline memory tomemory(void *base, MPI_Aint size): return tobuffer(base, size) #------------------------------------------------------------------------------ diff --git a/src/MPI/stdlib.pxi b/src/MPI/stdlib.pxi index 012e742c0..5d4e4a882 100644 --- a/src/MPI/stdlib.pxi +++ b/src/MPI/stdlib.pxi @@ -1,5 +1,10 @@ #------------------------------------------------------------------------------ +cdef extern from * nogil: # "string.h" + void *memset(void *, int, size_t) + void *memcpy(void *, void *, size_t) + void *memmove(void *, void *, size_t) + cdef extern from * nogil: # "stdio.h" ctypedef struct FILE FILE *stdin, *stdout, *stderr diff --git a/src/pycompat.h b/src/pycompat.h index 0632a2d51..cb5be7200 100644 --- a/src/pycompat.h +++ b/src/pycompat.h @@ -30,6 +30,24 @@ _PyLong_AsByteArray(PyLongObject* v, /* ------------------------------------------------------------------------- */ +#ifdef PYPY_VERSION +#ifdef PySlice_GetIndicesEx +#undef PySlice_GetIndicesEx +#define PySlice_GetIndicesEx(s, n, start, stop, step, length) \ +PyPySlice_GetIndicesEx((PySliceObject *)(s), n, start, stop, step, length) +#else +#define PySlice_GetIndicesEx(s, n, start, stop, step, length) \ +PySlice_GetIndicesEx((PySliceObject *)(s), n, start, stop, step, length) +#endif +#else +#if PY_VERSION_HEX < 0x03020000 +#define PySlice_GetIndicesEx(s, n, start, stop, step, length) \ +PySlice_GetIndicesEx((PySliceObject *)(s), n, start, stop, step, length) +#endif +#endif + +/* ------------------------------------------------------------------------- */ + #if !defined(WITH_THREAD) #undef PyGILState_Ensure #define PyGILState_Ensure() ((PyGILState_STATE)0) diff --git a/test/test_mem.py b/test/test_mem.py index 207ad14b5..60d816ea4 100644 --- a/test/test_mem.py +++ b/test/test_mem.py @@ -23,5 +23,65 @@ def testMemory2(self): except NotImplementedError: return + def testMemorySequence(self): + n = 16 + try: + mem = MPI.Alloc_mem(n, MPI.INFO_NULL) + except NotImplementedError: + return + try: + mem.address + except AttributeError: + MPI.Free_mem(mem) + return + try: + self.assertEqual(len(mem), n) + self.assertTrue(mem.address != 0) + self.assertEqual(mem.nbytes, n) + self.assertFalse(mem.readonly) + def getitem(): return mem[n] + def delitem(): del mem[n] + def setitem1(): mem[n] = 0 + def setitem2(): mem[::2] = 0 + def setitem3(): mem[None] = 0 + self.assertRaises(IndexError, getitem) + self.assertRaises(Exception, delitem) + self.assertRaises(IndexError, setitem1) + self.assertRaises(IndexError, setitem2) + self.assertRaises(TypeError, setitem3) + for i in range(n): + mem[i] = i + for i in range(n): + self.assertEqual(mem[i], i) + mem[:] = 0 + for i in range(n): + self.assertEqual(mem[i], 0) + mem[:] = 255 + for i in range(n): + self.assertEqual(mem[i], 255) + mem[:n//2] = 1 + mem[n//2:] = 0 + for i in range(n//2): + self.assertEqual(mem[i], 1) + for i in range(n//2, n): + self.assertEqual(mem[i], 0) + mem[:] = 0 + mem[1:5] = b"abcd" + mem[10:13] = b"xyz" + self.assertEqual(mem[0], 0) + for i, c in enumerate("abcd"): + self.assertEqual(mem[1+i], ord(c)) + for i in range(5, 10): + self.assertEqual(mem[i], 0) + for i, c in enumerate("xyz"): + self.assertEqual(mem[10+i], ord(c)) + for i in range(13, n): + self.assertEqual(mem[i], 0) + finally: + MPI.Free_mem(mem) + self.assertEqual(mem.address, 0) + self.assertEqual(mem.nbytes, 0) + self.assertFalse(mem.readonly) + if __name__ == '__main__': unittest.main()