Merge branch 'pickling_memoryview' into 'master'

pickle basevector without allocating new memory

See merge request jschoeberl/ngsolve!306

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); } )

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");
 }
 
 

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___

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"]