[llvm] [llvm][transforms] Add a new algorithm to SplitModule (PR #95941)

[Ilia Sergachev via llvm-commits]

================
@@ -268,6 +268,38 @@ void llvm::SplitModule(
   ClusterIDMapType ClusterIDMap;
   findPartitions(M, ClusterIDMap, N);
 
+  // Find empty modules and functions not mapped to modules in ClusterIDMap.
+  // Map these functions to the empty modules 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 (TryToAvoidEmptyModules) {
+    DenseSet<unsigned> NonEmptyModules;
+    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
+        NonEmptyModules.insert(It->second);
+    }
+    SmallVector<unsigned> EmptyModules;
+    for (unsigned I = 0; I < N; ++I) {
+      if (!NonEmptyModules.contains(I))
+        EmptyModules.push_back(I);
+    }
+    auto NextEmptyModuleIt = EmptyModules.begin();
+    for (const auto F : UnmappedFunctions) {
+      if (NextEmptyModuleIt == EmptyModules.end())
+        break;
+      ClusterIDMap.insert({F, *NextEmptyModuleIt});
+      ++NextEmptyModuleIt;
+    }
----------------
sergachev wrote:

> I don't quite get this loop, seems like you won't distribute all the unmapped functions, but only at most a number equal to the empty modules? All the others will be distributed by isInPartition() if I read this loop correctly.

Correct. Our main use case for this is to request the number of modules equal to the number of functions - then (provided that the functions are independent) the algorithm generates exactly 1 function per module. The behavior corresponds to the name of the flag - TryToAvoidEmptyModules.

> Why don't you just circle back to EmptyModules.begin() instead of breaking out of the loop?

Can do that too - this would become a different algorithm, completely overriding the hash-based distribution. For our goal (1 function per module) this would work equally. Do you think we should change it this way?

> Also should we skip the loop over UnmappedFunctions if EmptyModules.empty() ?

Indeed. Fixed.


https://github.com/llvm/llvm-project/pull/95941