#1: inform: fix some bugs, incl. termination, get inform working

Author: lifflander

examples/test_example.cc

@@ -17,35 +17,38 @@ int main(int argc, char** argv) {
   auto comm = vt_lb::comm::CommMPI();
   comm.init(argc, argv);
 
-  auto cls = std::make_unique<MyClass>();
-  auto handle = comm.registerInstanceCollective(cls.get());
-  auto rank = comm.getRank();
+  //auto cls = std::make_unique<MyClass>();
+  //auto handle = comm.registerInstanceCollective(cls.get());
+  //auto rank = comm.getRank();
 
-  int value = 10;
-  int recv_value = 0;
-  handle.reduce(1, MPI_INT, MPI_SUM, &value, &recv_value, 1);
+  // int value = 10;
+  // int recv_value = 0;
+  // handle.reduce(1, MPI_INT, MPI_SUM, &value, &recv_value, 1);
 
-  fmt::print("Rank {}: reduced value is {}\n", rank, recv_value);
+  // fmt::print("Rank {}: reduced value is {}\n", rank, recv_value);
 
-  if (rank == 0) {
-    handle[1].send<&MyClass::myHandler2>(std::string{"hello from rank 0"});
-  }
-  if (rank == 1) {
-    handle[0].send<&MyClass::myHandler>(2, 10.3);
-  }
-
-  while (comm.poll()) {
-  }
+  // if (rank == 0) {
+  //   handle[1].send<&MyClass::myHandler2>(std::string{"hello from rank 0"});
+  // }
+  // if (rank == 1) {
+  //   handle[0].send<&MyClass::myHandler>(2, 10.3);
+  // }
 
-  printf("out of poll\n");
+  printf("Running runLB\n");
+  //comm.barrier();
 
   vt_lb::runLB(
     vt_lb::DriverAlgoEnum::TemperedLB,
     comm,
-    vt_lb::algo::temperedlb::Configuration{},
+    vt_lb::algo::temperedlb::Configuration{comm.numRanks()},
     nullptr
   );
 
+  while (comm.poll()) {
+  }
+
+  printf("out of poll\n");
+
   comm.finalize();
   return 0;
 }

src/vt-lb/algo/temperedlb/temperedlb.h

@@ -48,6 +48,7 @@
 #include <vt-lb/algo/baselb/baselb.h>
 
 #include <limits>
+#include <random>
 
 namespace vt_lb::algo::temperedlb {
 
@@ -63,14 +64,22 @@ struct WorkModel {
 };
 
 struct Configuration {
+  Configuration() = default;
+  Configuration(int num_ranks) {
+    f_ = 2;
+    k_max_ = std::ceil(std::sqrt(std::log(num_ranks)/std::log(2.0)));
+  }
+
   /// @brief  Number of trials to perform
   int num_trials_ = 1;
   /// @brief  Number of iterations per trial
   int num_iters_ = 10;
   /// @brief  Fanout for information propagation
   int f_ = 2;
   /// @brief  Number of rounds of information propagation
-  int k_max_ = 0;
+  int k_max_ = 1;
+
+  bool async_ip_ = true;
 
   /// @brief  Work model parameters (rank-alpha, beta, gamma, delta)
   WorkModel work_model_;
@@ -79,6 +88,71 @@ struct Configuration {
   double converge_tolerance_ = 0.01;
 };
 
+template <typename CommT, typename DataT, typename JoinT>
+struct InformationPropagation {
+  using ThisType = InformationPropagation<CommT, DataT, JoinT>;
+  using JoinedDataType = std::unordered_map<int, DataT>;
+
+  InformationPropagation(CommT& comm, int f, int k_max)
+    : comm_(comm.clone()), f_(f), k_max_(k_max)
+  {
+    handle_ = comm_.template registerInstanceCollective<ThisType>(this);
+  }
+
+  void run(DataT initial_data) {
+    // Insert my own rank to avoid self-selection
+    already_selected_.insert(comm_.getRank());
+
+    local_data_[comm_.getRank()] = initial_data;
+
+    sendToFanout(1, local_data_);
+
+    while (comm_.poll()) {
+      // do nothing
+    }
+
+    printf("%d: done with poll: local_data size=%zu\n", comm_.getRank(), local_data_.size());
+  }
+
+  void sendToFanout(int round, JoinedDataType const& data) {
+    int num_ranks = comm_.numRanks();
+
+    for (int i = 1; i <= f_; ++i) {
+      if (already_selected_.size() >= static_cast<size_t>(num_ranks)) {
+        return;
+      }
+
+      std::uniform_int_distribution<int> dist(0, num_ranks - 1);
+      int target = -1;
+      do {
+        target = dist(gen_select_);
+      } while (already_selected_.find(target) != already_selected_.end());
+
+      already_selected_.insert(target);
+
+      //printf("rank %d sending to rank %d\n", comm_.getRank(), target);
+      handle_[target].template send<&ThisType::infoPropagateHandler>(round, data);
+    }
+  }
+
+  void infoPropagateHandler(int round, JoinedDataType incoming_data) {
+    // Process incoming data and add to local data
+     local_data_.insert(incoming_data.begin(), incoming_data.end());
+    if (round < k_max_) {
+      sendToFanout(round + 1, local_data_);
+    }
+  }
+
+private:
+  CommT comm_;
+  int f_ = 2;
+  int k_max_ = 0;
+  std::unordered_set<int> already_selected_;
+  std::unordered_map<int, DataT> local_data_;
+  std::mt19937 gen_select_{std::random_device{}()};
+  typename CommT::template HandleType<ThisType> handle_;
+};
+
 template <typename CommT>
 struct TemperedLB : baselb::BaseLB {
 
@@ -97,11 +171,22 @@ struct TemperedLB : baselb::BaseLB {
   { }
 
   void makeHandle() {
+    // printf("makeHandle\n");
     handle_ = comm_.template registerInstanceCollective<TemperedLB<CommT>>(this);
   }
 
   void run() {
     // Implementation of the TemperedLB algorithm would go here
+    auto& wm = config_.work_model_;
+    if (wm.beta == 0.0 && wm.gamma == 0.0 && wm.delta == 0.0) {
+      using LoadType = double;
+      printf("start InformationPropagation\n");
+      auto ip = InformationPropagation<CommT, LoadType, TemperedLB<CommT>>(
+        comm_, config_.f_, config_.k_max_
+      );
+      ip.run(10.0);
+    }
+
   }
 
 private:

src/vt-lb/comm/comm_mpi.h

@@ -174,23 +174,51 @@ struct CommMPI {
    * \brief Construct a new CommMPI instance
    * \param comm The MPI communicator to use (defaults to MPI_COMM_WORLD)
    */
-  CommMPI(MPI_Comm comm = MPI_COMM_WORLD) : comm_(comm) {}
+  explicit CommMPI(MPI_Comm comm = MPI_COMM_WORLD) : comm_(comm) { }
 
+  CommMPI(CommMPI const&) = delete;
+  CommMPI(CommMPI&&) = delete;
+  CommMPI& operator=(CommMPI const&) = delete;
+  CommMPI& operator=(CommMPI&&) = delete;
+
+private:
+  CommMPI(MPI_Comm comm, int rank, int size)
+    : comm_(comm), cached_rank_(rank), cached_size_(size)
+  {
+    initTermination();
+  }
+
+public:
  /**
    * \brief Initialize MPI
    * \param argc Pointer to argument count
    * \param argv Pointer to argument array
    */
   void init(int& argc, char**& argv) {
     MPI_Init(&argc, &argv);
+    MPI_Comm_rank(comm_, &cached_rank_);
+    MPI_Comm_size(comm_, &cached_size_);
     initTermination();
+    // printf("%d: Initialized MPI with %d ranks\n", cached_rank_, cached_size_);
+  }
+
+  CommMPI clone() const {
+    MPI_Comm new_comm;
+    MPI_Comm_dup(comm_, &new_comm);
+    return CommMPI{new_comm, cached_rank_, cached_size_};
   }
 
   /**
    * \brief Finalize MPI
    */
   void finalize() {
+    // printf("%d: Finalizing MPI\n", cached_rank_);
     MPI_Finalize();
+    comm_ = MPI_COMM_NULL;
+  }
+
+  void barrier() {
+    MPI_Barrier(comm_);
   }
 
   /**
@@ -287,6 +315,13 @@ struct CommMPI {
     buf_interpreter.classIndex() = idx;
     buf_interpreter.isTermination() = is_termination_msg ? 1 : 0;
 
+    // printf("%d: MPI_Send to %d handler_index=%d class_index=%d is_termination=%d\n",
+    //   cached_rank_,
+    //   dest,
+    //   buf_interpreter.handlerIndex(),
+    //   buf_interpreter.classIndex(),
+    //   buf_interpreter.isTermination()
+    // );
     MPI_Request req;
     MPI_Isend(
       send_ptr, send_ptr_size, MPI_BYTE, dest, 0, comm_,
@@ -310,6 +345,7 @@ struct CommMPI {
     {
       int flag = 0;
       MPI_Status status;
+      // printf("%d: MPI_Iprobe\n", cached_rank_);
       MPI_Iprobe(MPI_ANY_SOURCE, 0, comm_, &flag, &status);
       if (flag) {
         int count = 0;
@@ -326,12 +362,19 @@ struct CommMPI {
           throw std::runtime_error("Buffer alignment error");
         }
 
+        // printf("%d: MPI_Recv from %d of size %d\n", cached_rank_, status.MPI_SOURCE, count);
         MPI_Recv(buf.data(), count, MPI_BYTE, status.MPI_SOURCE, status.MPI_TAG, comm_, MPI_STATUS_IGNORE);
         BufferIntInterpreter buf_interpreter(buf.data());
         int handler_index = buf_interpreter.handlerIndex();
         int class_index = buf_interpreter.classIndex();
         bool is_termination = buf_interpreter.isTermination() != 0;
 
+        // printf("%d: Received message: handler_index=%d class_index=%d is_termination=%d\n",
+        //   cached_rank_,
+        //   handler_index,
+        //   class_index,
+        //   is_termination ? 1 : 0
+        // );
         auto mem_fn = detail::getMember(handler_index);
         assert(class_map_.find(class_index) != class_map_.end() && "Class index not found");
         mem_fn->dispatch(
@@ -367,6 +410,10 @@ struct CommMPI {
    */
   int numRanks() const {
     int size = 0;
+    // printf("%d: numRanks\n", cached_rank_);
+    if (comm_ == MPI_COMM_NULL) {
+      throw std::runtime_error("Communicator is not initialized");
+    }
     MPI_Comm_size(comm_, &size);
     return size;
   }
@@ -377,18 +424,25 @@ struct CommMPI {
    */
   int getRank() const {
     int rank = 0;
+    // printf("%d: getRank\n", cached_rank_);
+    if (comm_ == MPI_COMM_NULL) {
+      throw std::runtime_error("Communicator is not initialized");
+    }
     MPI_Comm_rank(comm_, &rank);
     return rank;
   }
 
 private:
   void initTermination();
 
-  MPI_Comm comm_;
+  MPI_Comm comm_ = MPI_COMM_NULL;
   std::list<std::tuple<MPI_Request, std::unique_ptr<char[]>>> pending_;
   int next_class_index_ = 0;
   std::unordered_map<int, void*> class_map_;
   std::unique_ptr<detail::TerminationDetector> termination_detector_;
+
+  int cached_rank_ = -1;
+  int cached_size_ = -1;
 };
 
 inline void CommMPI::initTermination() {

src/vt-lb/comm/termination.cc

@@ -119,6 +119,9 @@ void TerminationDetector::onResponse(uint64_t in_sent, uint64_t in_recv) {
     if (rank_ == 0) {
       // Root checks for termination
 
+      global_sent1_ += sent_;
+      global_recv1_ += recv_;
+
 #if DEBUG_TERMINATION
       printf("Root total: s1=%lld, r1=%lld, s2=%lld, r2=%lld\n",
              global_sent1_, global_recv1_, global_sent2_, global_recv2_);
@@ -167,7 +170,9 @@ void TerminationDetector::notifyMessageReceive() {
 }
 
 void TerminationDetector::terminated() {
+#if DEBUG_TERMINATION
   printf("%d: Terminated!\n", rank_);
+#endif
   terminated_ = true;
   for (int i = 0; i < num_children_; i++) {
     handle_[first_child_ + i].sendTerm<&TerminationDetector::terminated>();