Merge branch 'pickling_memoryview' into 'master'

pickle basevector without allocating new memory See merge request jschoeberl/ngsolve!306
NGSolve · Mar 1, 2018 · f07495d1e3c63e326f9b255661b8b28aa0073be6 · f07495d
2 parents 7ce7b6b + 895cf3d
commit f07495d1e3c63e326f9b255661b8b28aa0073be6
Unified Split

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
diff --git a/linalg/python_linalg.cpp b/linalg/python_linalg.cpp
@@ -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); } )
diff --git a/ngstd/python_ngstd.cpp b/ngstd/python_ngstd.cpp
@@ -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");
 }


diff --git a/ngstd/python_ngstd.hpp b/ngstd/python_ngstd.hpp
@@ -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___
diff --git a/python/__init__.py b/python/__init__.py
@@ -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"]