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