Initialize the local_comm as well in CommSpec.

Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com>
acezen · Sep 25, 2020 · f721f02 · f721f02
1 parent 66b1df7
commit f721f02
Showing 1 changed file with 30 additions and 4 deletions.
diff --git a/grape/worker/comm_spec.h b/grape/worker/comm_spec.h
@@ -41,6 +41,7 @@ class CommSpec {
         fid_(0),
         fnum_(1),
         comm_(NULL_COMM),
+        local_comm_(NULL_COMM),
         owner_(false) {}
 
   CommSpec(const CommSpec& comm_spec)
@@ -51,17 +52,28 @@ class CommSpec {
         fid_(comm_spec.fid_),
         fnum_(comm_spec.fnum_),
         comm_(comm_spec.comm_),
+        local_comm_(comm_spec.local_comm_),
         owner_(false) {}
 
   ~CommSpec() {
-    if (owner_ && ValidComm(comm_)) {
-      MPI_Comm_free(&comm_);
+    if (owner_) {
+      if (ValidComm(comm_)) {
+        MPI_Comm_free(&comm_);
+      }
+      if (ValidComm(local_comm_)) {
+        MPI_Comm_free(&local_comm_);
+      }
     }
   }
 
   CommSpec& operator=(const CommSpec& rhs) {
-    if (owner_ && ValidComm(comm_)) {
-      MPI_Comm_free(&comm_);
+    if (owner_) {
+      if (ValidComm(comm_)) {
+        MPI_Comm_free(&comm_);
+      }
+      if (ValidComm(local_comm_)) {
+        MPI_Comm_free(&local_comm_);
+      }
     }
 
     worker_num_ = rhs.worker_num_;
@@ -71,6 +83,7 @@ class CommSpec {
     fid_ = rhs.fid_;
     fnum_ = rhs.fnum_;
     comm_ = rhs.comm_;
+    local_comm_ = rhs.local_comm_;
     owner_ = false;
 
     return *this;
@@ -91,7 +104,9 @@ class CommSpec {
 
   void Dup() {
     MPI_Comm old_comm = comm_;
+    MPI_Comm old_local_comm = local_comm_;
     MPI_Comm_dup(old_comm, &comm_);
+    MPI_Comm_dup(old_local_comm, &local_comm_);
     owner_ = true;
   }
 
@@ -113,6 +128,8 @@ class CommSpec {
 
   inline MPI_Comm comm() const { return comm_; }
 
+  inline MPI_Comm local_comm() const { return local_comm_; }
+
  private:
   void initLocalInfo() {
     char hn[MPI_MAX_PROCESSOR_NAME];
@@ -142,6 +159,14 @@ class CommSpec {
         ++local_num_;
       }
     }
+
+    std::sort(worker_host_names.begin(), worker_host_names.end());
+    std::map<std::string, int> hostname_to_host_id;
+    for (size_t idx = 0; idx < worker_host_names.size(); ++idx) {
+      hostname_to_host_id[worker_host_names[idx]] = idx;
+    }
+    int color = hostname_to_host_id[worker_host_names[worker_id_]];
+    MPI_Comm_split(comm_, color, worker_id_, &local_comm_);
   }
 
   int worker_num_;
@@ -154,6 +179,7 @@ class CommSpec {
   fid_t fnum_;
 
   MPI_Comm comm_;
+  MPI_Comm local_comm_;
   bool owner_;
 };