Permalink
Browse files

Merge branch 'pickling_memoryview' into 'master'

pickle basevector without allocating new memory

See merge request jschoeberl/ngsolve!306
  • Loading branch information...
JSchoeberl committed Mar 1, 2018
2 parents 7ce7b6b + 895cf3d commit f07495d1e3c63e326f9b255661b8b28aa0073be6
Showing with 72 additions and 11 deletions.
  1. +18 −11 linalg/python_linalg.cpp
  2. +27 −0 ngstd/python_ngstd.cpp
  3. +18 −0 ngstd/python_ngstd.hpp
  4. +9 −0 python/__init__.py
@@ -75,20 +75,27 @@ void NGS_DLL_HEADER ExportNgla(py::module &m) {
"size"_a, "complex"_a=false, "entrysize"_a=1)
.def(py::pickle([] (const BaseVector& bv)
{
py::list lst;
for(auto val : bv.FVDouble())
lst.append(val);
return py::make_tuple(bv.Size(),bv.IsComplex(),bv.EntrySize(),lst);
MemoryView mv((void*) &bv.FVDouble()[0], sizeof(double) * bv.FVDouble().Size());
return py::make_tuple(bv.Size(),bv.IsComplex(),bv.EntrySize(),mv);
},
[] (py::tuple state) -> shared_ptr<BaseVector>
{
auto bv = CreateBaseVector(state[0].cast<size_t>(),
state[1].cast<bool>(),
state[2].cast<size_t>());
auto lst = state[3].cast<py::list>();
for(auto i : Range(py::len(lst)))
bv.FVDouble()[i] = lst[i].cast<double>();
return bv;
auto mv = state[3].cast<MemoryView>();
shared_ptr<BaseVector> bv;
if (state[1].cast<bool>())
{
// create basevector with owning pointer and afterwards assign it to mem
auto bptr = make_shared<S_BaseVectorPtr<Complex>>(0, state[2].cast<size_t>());
bptr->AssignMemory(state[0].cast<size_t>(), mv.Ptr());
return bptr;
}
else
{
// create basevector with owning pointer and afterwards assign it to mem
auto bptr = make_shared<S_BaseVectorPtr<double>>(0, state[2].cast<size_t>());
bptr->AssignMemory(state[0].cast<size_t>(), mv.Ptr());
return bptr;
}
}
))
.def("__str__", [](BaseVector &self) { return ToString<BaseVector>(self); } )
@@ -1,6 +1,7 @@
#ifdef NGS_PYTHON
#include "python_ngstd.hpp"
#include <Python.h>
#ifdef PARALLEL
bool MPIManager::initialized_by_me = false;
@@ -457,6 +458,32 @@ void NGS_DLL_HEADER ExportNgstd(py::module & m) {
.def("__exit__", &ParallelContextManager::Exit)
.def("__timing__", &TaskManager::Timing)
;
m.def("_PickleMemory", [](py::object pickler, MemoryView& view)
{
py::buffer_info bi((char*) view.Ptr(), view.Size());
pickler.attr("write")(py::bytes("\xf0"));
size_t size = view.Size();
pickler.attr("write")(py::bytes((char*) & size, sizeof(size_t)));
pickler.attr("write")(py::memoryview(bi));
});
m.def("_UnpickleMemory", [](py::object unpickler)
{
auto size = *(size_t*) PyBytes_AsString(unpickler.attr("read")(sizeof(size_t)).ptr());
char* mem = new char[size];
constexpr int BUFFER_SIZE = 8 * 1024 * 1024; // read 8 MB
size_t n = 0;
while (n + BUFFER_SIZE < size)
{
auto buffer = unpickler.attr("read")(BUFFER_SIZE);
memcpy(&mem[n], PyBytes_AsString(buffer.ptr()), BUFFER_SIZE);
n += BUFFER_SIZE;
}
auto buffer = unpickler.attr("read")(size-n);
memcpy(&mem[n], PyBytes_AsString(buffer.ptr()), size-n);
unpickler.attr("append")(MemoryView(mem,size));
});
py::class_<MemoryView>(m, "_MemoryView");
}
@@ -375,5 +375,23 @@ const char* docu_string(const char* str);
PYBIND11_DECLARE_HOLDER_TYPE(T, std::shared_ptr<T>);
namespace ngstd
{
// MemoryView for pickling without copying
// doesn't provide memory management, it has to be done from the outside!
class MemoryView
{
private:
void* ptr;
size_t size;
public:
MemoryView(void* aptr, size_t asize) : ptr(aptr), size(asize) { ; }
size_t Size() const { return size; }
void* Ptr() const { return ptr; }
};
}
#endif // NGS_PYTHON
#endif // PYTHON_NGSTD_HPP___
@@ -91,6 +91,15 @@ def TmpRedraw(*args, **kwargs):
finite element shape functions, and element-matrix/vector integrators
"""
# register our own memory pickler
import pickle
import ngsolve
pickle._Pickler.dispatch[ngsolve.ngstd._MemoryView] = ngsolve.ngstd._PickleMemory
pickle._Unpickler.dispatch[b"\xf0"[0]] = ngsolve.ngstd._UnpickleMemory
# use the python pickler and not cPickle one (cause we can't patch it)
pickle.Pickler, pickle.Unpickler = pickle._Pickler, pickle._Unpickler
pickle.dump, pickle.dumps, pickle.load, pickle.loads = pickle._dump, pickle._dumps, pickle._load, pickle._loads
__all__ = ngstd.__all__ + bla.__all__ +la.__all__ + fem.__all__ + comp.__all__ + solve.__all__ + utils.__all__ + ["Timing"]

0 comments on commit f07495d

Please sign in to comment.