1 //===-- CFGMST.h - Minimum Spanning Tree for CFG ----------------*- C++ -*-===// 2 // 3 // The LLVM Compiler Infrastructure 4 // 5 // This file is distributed under the University of Illinois Open Source 6 // License. See LICENSE.TXT for details. 7 // 8 //===----------------------------------------------------------------------===// 9 // 10 // This file implements a Union-find algorithm to compute Minimum Spanning Tree 11 // for a given CFG. 12 // 13 //===----------------------------------------------------------------------===// 14 15 #include "llvm/ADT/DenseMap.h" 16 #include "llvm/ADT/STLExtras.h" 17 #include "llvm/Analysis/BlockFrequencyInfo.h" 18 #include "llvm/Analysis/BranchProbabilityInfo.h" 19 #include "llvm/Analysis/CFG.h" 20 #include "llvm/Support/BranchProbability.h" 21 #include "llvm/Support/Debug.h" 22 #include "llvm/Support/raw_ostream.h" 23 #include "llvm/Transforms/Utils/BasicBlockUtils.h" 24 #include <string> 25 #include <utility> 26 #include <vector> 27 28 namespace llvm { 29 30 #define DEBUG_TYPE "cfgmst" 31 32 /// \brief An union-find based Minimum Spanning Tree for CFG 33 /// 34 /// Implements a Union-find algorithm to compute Minimum Spanning Tree 35 /// for a given CFG. 36 template <class Edge, class BBInfo> class CFGMST { 37 public: 38 Function &F; 39 40 // Store all the edges in CFG. It may contain some stale edges 41 // when Removed is set. 42 std::vector<std::unique_ptr<Edge>> AllEdges; 43 44 // This map records the auxiliary information for each BB. 45 DenseMap<const BasicBlock *, std::unique_ptr<BBInfo>> BBInfos; 46 47 // Find the root group of the G and compress the path from G to the root. 48 BBInfo *findAndCompressGroup(BBInfo *G) { 49 if (G->Group != G) 50 G->Group = findAndCompressGroup(static_cast<BBInfo *>(G->Group)); 51 return static_cast<BBInfo *>(G->Group); 52 } 53 54 // Union BB1 and BB2 into the same group and return true. 55 // Returns false if BB1 and BB2 are already in the same group. 56 bool unionGroups(const BasicBlock *BB1, const BasicBlock *BB2) { 57 BBInfo *BB1G = findAndCompressGroup(&getBBInfo(BB1)); 58 BBInfo *BB2G = findAndCompressGroup(&getBBInfo(BB2)); 59 60 if (BB1G == BB2G) 61 return false; 62 63 // Make the smaller rank tree a direct child or the root of high rank tree. 64 if (BB1G->Rank < BB2G->Rank) 65 BB1G->Group = BB2G; 66 else { 67 BB2G->Group = BB1G; 68 // If the ranks are the same, increment root of one tree by one. 69 if (BB1G->Rank == BB2G->Rank) 70 BB1G->Rank++; 71 } 72 return true; 73 } 74 75 // Give BB, return the auxiliary information. 76 BBInfo &getBBInfo(const BasicBlock *BB) const { 77 auto It = BBInfos.find(BB); 78 assert(It->second.get() != nullptr); 79 return *It->second.get(); 80 } 81 82 // Traverse the CFG using a stack. Find all the edges and assign the weight. 83 // Edges with large weight will be put into MST first so they are less likely 84 // to be instrumented. 85 void buildEdges() { 86 DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n"); 87 88 const BasicBlock *BB = &(F.getEntryBlock()); 89 uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2); 90 // Add a fake edge to the entry. 91 addEdge(nullptr, BB, EntryWeight); 92 93 // Special handling for single BB functions. 94 if (succ_empty(BB)) { 95 addEdge(BB, nullptr, EntryWeight); 96 return; 97 } 98 99 static const uint32_t CriticalEdgeMultiplier = 1000; 100 101 for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) { 102 TerminatorInst *TI = BB->getTerminator(); 103 uint64_t BBWeight = 104 (BFI != nullptr ? BFI->getBlockFreq(&*BB).getFrequency() : 2); 105 uint64_t Weight = 2; 106 if (int successors = TI->getNumSuccessors()) { 107 for (int i = 0; i != successors; ++i) { 108 BasicBlock *TargetBB = TI->getSuccessor(i); 109 bool Critical = isCriticalEdge(TI, i); 110 uint64_t scaleFactor = BBWeight; 111 if (Critical) { 112 if (scaleFactor < UINT64_MAX / CriticalEdgeMultiplier) 113 scaleFactor *= CriticalEdgeMultiplier; 114 else 115 scaleFactor = UINT64_MAX; 116 } 117 if (BPI != nullptr) 118 Weight = BPI->getEdgeProbability(&*BB, TargetBB).scale(scaleFactor); 119 addEdge(&*BB, TargetBB, Weight).IsCritical = Critical; 120 DEBUG(dbgs() << " Edge: from " << BB->getName() << " to " 121 << TargetBB->getName() << " w=" << Weight << "\n"); 122 } 123 } else { 124 addEdge(&*BB, nullptr, BBWeight); 125 DEBUG(dbgs() << " Edge: from " << BB->getName() << " to exit" 126 << " w = " << BBWeight << "\n"); 127 } 128 } 129 } 130 131 // Sort CFG edges based on its weight. 132 void sortEdgesByWeight() { 133 std::stable_sort(AllEdges.begin(), AllEdges.end(), 134 [](const std::unique_ptr<Edge> &Edge1, 135 const std::unique_ptr<Edge> &Edge2) { 136 return Edge1->Weight > Edge2->Weight; 137 }); 138 } 139 140 // Traverse all the edges and compute the Minimum Weight Spanning Tree 141 // using union-find algorithm. 142 void computeMinimumSpanningTree() { 143 // First, put all the critical edge with landing-pad as the Dest to MST. 144 // This works around the insufficient support of critical edges split 145 // when destination BB is a landing pad. 146 for (auto &Ei : AllEdges) { 147 if (Ei->Removed) 148 continue; 149 if (Ei->IsCritical) { 150 if (Ei->DestBB && Ei->DestBB->isLandingPad()) { 151 if (unionGroups(Ei->SrcBB, Ei->DestBB)) 152 Ei->InMST = true; 153 } 154 } 155 } 156 157 for (auto &Ei : AllEdges) { 158 if (Ei->Removed) 159 continue; 160 if (unionGroups(Ei->SrcBB, Ei->DestBB)) 161 Ei->InMST = true; 162 } 163 } 164 165 // Dump the Debug information about the instrumentation. 166 void dumpEdges(raw_ostream &OS, const Twine &Message) const { 167 if (!Message.str().empty()) 168 OS << Message << "\n"; 169 OS << " Number of Basic Blocks: " << BBInfos.size() << "\n"; 170 for (auto &BI : BBInfos) { 171 const BasicBlock *BB = BI.first; 172 OS << " BB: " << (BB == nullptr ? "FakeNode" : BB->getName()) << " " 173 << BI.second->infoString() << "\n"; 174 } 175 176 OS << " Number of Edges: " << AllEdges.size() 177 << " (*: Instrument, C: CriticalEdge, -: Removed)\n"; 178 uint32_t Count = 0; 179 for (auto &EI : AllEdges) 180 OS << " Edge " << Count++ << ": " << getBBInfo(EI->SrcBB).Index << "-->" 181 << getBBInfo(EI->DestBB).Index << EI->infoString() << "\n"; 182 } 183 184 // Add an edge to AllEdges with weight W. 185 Edge &addEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W) { 186 uint32_t Index = BBInfos.size(); 187 auto Iter = BBInfos.end(); 188 bool Inserted; 189 std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Src, nullptr)); 190 if (Inserted) { 191 // Newly inserted, update the real info. 192 Iter->second = std::move(llvm::make_unique<BBInfo>(Index)); 193 Index++; 194 } 195 std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Dest, nullptr)); 196 if (Inserted) 197 // Newly inserted, update the real info. 198 Iter->second = std::move(llvm::make_unique<BBInfo>(Index)); 199 AllEdges.emplace_back(new Edge(Src, Dest, W)); 200 return *AllEdges.back(); 201 } 202 203 BranchProbabilityInfo *BPI; 204 BlockFrequencyInfo *BFI; 205 206 public: 207 CFGMST(Function &Func, BranchProbabilityInfo *BPI_ = nullptr, 208 BlockFrequencyInfo *BFI_ = nullptr) 209 : F(Func), BPI(BPI_), BFI(BFI_) { 210 buildEdges(); 211 sortEdgesByWeight(); 212 computeMinimumSpanningTree(); 213 } 214 }; 215 216 #undef DEBUG_TYPE // "cfgmst" 217 } // end namespace llvm 218
