fixing possible races

Author: fschlimb

src/MPIMediator.cpp

@@ -4,6 +4,7 @@
 #include <thread>
 #include <iostream>
 #include <unordered_map>
+#include <mutex>
 #include <bitsery/bitsery.h>
 #include <bitsery/adapter/buffer.h>
 #include <bitsery/traits/vector.h>
@@ -15,12 +16,13 @@
 using OutputAdapter = bitsery::OutputBufferAdapter<Buffer>;
 using InputAdapter = bitsery::InputBufferAdapter<Buffer>;
 using array_keeper_type = std::unordered_map<uint64_t, tensor_i::ptr_type>;
+using locker = std::lock_guard<std::mutex>;
 
 static array_keeper_type s_ak;
 static uint64_t s_last_id = 0;
 constexpr static int PULL_TAG = 4711;
 constexpr static int PUSH_TAG = 4712;
-
+static std::mutex ak_mutex;
 
 MPIMediator::MPIMediator()
     : _listener(&MPIMediator::listen, this)
@@ -40,11 +42,13 @@ MPIMediator::~MPIMediator()
     ser.adapter().flush();
     MPI_Send(buff.data(), buff.size(), MPI_CHAR, rank, PULL_TAG, MPI_COMM_WORLD);
     _listener.join();
+    locker _l(ak_mutex);
     s_ak.clear();
 }
 
 uint64_t MPIMediator::register_array(tensor_i::ptr_type ary)
 {
+    locker _l(ak_mutex);
     s_ak[++s_last_id] = ary;
     return s_last_id;
 }
@@ -114,6 +118,7 @@ void MPIMediator::listen()
         MPI_Irecv(buff.data(), buff.size(), MPI_CHAR, MPI_ANY_SOURCE, PULL_TAG, comm, &request_in);
 
         // Now find the array in question and send back its bufferized slice
+        locker _l(ak_mutex);
         auto x = s_ak.find(id);
         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

src/MPITransceiver.cpp

@@ -56,6 +56,13 @@ static MPI_Op to_mpi(RedOpType o)
     }
 }
 
+
+
+void MPITransceiver::barrier()
+{
+    MPI_Barrier(MPI_COMM_WORLD);
+}
+
 void MPITransceiver::bcast(void * ptr, size_t N, rank_type root)
 {
     MPI_Bcast(ptr, N, MPI_CHAR, root, MPI_COMM_WORLD);

src/include/ddptensor/MPITransceiver.hpp

@@ -19,6 +19,8 @@ class MPITransceiver : public Transceiver
         return _rank;
     }
 
+
+    virtual void barrier();
     virtual void bcast(void * ptr, size_t N, rank_type root);
     virtual void reduce_all(void * inout, DType T, size_t N, RedOpType op);
 

src/include/ddptensor/Transceiver.hpp

@@ -12,6 +12,9 @@ class Transceiver
     virtual rank_type nranks() const = 0;
     virtual rank_type rank() const = 0;
 
+    // Barrier
+    virtual void barrier() = 0;
+
     // Broadcast data from root to all other processes
     // @param[inout] ptr   on root: pointer to data to be sent
     //                     on all other processes: pointer to buffer to store received data

src/include/ddptensor/ddptensor_impl.hpp

@@ -292,7 +292,7 @@ class dtensor_impl : public tensor_i
             auto ptr = buff.ptr;
             auto pylen = VPROD(buff.shape);
             assert(buff.itemsize == sizeof(T));
-            theTransceiver->reduce_all(ptr, DTYPE<T>::value, pylen, red_op(op));
+            theTransceiver->reduce_all(ptr, dtype(), pylen, red_op(op));
             return create_dtensor(pvslice(), new_shape, ary, REPLICATED);
         }
 
@@ -404,6 +404,7 @@ class dtensor_impl : public tensor_i
         NDSlice my_norm_slice = g_slc_view.map_slice(my_slice);
         std::cerr << "my_norm_slice: " << my_norm_slice << std::endl;
 
+        theTransceiver->barrier();
         _set_slice(cast(val), my_norm_slice, this, my_slice);
     }
 
@@ -431,9 +432,13 @@ class dtensor_impl : public tensor_i
     py::object get_slice(const NDSlice & slice) const
     {
         auto shp = slice.shape();
-        auto out = create_dtensor(PVSlice(shp, NOSPLIT), shp, DTYPE<T>::value, "empty");
-        _set_slice(this, slice, cast(out), {shp});
-        return cast(out)->_pyarray;
+        // Create dtensor without creating id: do not use create_dtensor
+        py::dict kwa;
+        kwa["dtype"] = get_impl_dtype<T>();
+        auto ary = _array_ns.attr("empty")(_make_tuple(shp), kwa);
+        auto out = dtensor_impl<T>(PVSlice(shp, NOSPLIT), shp, ary, theTransceiver->rank());
+        _set_slice(this, slice, &out, {shp});
+        return out._pyarray;
     }
 
     std::string __repr__() const