allow arrays to be GC'ed, using weak_ptr when registering them

Author: fschlimb

src/MPIMediator.cpp

@@ -15,11 +15,11 @@
 
 using OutputAdapter = bitsery::OutputBufferAdapter<Buffer>;
 using InputAdapter = bitsery::InputBufferAdapter<Buffer>;
-using array_keeper_type = std::unordered_map<uint64_t, tensor_i::ptr_type>;
+using array_keeper_type = std::unordered_map<uint64_t, tensor_i::ptr_type::weak_type>;
 using locker = std::lock_guard<std::mutex>;
 
 static array_keeper_type s_ak;
-static uint64_t s_last_id = 0;
+static uint64_t s_last_id = Mediator::LOCAL_ONLY;
 constexpr static int PULL_TAG = 4711;
 constexpr static int PUSH_TAG = 4712;
 static std::mutex ak_mutex;
@@ -53,6 +53,11 @@ uint64_t MPIMediator::register_array(tensor_i::ptr_type ary)
     return s_last_id;
 }
 
+uint64_t MPIMediator::unregister_array(uint64_t id)
+{
+    s_ak.erase(id);
+}
+
 void MPIMediator::pull(rank_type from, const tensor_i & ary, const NDSlice & slice, void * rbuff)
 {
     MPI_Comm comm = MPI_COMM_WORLD;
@@ -123,8 +128,9 @@ void MPIMediator::listen()
         if(x == s_ak.end()) throw(std::runtime_error("Encountered pull request for unknown tensor."));
         // Wait for previous answer to complete so that we can re-use the buffer
         MPI_Wait(&request_out, MPI_STATUS_IGNORE);
-        x->second->bufferize(slice, rbuff);
-        if(slice.size() * x->second->item_size() != rbuff.size()) throw(std::runtime_error("Got unexpected buffer size."));
+        auto ptr = x->second.lock();
+        ptr->bufferize(slice, rbuff);
+        if(slice.size() * ptr->item_size() != rbuff.size()) throw(std::runtime_error("Got unexpected buffer size."));
         MPI_Isend(rbuff.data(), rbuff.size(), MPI_CHAR, requester, PUSH_TAG, comm, &request_out);
     } while(true);
     // MPI_Cancel(&request_in);

src/include/ddptensor/MPIMediator.hpp

@@ -13,6 +13,7 @@ class MPIMediator : public Mediator
     MPIMediator();
     virtual ~MPIMediator();
     virtual uint64_t register_array(tensor_i::ptr_type ary);
+    virtual uint64_t unregister_array(uint64_t);
     virtual void pull(rank_type from, const tensor_i & ary, const NDSlice & slice, void * buffer);
 
 protected:

src/include/ddptensor/Mediator.hpp

@@ -10,8 +10,10 @@ class NDSlice;
 class Mediator
 {
 public:
+    enum : uint64_t {LOCAL_ONLY = 0};
     virtual ~Mediator() {}
     virtual uint64_t register_array(tensor_i::ptr_type ary) = 0;
+    virtual uint64_t unregister_array(uint64_t) = 0;
     virtual void pull(rank_type from, const tensor_i & ary, const NDSlice & slice, void * buffer) = 0;
 };
 

src/include/ddptensor/x.hpp

@@ -42,7 +42,7 @@ namespace x
     template<typename T>
     class DPTensorX : public DPTensorBaseX
     {
-        uint64_t _id = 0;
+        uint64_t _id = Mediator::LOCAL_ONLY;
         mutable rank_type _owner;
         PVSlice _slice;
         xt::xstrided_slice_vector _lslice;
@@ -117,6 +117,11 @@ namespace x
             _xarray = org;
         }
 
+        ~DPTensorX()
+        {
+            if(_id != Mediator::LOCAL_ONLY && theMediator) theMediator->unregister_array(_id);
+        }
+
         bool is_sliced() const
         {
             return _slice.is_sliced();

test/test_linalg.py

@@ -18,3 +18,5 @@
 b = np.reshape(b, (6,5))
 c = np.dot(a, b)
 print(c)
+
+dt.fini()

test/test_x.py

@@ -1,5 +1,6 @@
 import ddptensor as dt
 a = dt.ones([4,4], dt.float64)
+a = dt.ones([4,4], dt.float64)
 b = dt.ones([4,4], dt.float64)
 a += b
 print(a)
@@ -9,4 +10,5 @@
 print(a[0:1,0:1], float(a[0:1,0:1]), bool(a[0:1,0:1]), int(a[0:1,0:1]))
 print(a[0:2,0:2])
 print(float(a[1:2, 1:2]))
+
 dt.fini()