1 //===-- CFGMST.h - Minimum Spanning Tree for CFG ----------------*- C++ -*-===// 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 a Union-find algorithm to compute Minimum Spanning Tree 10 // for a given CFG. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #ifndef LLVM_TRANSFORMS_INSTRUMENTATION_CFGMST_H 15 #define LLVM_TRANSFORMS_INSTRUMENTATION_CFGMST_H 16 17 #include "llvm/ADT/DenseMap.h" 18 #include "llvm/ADT/STLExtras.h" 19 #include "llvm/Analysis/BlockFrequencyInfo.h" 20 #include "llvm/Analysis/BranchProbabilityInfo.h" 21 #include "llvm/Analysis/CFG.h" 22 #include "llvm/Analysis/LoopInfo.h" 23 #include "llvm/IR/Instructions.h" 24 #include "llvm/IR/IntrinsicInst.h" 25 #include "llvm/Support/BranchProbability.h" 26 #include "llvm/Support/Debug.h" 27 #include "llvm/Support/raw_ostream.h" 28 #include "llvm/Transforms/Utils/BasicBlockUtils.h" 29 #include <utility> 30 #include <vector> 31 32 #define DEBUG_TYPE "cfgmst" 33 34 namespace llvm { 35 36 /// An union-find based Minimum Spanning Tree for CFG 37 /// 38 /// Implements a Union-find algorithm to compute Minimum Spanning Tree 39 /// for a given CFG. 40 template <class Edge, class BBInfo> class CFGMST { 41 Function &F; 42 43 // Store all the edges in CFG. It may contain some stale edges 44 // when Removed is set. 45 std::vector<std::unique_ptr<Edge>> AllEdges; 46 47 // This map records the auxiliary information for each BB. 48 DenseMap<const BasicBlock *, std::unique_ptr<BBInfo>> BBInfos; 49 50 // Whehter the function has an exit block with no successors. 51 // (For function with an infinite loop, this block may be absent) 52 bool ExitBlockFound = false; 53 54 BranchProbabilityInfo *const BPI; 55 BlockFrequencyInfo *const BFI; 56 LoopInfo *const LI; 57 58 // If function entry will be always instrumented. 59 const bool InstrumentFuncEntry; 60 61 // If true loop entries will be always instrumented. 62 const bool InstrumentLoopEntries; 63 64 // Find the root group of the G and compress the path from G to the root. 65 BBInfo *findAndCompressGroup(BBInfo *G) { 66 if (G->Group != G) 67 G->Group = findAndCompressGroup(static_cast<BBInfo *>(G->Group)); 68 return static_cast<BBInfo *>(G->Group); 69 } 70 71 // Union BB1 and BB2 into the same group and return true. 72 // Returns false if BB1 and BB2 are already in the same group. 73 bool unionGroups(const BasicBlock *BB1, const BasicBlock *BB2) { 74 BBInfo *BB1G = findAndCompressGroup(&getBBInfo(BB1)); 75 BBInfo *BB2G = findAndCompressGroup(&getBBInfo(BB2)); 76 77 if (BB1G == BB2G) 78 return false; 79 80 // Make the smaller rank tree a direct child or the root of high rank tree. 81 if (BB1G->Rank < BB2G->Rank) 82 BB1G->Group = BB2G; 83 else { 84 BB2G->Group = BB1G; 85 // If the ranks are the same, increment root of one tree by one. 86 if (BB1G->Rank == BB2G->Rank) 87 BB1G->Rank++; 88 } 89 return true; 90 } 91 92 void handleCoroSuspendEdge(Edge *E) { 93 // We must not add instrumentation to the BB representing the 94 // "suspend" path, else CoroSplit won't be able to lower 95 // llvm.coro.suspend to a tail call. We do want profiling info for 96 // the other branches (resume/destroy). So we do 2 things: 97 // 1. we prefer instrumenting those other edges by setting the weight 98 // of the "suspend" edge to max, and 99 // 2. we mark the edge as "Removed" to guarantee it is not considered 100 // for instrumentation. That could technically happen: 101 // (from test/Transforms/Coroutines/coro-split-musttail.ll) 102 // 103 // %suspend = call i8 @llvm.coro.suspend(token %save, i1 false) 104 // switch i8 %suspend, label %exit [ 105 // i8 0, label %await.ready 106 // i8 1, label %exit 107 // ] 108 if (!E->DestBB) 109 return; 110 assert(E->SrcBB); 111 if (llvm::isPresplitCoroSuspendExitEdge(*E->SrcBB, *E->DestBB)) 112 E->Removed = true; 113 } 114 115 // Traverse the CFG using a stack. Find all the edges and assign the weight. 116 // Edges with large weight will be put into MST first so they are less likely 117 // to be instrumented. 118 void buildEdges() { 119 LLVM_DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n"); 120 121 BasicBlock *Entry = &(F.getEntryBlock()); 122 uint64_t EntryWeight = 123 (BFI != nullptr ? BFI->getEntryFreq().getFrequency() : 2); 124 // If we want to instrument the entry count, lower the weight to 0. 125 if (InstrumentFuncEntry) 126 EntryWeight = 0; 127 Edge *EntryIncoming = nullptr, *EntryOutgoing = nullptr, 128 *ExitOutgoing = nullptr, *ExitIncoming = nullptr; 129 uint64_t MaxEntryOutWeight = 0, MaxExitOutWeight = 0, MaxExitInWeight = 0; 130 131 // Add a fake edge to the entry. 132 EntryIncoming = &addEdge(nullptr, Entry, EntryWeight); 133 LLVM_DEBUG(dbgs() << " Edge: from fake node to " << Entry->getName() 134 << " w = " << EntryWeight << "\n"); 135 136 // Special handling for single BB functions. 137 if (succ_empty(Entry)) { 138 addEdge(Entry, nullptr, EntryWeight); 139 return; 140 } 141 142 static const uint32_t CriticalEdgeMultiplier = 1000; 143 144 for (BasicBlock &BB : F) { 145 Instruction *TI = BB.getTerminator(); 146 uint64_t BBWeight = 147 (BFI != nullptr ? BFI->getBlockFreq(&BB).getFrequency() : 2); 148 uint64_t Weight = 2; 149 if (int successors = TI->getNumSuccessors()) { 150 for (int i = 0; i != successors; ++i) { 151 BasicBlock *TargetBB = TI->getSuccessor(i); 152 bool Critical = isCriticalEdge(TI, i); 153 uint64_t scaleFactor = BBWeight; 154 if (Critical) { 155 if (scaleFactor < UINT64_MAX / CriticalEdgeMultiplier) 156 scaleFactor *= CriticalEdgeMultiplier; 157 else 158 scaleFactor = UINT64_MAX; 159 } 160 if (BPI != nullptr) 161 Weight = BPI->getEdgeProbability(&BB, TargetBB).scale(scaleFactor); 162 // If InstrumentLoopEntries is on and the current edge leads to a loop 163 // (i.e., TargetBB is a loop head and BB is outside its loop), set 164 // Weight to be minimal, so that the edge won't be chosen for the MST 165 // and will be instrumented. 166 if (InstrumentLoopEntries && LI->isLoopHeader(TargetBB)) { 167 Loop *TargetLoop = LI->getLoopFor(TargetBB); 168 assert(TargetLoop); 169 if (!TargetLoop->contains(&BB)) 170 Weight = 0; 171 } 172 if (Weight == 0) 173 Weight++; 174 auto *E = &addEdge(&BB, TargetBB, Weight); 175 E->IsCritical = Critical; 176 handleCoroSuspendEdge(E); 177 LLVM_DEBUG(dbgs() << " Edge: from " << BB.getName() << " to " 178 << TargetBB->getName() << " w=" << Weight << "\n"); 179 180 // Keep track of entry/exit edges: 181 if (&BB == Entry) { 182 if (Weight > MaxEntryOutWeight) { 183 MaxEntryOutWeight = Weight; 184 EntryOutgoing = E; 185 } 186 } 187 188 auto *TargetTI = TargetBB->getTerminator(); 189 if (TargetTI && !TargetTI->getNumSuccessors()) { 190 if (Weight > MaxExitInWeight) { 191 MaxExitInWeight = Weight; 192 ExitIncoming = E; 193 } 194 } 195 } 196 } else { 197 ExitBlockFound = true; 198 Edge *ExitO = &addEdge(&BB, nullptr, BBWeight); 199 if (BBWeight > MaxExitOutWeight) { 200 MaxExitOutWeight = BBWeight; 201 ExitOutgoing = ExitO; 202 } 203 LLVM_DEBUG(dbgs() << " Edge: from " << BB.getName() << " to fake exit" 204 << " w = " << BBWeight << "\n"); 205 } 206 } 207 208 // Entry/exit edge adjustment heurisitic: 209 // prefer instrumenting entry edge over exit edge 210 // if possible. Those exit edges may never have a chance to be 211 // executed (for instance the program is an event handling loop) 212 // before the profile is asynchronously dumped. 213 // 214 // If EntryIncoming and ExitOutgoing has similar weight, make sure 215 // ExitOutging is selected as the min-edge. Similarly, if EntryOutgoing 216 // and ExitIncoming has similar weight, make sure ExitIncoming becomes 217 // the min-edge. 218 uint64_t EntryInWeight = EntryWeight; 219 220 if (EntryInWeight >= MaxExitOutWeight && 221 EntryInWeight * 2 < MaxExitOutWeight * 3) { 222 EntryIncoming->Weight = MaxExitOutWeight; 223 ExitOutgoing->Weight = EntryInWeight + 1; 224 } 225 226 if (MaxEntryOutWeight >= MaxExitInWeight && 227 MaxEntryOutWeight * 2 < MaxExitInWeight * 3) { 228 EntryOutgoing->Weight = MaxExitInWeight; 229 ExitIncoming->Weight = MaxEntryOutWeight + 1; 230 } 231 } 232 233 // Sort CFG edges based on its weight. 234 void sortEdgesByWeight() { 235 llvm::stable_sort(AllEdges, [](const std::unique_ptr<Edge> &Edge1, 236 const std::unique_ptr<Edge> &Edge2) { 237 return Edge1->Weight > Edge2->Weight; 238 }); 239 } 240 241 // Traverse all the edges and compute the Minimum Weight Spanning Tree 242 // using union-find algorithm. 243 void computeMinimumSpanningTree() { 244 // First, put all the critical edge with landing-pad as the Dest to MST. 245 // This works around the insufficient support of critical edges split 246 // when destination BB is a landing pad. 247 for (auto &Ei : AllEdges) { 248 if (Ei->Removed) 249 continue; 250 if (Ei->IsCritical) { 251 if (Ei->DestBB && Ei->DestBB->isLandingPad()) { 252 if (unionGroups(Ei->SrcBB, Ei->DestBB)) 253 Ei->InMST = true; 254 } 255 } 256 } 257 258 for (auto &Ei : AllEdges) { 259 if (Ei->Removed) 260 continue; 261 // If we detect infinite loops, force 262 // instrumenting the entry edge: 263 if (!ExitBlockFound && Ei->SrcBB == nullptr) 264 continue; 265 if (unionGroups(Ei->SrcBB, Ei->DestBB)) 266 Ei->InMST = true; 267 } 268 } 269 270 [[maybe_unused]] bool validateLoopEntryInstrumentation() { 271 if (!InstrumentLoopEntries) 272 return true; 273 for (auto &Ei : AllEdges) { 274 if (Ei->Removed) 275 continue; 276 if (Ei->DestBB && LI->isLoopHeader(Ei->DestBB) && 277 !LI->getLoopFor(Ei->DestBB)->contains(Ei->SrcBB) && Ei->InMST) 278 return false; 279 } 280 return true; 281 } 282 283 public: 284 // Dump the Debug information about the instrumentation. 285 void dumpEdges(raw_ostream &OS, const Twine &Message) const { 286 if (!Message.str().empty()) 287 OS << Message << "\n"; 288 OS << " Number of Basic Blocks: " << BBInfos.size() << "\n"; 289 for (auto &BI : BBInfos) { 290 const BasicBlock *BB = BI.first; 291 OS << " BB: " << (BB == nullptr ? "FakeNode" : BB->getName()) << " " 292 << BI.second->infoString() << "\n"; 293 } 294 295 OS << " Number of Edges: " << AllEdges.size() 296 << " (*: Instrument, C: CriticalEdge, -: Removed)\n"; 297 uint32_t Count = 0; 298 for (auto &EI : AllEdges) 299 OS << " Edge " << Count++ << ": " << getBBInfo(EI->SrcBB).Index << "-->" 300 << getBBInfo(EI->DestBB).Index << EI->infoString() << "\n"; 301 } 302 303 // Add an edge to AllEdges with weight W. 304 Edge &addEdge(BasicBlock *Src, BasicBlock *Dest, uint64_t W) { 305 uint32_t Index = BBInfos.size(); 306 auto Iter = BBInfos.end(); 307 bool Inserted; 308 std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Src, nullptr)); 309 if (Inserted) { 310 // Newly inserted, update the real info. 311 Iter->second = std::make_unique<BBInfo>(Index); 312 Index++; 313 } 314 std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Dest, nullptr)); 315 if (Inserted) 316 // Newly inserted, update the real info. 317 Iter->second = std::make_unique<BBInfo>(Index); 318 AllEdges.emplace_back(new Edge(Src, Dest, W)); 319 return *AllEdges.back(); 320 } 321 322 CFGMST(Function &Func, bool InstrumentFuncEntry, bool InstrumentLoopEntries, 323 BranchProbabilityInfo *BPI = nullptr, 324 BlockFrequencyInfo *BFI = nullptr, LoopInfo *LI = nullptr) 325 : F(Func), BPI(BPI), BFI(BFI), LI(LI), 326 InstrumentFuncEntry(InstrumentFuncEntry), 327 InstrumentLoopEntries(InstrumentLoopEntries) { 328 assert(!(InstrumentLoopEntries && !LI) && 329 "expected a LoopInfo to instrumenting loop entries"); 330 buildEdges(); 331 sortEdgesByWeight(); 332 computeMinimumSpanningTree(); 333 assert(validateLoopEntryInstrumentation() && 334 "Loop entries should not be in MST when " 335 "InstrumentLoopEntries is on"); 336 if (AllEdges.size() > 1 && InstrumentFuncEntry) 337 std::iter_swap(std::move(AllEdges.begin()), 338 std::move(AllEdges.begin() + AllEdges.size() - 1)); 339 } 340 341 const std::vector<std::unique_ptr<Edge>> &allEdges() const { 342 return AllEdges; 343 } 344 345 std::vector<std::unique_ptr<Edge>> &allEdges() { return AllEdges; } 346 347 size_t numEdges() const { return AllEdges.size(); } 348 349 size_t bbInfoSize() const { return BBInfos.size(); } 350 351 // Give BB, return the auxiliary information. 352 BBInfo &getBBInfo(const BasicBlock *BB) const { 353 auto It = BBInfos.find(BB); 354 assert(It->second.get() != nullptr); 355 return *It->second.get(); 356 } 357 358 // Give BB, return the auxiliary information if it's available. 359 BBInfo *findBBInfo(const BasicBlock *BB) const { 360 auto It = BBInfos.find(BB); 361 if (It == BBInfos.end()) 362 return nullptr; 363 return It->second.get(); 364 } 365 }; 366 367 } // end namespace llvm 368 369 #undef DEBUG_TYPE // "cfgmst" 370 371 #endif // LLVM_TRANSFORMS_INSTRUMENTATION_CFGMST_H 372