[BOLT] Add randomN split strategy

This adds a strategy to split functions into a random number of
fragments at randomly chosen split points.

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D130647

bolt/include/bolt/Passes/SplitFunctions.h

@@ -24,6 +24,9 @@ enum SplitFunctionsStrategy : char {
   /// Split each function into a hot and cold fragment at a randomly chosen
   /// split point (ignoring any available profiling information).
   Random2,
+  /// Split each function into N fragments at a randomly chosen split points
+  /// (ignoring any available profiling information).
+  RandomN,
   /// Split all basic blocks of each function into fragments such that each
   /// fragment contains exactly a single basic block.
   All

bolt/lib/Passes/LoopInversionPass.cpp

@@ -44,7 +44,8 @@ bool LoopInversionPass::runOnFunction(BinaryFunction &BF) {
     const unsigned BBIndex = BB->getLayoutIndex();
     const unsigned SuccBBIndex = SuccBB->getLayoutIndex();
     if (SuccBB == PredBB && BB != SuccBB && BBIndex != 0 && SuccBBIndex != 0 &&
-        SuccBB->succ_size() == 2 && BB->isCold() == SuccBB->isCold()) {
+        SuccBB->succ_size() == 2 &&
+        BB->getFragmentNum() == SuccBB->getFragmentNum()) {
       // Get the second successor (after loop BB)
       BinaryBasicBlock *SecondSucc = nullptr;
       for (BinaryBasicBlock *Succ : SuccBB->successors()) {

bolt/lib/Passes/SplitFunctions.cpp

@@ -11,13 +11,19 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/SplitFunctions.h"
+#include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <algorithm>
 #include <iterator>
+#include <numeric>
 #include <random>
 #include <vector>
 
@@ -88,6 +94,10 @@ static cl::opt<SplitFunctionsStrategy> SplitStrategy(
         SplitFunctionsStrategy::Random2, "random2",
         "split each function into a hot and cold fragment at a randomly chosen "
         "split point (ignoring any available profiling information)")),
+    cl::values(clEnumValN(
+        SplitFunctionsStrategy::RandomN, "randomN",
+        "split each function into N fragments at a randomly chosen split "
+        "points (ignoring any available profiling information)")),
     cl::values(clEnumValN(
         SplitFunctionsStrategy::All, "all",
         "split all basic blocks of each function into fragments such that each "
@@ -139,7 +149,7 @@ struct SplitRandom2 {
     using DiffT = typename std::iterator_traits<It>::difference_type;
     const DiffT NumOutlineableBlocks = End - Start;
 
-    // We want to split at least one block unless there are not blocks that can
+    // We want to split at least one block unless there are no blocks that can
     // be outlined
     const auto MinimumSplit = std::min<DiffT>(NumOutlineableBlocks, 1);
     std::uniform_int_distribution<DiffT> Dist(MinimumSplit,
@@ -155,6 +165,51 @@ struct SplitRandom2 {
   }
 };
 
+struct SplitRandomN {
+  std::minstd_rand0 *Gen;
+
+  explicit SplitRandomN(std::minstd_rand0 &Gen) : Gen(&Gen) {}
+
+  bool canSplit(const BinaryFunction &BF) { return true; }
+  bool canOutline(const BinaryBasicBlock &BB) { return true; }
+
+  template <typename It> void partition(It Start, It End) const {
+    using DiffT = typename std::iterator_traits<It>::difference_type;
+    const DiffT NumOutlineableBlocks = End - Start;
+
+    // We want to split at least one fragment if possible
+    const auto MinimumSplits = std::min<DiffT>(NumOutlineableBlocks, 1);
+    std::uniform_int_distribution<DiffT> Dist(MinimumSplits,
+                                              NumOutlineableBlocks);
+    // Choose how many splits to perform
+    const DiffT NumSplits = Dist(*Gen);
+
+    // Draw split points from a lottery
+    SmallVector<unsigned, 0> Lottery(NumOutlineableBlocks);
+    std::iota(Lottery.begin(), Lottery.end(), 0u);
+    std::shuffle(Lottery.begin(), Lottery.end(), *Gen);
+    Lottery.resize(NumSplits);
+    llvm::sort(Lottery);
+
+    // Add one past the end entry to lottery
+    Lottery.push_back(NumOutlineableBlocks);
+
+    unsigned LotteryIndex = 0;
+    unsigned BBPos = 0;
+    for (BinaryBasicBlock *const BB : make_range(Start, End)) {
+      // Check whether to start new fragment
+      if (BBPos >= Lottery[LotteryIndex])
+        ++LotteryIndex;
+
+      // Because LotteryIndex is 0 based and cold fragments are 1 based, we can
+      // use the index to assign fragments.
+      BB->setFragmentNum(FragmentNum(LotteryIndex));
+
+      ++BBPos;
+    }
+  }
+};
+
 struct SplitAll {
   bool canSplit(const BinaryFunction &BF) { return true; }
   bool canOutline(const BinaryBasicBlock &BB) { return true; }
@@ -203,6 +258,12 @@ void SplitFunctions::runOnFunctions(BinaryContext &BC) {
     // order.
     ForceSequential = true;
     break;
+  case SplitFunctionsStrategy::RandomN:
+    WorkFun = [&](BinaryFunction &BF) {
+      splitFunction(BF, SplitRandomN(RandGen));
+    };
+    ForceSequential = true;
+    break;
   case SplitFunctionsStrategy::All:
     WorkFun = [&](BinaryFunction &BF) { splitFunction<SplitAll>(BF); };
     break;

bolt/test/X86/split-random.s

@@ -2,8 +2,13 @@
 
 # RUN: llvm-mc --filetype=obj --triple x86_64-unknown-unknown %s -o %t.o
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
-# RUN: llvm-bolt %t.exe -o %t.out --split-functions --split-strategy=random2 \
-# RUN:     --print-finalized --print-only=two_block --bolt-seed=7 2>&1 | \
+# RUN: llvm-bolt %t.exe -o %t.random2 --split-functions \
+# RUN:         --split-strategy=random2 --print-finalized \
+# RUN:         --print-only=two_block --bolt-seed=7 2>&1 | \
+# RUN:     FileCheck %s
+# RUN: llvm-bolt %t.exe -o %t.randomN --split-functions \
+# RUN:         --split-strategy=randomN --print-finalized \
+# RUN:         --print-only=two_block --bolt-seed=7 2>&1 | \
 # RUN:     FileCheck %s
 
 # CHECK: Binary Function "two_block"