llvm/lib/Support/BalancedPartitioning.cpp

   1 //===- BalancedPartitioning.cpp -------------------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements BalancedPartitioning, a recursive balanced graph
  10 // partitioning algorithm.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/Support/BalancedPartitioning.h"
  15 #include "llvm/Support/Debug.h"
  16 #include "llvm/Support/Format.h"
  17 #include "llvm/Support/FormatVariadic.h"
  18 #include "llvm/Support/ThreadPool.h"
  19
  20 using namespace llvm;
  21 #define DEBUG_TYPE "balanced-partitioning"
  22
  23 void BPFunctionNode::dump(raw_ostream &OS) const {
  24   OS << formatv("{{ID={0} Utilities={{{1:$[,]}} Bucket={2}}", Id,
  25                 make_range(UtilityNodes.begin(), UtilityNodes.end()), Bucket);
  26 }
  27
  28 template <typename Func>
  29 void BalancedPartitioning::BPThreadPool::async(Func &&F) {
  30 #if LLVM_ENABLE_THREADS
  31   // This new thread could spawn more threads, so mark it as active
  32   ++NumActiveThreads;
  33   TheThreadPool.async([=]() {
  34     // Run the task
  35     F();
  36
  37     // This thread will no longer spawn new threads, so mark it as inactive
  38     if (--NumActiveThreads == 0) {
  39       // There are no more active threads, so mark as finished and notify
  40       {
  41         std::unique_lock<std::mutex> lock(mtx);
  42         assert(!IsFinishedSpawning);
  43         IsFinishedSpawning = true;
  44       }
  45       cv.notify_one();
  46     }
  47   });
  48 #else
  49   llvm_unreachable("threads are disabled");
  50 #endif
  51 }
  52
  53 void BalancedPartitioning::BPThreadPool::wait() {
  54 #if LLVM_ENABLE_THREADS
  55   // TODO: We could remove the mutex and condition variable and use
  56   // std::atomic::wait() instead, but that isn't available until C++20
  57   {
  58     std::unique_lock<std::mutex> lock(mtx);
  59     cv.wait(lock, [&]() { return IsFinishedSpawning; });
  60     assert(IsFinishedSpawning && NumActiveThreads == 0);
  61   }
  62   // Now we can call ThreadPool::wait() since all tasks have been submitted
  63   TheThreadPool.wait();
  64 #else
  65   llvm_unreachable("threads are disabled");
  66 #endif
  67 }
  68
  69 BalancedPartitioning::BalancedPartitioning(
  70     const BalancedPartitioningConfig &Config)
  71     : Config(Config) {
  72   // Pre-computing log2 values
  73   Log2Cache[0] = 0.0;
  74   for (unsigned I = 1; I < LOG_CACHE_SIZE; I++)
  75     Log2Cache[I] = std::log2(I);
  76 }
  77
  78 void BalancedPartitioning::run(std::vector<BPFunctionNode> &Nodes) const {
  79   LLVM_DEBUG(
  80       dbgs() << format(
  81           "Partitioning %d nodes using depth %d and %d iterations per split\n",
  82           Nodes.size(), Config.SplitDepth, Config.IterationsPerSplit));
  83   std::optional<BPThreadPool> TP;
  84 #if LLVM_ENABLE_THREADS
  85   ThreadPool TheThreadPool;
  86   if (Config.TaskSplitDepth > 1)
  87     TP.emplace(TheThreadPool);
  88 #endif
  89
  90   // Record the input order
  91   for (unsigned I = 0; I < Nodes.size(); I++)
  92     Nodes[I].InputOrderIndex = I;
  93
  94   auto NodesRange = llvm::make_range(Nodes.begin(), Nodes.end());
  95   auto BisectTask = [=, &TP]() {
  96     bisect(NodesRange, /*RecDepth=*/0, /*RootBucket=*/1, /*Offset=*/0, TP);
  97   };
  98   if (TP) {
  99     TP->async(std::move(BisectTask));
 100     TP->wait();
 101   } else {
 102     BisectTask();
 103   }
 104
 105   llvm::stable_sort(NodesRange, [](const auto &L, const auto &R) {
 106     return L.Bucket < R.Bucket;
 107   });
 108
 109   LLVM_DEBUG(dbgs() << "Balanced partitioning completed\n");
 110 }
 111
 112 void BalancedPartitioning::bisect(const FunctionNodeRange Nodes,
 113                                   unsigned RecDepth, unsigned RootBucket,
 114                                   unsigned Offset,
 115                                   std::optional<BPThreadPool> &TP) const {
 116   unsigned NumNodes = std::distance(Nodes.begin(), Nodes.end());
 117   if (NumNodes <= 1 || RecDepth >= Config.SplitDepth) {
 118     // We've reach the lowest level of the recursion tree. Fall back to the
 119     // original order and assign to buckets.
 120     llvm::sort(Nodes, [](const auto &L, const auto &R) {
 121       return L.InputOrderIndex < R.InputOrderIndex;
 122     });
 123     for (auto &N : Nodes)
 124       N.Bucket = Offset++;
 125     return;
 126   }
 127
 128   LLVM_DEBUG(dbgs() << format("Bisect with %d nodes and root bucket %d\n",
 129                               NumNodes, RootBucket));
 130
 131   std::mt19937 RNG(RootBucket);
 132
 133   unsigned LeftBucket = 2 * RootBucket;
 134   unsigned RightBucket = 2 * RootBucket + 1;
 135
 136   // Split into two and assign to the left and right buckets
 137   split(Nodes, LeftBucket);
 138
 139   runIterations(Nodes, RecDepth, LeftBucket, RightBucket, RNG);
 140
 141   // Split nodes wrt the resulting buckets
 142   auto NodesMid =
 143       llvm::partition(Nodes, [&](auto &N) { return N.Bucket == LeftBucket; });
 144   unsigned MidOffset = Offset + std::distance(Nodes.begin(), NodesMid);
 145
 146   auto LeftNodes = llvm::make_range(Nodes.begin(), NodesMid);
 147   auto RightNodes = llvm::make_range(NodesMid, Nodes.end());
 148
 149   auto LeftRecTask = [=, &TP]() {
 150     bisect(LeftNodes, RecDepth + 1, LeftBucket, Offset, TP);
 151   };
 152   auto RightRecTask = [=, &TP]() {
 153     bisect(RightNodes, RecDepth + 1, RightBucket, MidOffset, TP);
 154   };
 155
 156   if (TP && RecDepth < Config.TaskSplitDepth && NumNodes >= 4) {
 157     TP->async(std::move(LeftRecTask));
 158     TP->async(std::move(RightRecTask));
 159   } else {
 160     LeftRecTask();
 161     RightRecTask();
 162   }
 163 }
 164
 165 void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
 166                                          unsigned RecDepth, unsigned LeftBucket,
 167                                          unsigned RightBucket,
 168                                          std::mt19937 &RNG) const {
 169   unsigned NumNodes = std::distance(Nodes.begin(), Nodes.end());
 170   DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeIndex;
 171   for (auto &N : Nodes)
 172     for (auto &UN : N.UtilityNodes)
 173       ++UtilityNodeIndex[UN];
 174   // Remove utility nodes if they have just one edge or are connected to all
 175   // functions
 176   for (auto &N : Nodes)
 177     llvm::erase_if(N.UtilityNodes, [&](auto &UN) {
 178       return UtilityNodeIndex[UN] == 1 || UtilityNodeIndex[UN] == NumNodes;
 179     });
 180
 181   // Renumber utility nodes so they can be used to index into Signatures
 182   UtilityNodeIndex.clear();
 183   for (auto &N : Nodes)
 184     for (auto &UN : N.UtilityNodes)
 185       UN = UtilityNodeIndex.insert({UN, UtilityNodeIndex.size()}).first->second;
 186
 187   // Initialize signatures
 188   SignaturesT Signatures(/*Size=*/UtilityNodeIndex.size());
 189   for (auto &N : Nodes) {
 190     for (auto &UN : N.UtilityNodes) {
 191       assert(UN < Signatures.size());
 192       if (N.Bucket == LeftBucket) {
 193         Signatures[UN].LeftCount++;
 194       } else {
 195         Signatures[UN].RightCount++;
 196       }
 197     }
 198   }
 199
 200   for (unsigned I = 0; I < Config.IterationsPerSplit; I++) {
 201     unsigned NumMovedNodes =
 202         runIteration(Nodes, LeftBucket, RightBucket, Signatures, RNG);
 203     if (NumMovedNodes == 0)
 204       break;
 205   }
 206 }
 207
 208 unsigned BalancedPartitioning::runIteration(const FunctionNodeRange Nodes,
 209                                             unsigned LeftBucket,
 210                                             unsigned RightBucket,
 211                                             SignaturesT &Signatures,
 212                                             std::mt19937 &RNG) const {
 213   // Init signature cost caches
 214   for (auto &Signature : Signatures) {
 215     if (Signature.CachedGainIsValid)
 216       continue;
 217     unsigned L = Signature.LeftCount;
 218     unsigned R = Signature.RightCount;
 219     assert((L > 0 || R > 0) && "incorrect signature");
 220     float Cost = logCost(L, R);
 221     Signature.CachedGainLR = 0.f;
 222     Signature.CachedGainRL = 0.f;
 223     if (L > 0)
 224       Signature.CachedGainLR = Cost - logCost(L - 1, R + 1);
 225     if (R > 0)
 226       Signature.CachedGainRL = Cost - logCost(L + 1, R - 1);
 227     Signature.CachedGainIsValid = true;
 228   }
 229
 230   // Compute move gains
 231   typedef std::pair<float, BPFunctionNode *> GainPair;
 232   std::vector<GainPair> Gains;
 233   for (auto &N : Nodes) {
 234     bool FromLeftToRight = (N.Bucket == LeftBucket);
 235     float Gain = moveGain(N, FromLeftToRight, Signatures);
 236     Gains.push_back(std::make_pair(Gain, &N));
 237   }
 238
 239   // Collect left and right gains
 240   auto LeftEnd = llvm::partition(
 241       Gains, [&](const auto &GP) { return GP.second->Bucket == LeftBucket; });
 242   auto LeftRange = llvm::make_range(Gains.begin(), LeftEnd);
 243   auto RightRange = llvm::make_range(LeftEnd, Gains.end());
 244
 245   // Sort gains in descending order
 246   auto LargerGain = [](const auto &L, const auto &R) {
 247     return L.first > R.first;
 248   };
 249   llvm::stable_sort(LeftRange, LargerGain);
 250   llvm::stable_sort(RightRange, LargerGain);
 251
 252   unsigned NumMovedDataVertices = 0;
 253   for (auto [LeftPair, RightPair] : llvm::zip(LeftRange, RightRange)) {
 254     auto &[LeftGain, LeftNode] = LeftPair;
 255     auto &[RightGain, RightNode] = RightPair;
 256     // Stop when the gain is no longer beneficial
 257     if (LeftGain + RightGain <= 0.f)
 258       break;
 259     // Try to exchange the nodes between buckets
 260     if (moveFunctionNode(*LeftNode, LeftBucket, RightBucket, Signatures, RNG))
 261       ++NumMovedDataVertices;
 262     if (moveFunctionNode(*RightNode, LeftBucket, RightBucket, Signatures, RNG))
 263       ++NumMovedDataVertices;
 264   }
 265   return NumMovedDataVertices;
 266 }
 267
 268 bool BalancedPartitioning::moveFunctionNode(BPFunctionNode &N,
 269                                             unsigned LeftBucket,
 270                                             unsigned RightBucket,
 271                                             SignaturesT &Signatures,
 272                                             std::mt19937 &RNG) const {
 273   // Sometimes we skip the move. This helps to escape local optima
 274   if (std::uniform_real_distribution<float>(0.f, 1.f)(RNG) <=
 275       Config.SkipProbability)
 276     return false;
 277
 278   bool FromLeftToRight = (N.Bucket == LeftBucket);
 279   // Update the current bucket
 280   N.Bucket = (FromLeftToRight ? RightBucket : LeftBucket);
 281
 282   // Update signatures and invalidate gain cache
 283   if (FromLeftToRight) {
 284     for (auto &UN : N.UtilityNodes) {
 285       auto &Signature = Signatures[UN];
 286       Signature.LeftCount--;
 287       Signature.RightCount++;
 288       Signature.CachedGainIsValid = false;
 289     }
 290   } else {
 291     for (auto &UN : N.UtilityNodes) {
 292       auto &Signature = Signatures[UN];
 293       Signature.LeftCount++;
 294       Signature.RightCount--;
 295       Signature.CachedGainIsValid = false;
 296     }
 297   }
 298   return true;
 299 }
 300
 301 void BalancedPartitioning::split(const FunctionNodeRange Nodes,
 302                                  unsigned StartBucket) const {
 303   unsigned NumNodes = std::distance(Nodes.begin(), Nodes.end());
 304   auto NodesMid = Nodes.begin() + (NumNodes + 1) / 2;
 305
 306   std::nth_element(Nodes.begin(), NodesMid, Nodes.end(), [](auto &L, auto &R) {
 307     return L.InputOrderIndex < R.InputOrderIndex;
 308   });
 309
 310   for (auto &N : llvm::make_range(Nodes.begin(), NodesMid))
 311     N.Bucket = StartBucket;
 312   for (auto &N : llvm::make_range(NodesMid, Nodes.end()))
 313     N.Bucket = StartBucket + 1;
 314 }
 315
 316 float BalancedPartitioning::moveGain(const BPFunctionNode &N,
 317                                      bool FromLeftToRight,
 318                                      const SignaturesT &Signatures) {
 319   float Gain = 0.f;
 320   for (auto &UN : N.UtilityNodes)
 321     Gain += (FromLeftToRight ? Signatures[UN].CachedGainLR
 322                              : Signatures[UN].CachedGainRL);
 323   return Gain;
 324 }
 325
 326 float BalancedPartitioning::logCost(unsigned X, unsigned Y) const {
 327   return -(X * log2Cached(X + 1) + Y * log2Cached(Y + 1));
 328 }
 329
 330 float BalancedPartitioning::log2Cached(unsigned i) const {
 331   return (i < LOG_CACHE_SIZE) ? Log2Cache[i] : std::log2(i);
 332 }