1 files changed, 62 insertions, 23 deletions

diff --git a/contrib/llvm-project/llvm/lib/Transforms/Utils/SplitModule.cpp b/contrib/llvm-project/llvm/lib/Transforms/Utils/SplitModule.cpp
index 9c39c26d8b7a..a30afadf0365 100644
--- a/contrib/llvm-project/llvm/lib/Transforms/Utils/SplitModule.cpp
+++ b/contrib/llvm-project/llvm/lib/Transforms/Utils/SplitModule.cpp
@@ -55,6 +55,18 @@ using ClusterMapType = EquivalenceClasses<const GlobalValue *>;
 using ComdatMembersType = DenseMap<const Comdat *, const GlobalValue *>;
 using ClusterIDMapType = DenseMap<const GlobalValue *, unsigned>;
 
+bool compareClusters(const std::pair<unsigned, unsigned> &A,
+                     const std::pair<unsigned, unsigned> &B) {
+  if (A.second || B.second)
+    return A.second > B.second;
+  return A.first > B.first;
+}
+
+using BalancingQueueType =
+    std::priority_queue<std::pair<unsigned, unsigned>,
+                        std::vector<std::pair<unsigned, unsigned>>,
+                        decltype(compareClusters) *>;
+
 } // end anonymous namespace
 
 static void addNonConstUser(ClusterMapType &GVtoClusterMap,
@@ -105,7 +117,8 @@ static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
   // At this point module should have the proper mix of globals and locals.
   // As we attempt to partition this module, we must not change any
   // locals to globals.
-  LLVM_DEBUG(dbgs() << "Partition module with (" << M.size() << ")functions\n");
+  LLVM_DEBUG(dbgs() << "Partition module with (" << M.size()
+                    << ") functions\n");
   ClusterMapType GVtoClusterMap;
   ComdatMembersType ComdatMembers;
 
@@ -153,21 +166,10 @@ static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
 
   // Assigned all GVs to merged clusters while balancing number of objects in
   // each.
-  auto CompareClusters = [](const std::pair<unsigned, unsigned> &a,
-                            const std::pair<unsigned, unsigned> &b) {
-    if (a.second || b.second)
-      return a.second > b.second;
-    else
-      return a.first > b.first;
-  };
-
-  std::priority_queue<std::pair<unsigned, unsigned>,
-                      std::vector<std::pair<unsigned, unsigned>>,
-                      decltype(CompareClusters)>
-      BalancinQueue(CompareClusters);
+  BalancingQueueType BalancingQueue(compareClusters);
   // Pre-populate priority queue with N slot blanks.
   for (unsigned i = 0; i < N; ++i)
-    BalancinQueue.push(std::make_pair(i, 0));
+    BalancingQueue.push(std::make_pair(i, 0));
 
   using SortType = std::pair<unsigned, ClusterMapType::iterator>;
 
@@ -177,11 +179,13 @@ static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
   // To guarantee determinism, we have to sort SCC according to size.
   // When size is the same, use leader's name.
   for (ClusterMapType::iterator I = GVtoClusterMap.begin(),
-                                E = GVtoClusterMap.end(); I != E; ++I)
+                                E = GVtoClusterMap.end();
+       I != E; ++I)
     if (I->isLeader())
       Sets.push_back(
           std::make_pair(std::distance(GVtoClusterMap.member_begin(I),
-                                       GVtoClusterMap.member_end()), I));
+                                       GVtoClusterMap.member_end()),
+                         I));
 
   llvm::sort(Sets, [](const SortType &a, const SortType &b) {
     if (a.first == b.first)
@@ -191,9 +195,9 @@ static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
   });
 
   for (auto &I : Sets) {
-    unsigned CurrentClusterID = BalancinQueue.top().first;
-    unsigned CurrentClusterSize = BalancinQueue.top().second;
-    BalancinQueue.pop();
+    unsigned CurrentClusterID = BalancingQueue.top().first;
+    unsigned CurrentClusterSize = BalancingQueue.top().second;
+    BalancingQueue.pop();
 
     LLVM_DEBUG(dbgs() << "Root[" << CurrentClusterID << "] cluster_size("
                       << I.first << ") ----> " << I.second->getData()->getName()
@@ -211,7 +215,7 @@ static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
       CurrentClusterSize++;
     }
     // Add this set size to the number of entries in this cluster.
-    BalancinQueue.push(std::make_pair(CurrentClusterID, CurrentClusterSize));
+    BalancingQueue.push(std::make_pair(CurrentClusterID, CurrentClusterSize));
   }
 }
 
@@ -251,7 +255,7 @@ static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {
 void llvm::SplitModule(
     Module &M, unsigned N,
     function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
-    bool PreserveLocals) {
+    bool PreserveLocals, bool RoundRobin) {
   if (!PreserveLocals) {
     for (Function &F : M)
       externalize(&F);
@@ -268,6 +272,41 @@ void llvm::SplitModule(
   ClusterIDMapType ClusterIDMap;
   findPartitions(M, ClusterIDMap, N);
 
+  // Find functions not mapped to modules in ClusterIDMap and count functions
+  // per module. Map unmapped functions using round-robin so that they skip
+  // being distributed by isInPartition() based on function name hashes below.
+  // This provides better uniformity of distribution of functions to modules
+  // in some cases - for example when the number of functions equals to N.
+  if (RoundRobin) {
+    DenseMap<unsigned, unsigned> ModuleFunctionCount;
+    SmallVector<const GlobalValue *> UnmappedFunctions;
+    for (const auto &F : M.functions()) {
+      if (F.isDeclaration() ||
+          F.getLinkage() != GlobalValue::LinkageTypes::ExternalLinkage)
+        continue;
+      auto It = ClusterIDMap.find(&F);
+      if (It == ClusterIDMap.end())
+        UnmappedFunctions.push_back(&F);
+      else
+        ++ModuleFunctionCount[It->second];
+    }
+    BalancingQueueType BalancingQueue(compareClusters);
+    for (unsigned I = 0; I < N; ++I) {
+      if (auto It = ModuleFunctionCount.find(I);
+          It != ModuleFunctionCount.end())
+        BalancingQueue.push(*It);
+      else
+        BalancingQueue.push({I, 0});
+    }
+    for (const auto *const F : UnmappedFunctions) {
+      const unsigned I = BalancingQueue.top().first;
+      const unsigned Count = BalancingQueue.top().second;
+      BalancingQueue.pop();
+      ClusterIDMap.insert({F, I});
+      BalancingQueue.push({I, Count + 1});
+    }
+  }
+
   // FIXME: We should be able to reuse M as the last partition instead of
   // cloning it. Note that the callers at the moment expect the module to
   // be preserved, so will need some adjustments as well.
@@ -275,8 +314,8 @@ void llvm::SplitModule(
     ValueToValueMapTy VMap;
     std::unique_ptr<Module> MPart(
         CloneModule(M, VMap, [&](const GlobalValue *GV) {
-          if (ClusterIDMap.count(GV))
-            return (ClusterIDMap[GV] == I);
+          if (auto It = ClusterIDMap.find(GV); It != ClusterIDMap.end())
+            return It->second == I;
           else
             return isInPartition(GV, I, N);
         }));