Support cluster split functionality

Author: MC952-arch

flagcx/core/cluster.cc

@@ -1,6 +1,29 @@
 #include "cluster.h"
 #include <cstring>
 
+flagcxResult_t parseClusterSplitList(const char *input,
+                                     std::vector<int> &output) {
+  if (input == NULL) {
+    return flagcxSuccess;
+  }
+
+  output.clear();
+  std::stringstream ss(input);
+  std::string token;
+
+  while (std::getline(ss, token, ',')) {
+    try {
+      int value = std::stoi(token);
+      output.push_back(value);
+    } catch (const std::exception &e) {
+      WARN("Invalid Cluster split info, its format should be like '2,4,8,...'");
+      return flagcxSystemError;
+    }
+  }
+
+  return flagcxSuccess;
+}
+
 flagcxResult_t flagcxCollectClusterInfos(const flagcxVendor *allData,
                                          flagcxCommunicatorType_t *type,
                                          int *homo_rank, int *homo_root_rank,
@@ -11,7 +34,7 @@ flagcxResult_t flagcxCollectClusterInfos(const flagcxVendor *allData,
   *homo_root_rank = 0;
   *homo_ranks = 1;
   *cluster_id = 0;
-  *cluster_inter_rank = -1;
+  *cluster_inter_rank = -1; // deprecated, to be removed
   *ncluster = 1;
   *type = flagcxCommunicatorHomo;
 
@@ -55,23 +78,48 @@ flagcxResult_t flagcxCollectClusterInfos(const flagcxVendor *allData,
   }
 
   if (*type == flagcxCommunicatorHybrid) {
-    const char *useDev = flagcxGetEnv("FLAGCX_USEDEV");
-    int useDev_;
-    if (useDev == NULL) {
-      useDev_ = -1;
-    } else {
-      useDev_ = std::stoi(useDev);
-    }
-    if (*homo_rank == useDev_) {
-      *cluster_inter_rank = rank;
-    }
-    if (*homo_ranks <= useDev_ && *homo_rank == *homo_ranks - 1) {
-      *cluster_inter_rank = rank;
-    }
     *cluster_id = currCluster;
     *ncluster = numClusters;
   }
 
+  // split and obtain sub-clusters
+  const char *clusterSplitInfo = flagcxGetEnv("FLAGCX_CLUSTER_SPLIT_LIST");
+  std::vector<int> clusterSplitList;
+  FLAGCXCHECK(parseClusterSplitList(clusterSplitInfo, clusterSplitList));
+  int subClusterId = 0;
+  for (int i = 0; i < currCluster; ++i) {
+    subClusterId += clusterSplitList[i];
+  }
+  int subHomoRanks = (*homo_ranks) / clusterSplitList[currCluster];
+  int hasRes =
+      (((*homo_rank) / subHomoRanks) >= clusterSplitList[currCluster]) ? 1 : 0;
+  subClusterId += (hasRes == 1) ? ((*homo_rank) / subHomoRanks) - 1
+                                : ((*homo_rank) / subHomoRanks);
+  int subHomoRank = (hasRes == 1) ? subHomoRanks + ((*homo_rank) % subHomoRanks)
+                                  : ((*homo_rank) % subHomoRanks);
+  int subHomoRootRank = rank - subHomoRank;
+  if (hasRes == 1 ||
+      ((*homo_rank) / subHomoRanks) == clusterSplitList[currCluster] - 1) {
+    subHomoRanks += (*homo_ranks) % clusterSplitList[currCluster];
+  }
+  int subNClusters = 0;
+  for (int i = 0; i < (*ncluster); ++i) {
+    subNClusters += clusterSplitList[i];
+  }
+  *homo_rank = subHomoRank;
+  *homo_root_rank = subHomoRootRank;
+  *homo_ranks = subHomoRanks;
+  *cluster_id = subClusterId;
+  *ncluster = subNClusters;
+  *type =
+      (subNClusters > 1) ? flagcxCommunicatorHybrid : flagcxCommunicatorHomo;
+
+  INFO(FLAGCX_INIT,
+       "rank = %d, nranks = %d, nclusters = %d, cluster_id = %d, "
+       "homo_rank = %d, homo_root_rank = %d, homo_ranks = %d, type = %d",
+       rank, nranks, *ncluster, *cluster_id, *homo_rank, *homo_root_rank,
+       *homo_ranks, *type);
+
   return flagcxSuccess;
 }
 

flagcx/core/cluster.h

@@ -4,8 +4,14 @@
 #include "adaptor.h"
 #include "flagcx.h"
 #include "param.h"
+#include <iostream>
 #include <map>
+#include <sstream>
 #include <string>
+#include <vector>
+
+flagcxResult_t parseClusterSplitList(const char *input,
+                                     std::vector<int> &output);
 
 flagcxResult_t flagcxCollectClusterInfos(const flagcxVendor *allData,
                                          flagcxCommunicatorType_t *type,