Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Browse files
Merge branch 'pickling_memoryview' into 'master'
pickle basevector without allocating new memory
See merge request jschoeberl/ngsolve!306
- Loading branch information
Showing
with
72 additions
and
11 deletions.
-
+18
−11
linalg/python_linalg.cpp
-
+27
−0
ngstd/python_ngstd.cpp
-
+18
−0
ngstd/python_ngstd.hpp
-
+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"] |
|
|
|
|
|