Load balancing for LTO

[david2050]

Summary

Use bitcode size as an estimated for compilation time and sort
compilation tasks in order decreasing sizee to provide better
load balance when there are large number of taks.

[mehdi_amini]

lib/LTO/LTO.cpp

1301

The ThinLink has access to much fine grain data: like number of function and number of instructions. Have you looked into using this?

1312

FYI: https://github.com/llvm/llvm-project/blob/master/llvm/lib/LTO/ThinLTOCodeGenerator.cpp#L923

I used a separate index vector in ThinLTOCodeGenerator to do the same, it may make the code a bit easier to track (less pair<> and .first->second kind of thing).

[david2050]

lib/LTO/LTO.cpp

1301

I did not explore alternatives. This was effective for the workloads that motivated it so I thought I would share.

1312

Your implementation is cleaner. I may not have time to retest.

[david2050]

copy & paste & modify :-) from Mehdi pointerwd

[modocache]

Hello! Would anyone be willing to review this patch? Is there anything that could be tweaked, maybe putting this behavior behind an option, that would make the change more amenable to any reviewers?

[tejohnson]

Thanks! I think this makes sense, and since this is now the same as what is done in the old LTO API (ThinLTOCodeGenerator.cpp), thinking it would be better to extract into a common helper in LTO.cpp that is called from both places (that way if we find a way to improve the sorting, e.g. some of the other metrics @mehdi_amini mentioned, it only has to be changed in one place).

[mehdi_amini]

LGTM, but it'd be nice if @tejohnson or @pcc would also approve.

[tejohnson]

In D60495#1563124, @mehdi_amini wrote:

LGTM, but it'd be nice if @tejohnson or @pcc would also approve.

I added a suggestion a couple hours ago, just for commoning the code between the two LTO APIs.

[mehdi_amini]

In D60495#1563127, @tejohnson wrote:

In D60495#1563124, @mehdi_amini wrote:

LGTM, but it'd be nice if @tejohnson or @pcc would also approve.

I added a suggestion a couple hours ago, just for commoning the code between the two LTO APIs.

Looks like I was so slow to type that our answers crossed ;)

It makes sense to me! Thanks.

[aganea]

Hello @david2050! Are you still planning on working on this? If no, can I take over the patch and land it?

[aganea]

Commited as https://github.com/llvm/llvm-project/commit/6537004913f3009d896bc30856698e7d22199ba7

lib/LTO/LTO.cpp

 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/Utils/FunctionImportUtils.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 
+#include <numeric>
 #include <set>
 
 using namespace llvm;
 using namespace lto;
 using namespace object;
 
 #define DEBUG_TYPE "lto"
 
Error LTO::runThinLTO(AddStreamFn AddStream, NativeObjectCache Cache) {
   };
   thinLTOResolvePrevailingInIndex(ThinLTO.CombinedIndex, isPrevailing,
                                   recordNewLinkage);
 
   std::unique_ptr<ThinBackendProc> BackendProc =
       ThinLTO.Backend(Conf, ThinLTO.CombinedIndex, ModuleToDefinedGVSummaries,
                       AddStream, Cache);
 
-  // Tasks 0 through ParallelCodeGenParallelismLevel-1 are reserved for combined
-  // module and parallel code generation partitions.
-  unsigned Task = RegularLTO.ParallelCodeGenParallelismLevel;
-  for (auto &Mod : ThinLTO.ModuleMap) {
-    if (Error E = BackendProc->start(Task, Mod.second, ImportLists[Mod.first],
-                                     ExportLists[Mod.first],
-                                     ResolvedODR[Mod.first], ThinLTO.ModuleMap))
+  // Compute the ordering we will process the inputs: the rough heuristic here
+  // is to sort them per size so that the largest module get schedule as soon as
+  // possible. This is purely a compile-time optimization.
+  std::vector<int> ModulesOrdering;
+  ModulesOrdering.resize(ThinLTO.ModuleMap.size());
+  std::iota(ModulesOrdering.begin(), ModulesOrdering.end(), 0);
+  auto ModuleIter = ThinLTO.ModuleMap.begin();
+  llvm::sort(ModulesOrdering, [ModuleIter](int LeftIndex, int RightIndex) {
+    auto LSize = ModuleIter[LeftIndex].second.getBuffer().size();
+    auto RSize = ModuleIter[RightIndex].second.getBuffer().size();
+    return LSize > RSize;
+  });
+
+  // Task numbers start at ParallelCodeGenParallelismLevel if an LTO
+  // module is present, as tasks 0 through ParallelCodeGenParallelismLevel-1
+  // are reserved for parallel code generation partitions.
+  unsigned FirstTask = RegularLTO.ParallelCodeGenParallelismLevel;
+  for (unsigned Task : ModulesOrdering) {
+    auto &Mod = ModuleIter[Task];
+    if (Error E = BackendProc->start(
+            FirstTask + Task, Mod.second, ImportLists[Mod.first],
+            ExportLists[Mod.first], ResolvedODR[Mod.first], ThinLTO.ModuleMap))
       return E;
-    ++Task;
   }
-
   return BackendProc->wait();
 }
 
 Expected<std::unique_ptr<ToolOutputFile>>
 lto::setupOptimizationRemarks(LLVMContext &Context,
                               StringRef LTORemarksFilename,
                               StringRef LTORemarksPasses,
                               bool LTOPassRemarksWithHotness, int Count) {