1 //===-- CFGMST.h - Minimum Spanning Tree for CFG ----------------*- C++ -*-===//
3 // The LLVM Compiler Infrastructure
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
8 //===----------------------------------------------------------------------===//
10 // This file implements a Union-find algorithm to compute Minimum Spanning Tree
13 //===----------------------------------------------------------------------===//
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"
29 #define DEBUG_TYPE "cfgmst"
31 /// \brief An union-find based Minimum Spanning Tree for CFG
33 /// Implements a Union-find algorithm to compute Minimum Spanning Tree
35 template <class Edge, class BBInfo> class CFGMST {
39 // Store all the edges in CFG. It may contain some stale edges
40 // when Removed is set.
41 std::vector<std::unique_ptr<Edge>> AllEdges;
43 // This map records the auxiliary information for each BB.
44 DenseMap<const BasicBlock *, std::unique_ptr<BBInfo>> BBInfos;
46 // Find the root group of the G and compress the path from G to the root.
47 BBInfo *findAndCompressGroup(BBInfo *G) {
49 G->Group = findAndCompressGroup(static_cast<BBInfo *>(G->Group));
50 return static_cast<BBInfo *>(G->Group);
53 // Union BB1 and BB2 into the same group and return true.
54 // Returns false if BB1 and BB2 are already in the same group.
55 bool unionGroups(const BasicBlock *BB1, const BasicBlock *BB2) {
56 BBInfo *BB1G = findAndCompressGroup(&getBBInfo(BB1));
57 BBInfo *BB2G = findAndCompressGroup(&getBBInfo(BB2));
62 // Make the smaller rank tree a direct child or the root of high rank tree.
63 if (BB1G->Rank < BB2G->Rank)
67 // If the ranks are the same, increment root of one tree by one.
68 if (BB1G->Rank == BB2G->Rank)
74 // Give BB, return the auxiliary information.
75 BBInfo &getBBInfo(const BasicBlock *BB) const {
76 auto It = BBInfos.find(BB);
77 assert(It->second.get() != nullptr);
78 return *It->second.get();
81 // Give BB, return the auxiliary information if it's available.
82 BBInfo *findBBInfo(const BasicBlock *BB) const {
83 auto It = BBInfos.find(BB);
84 if (It == BBInfos.end())
86 return It->second.get();
89 // Traverse the CFG using a stack. Find all the edges and assign the weight.
90 // Edges with large weight will be put into MST first so they are less likely
91 // to be instrumented.
93 DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n");
95 const BasicBlock *BB = &(F.getEntryBlock());
96 uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2);
97 // Add a fake edge to the entry.
98 addEdge(nullptr, BB, EntryWeight);
100 // Special handling for single BB functions.
101 if (succ_empty(BB)) {
102 addEdge(BB, nullptr, EntryWeight);
106 static const uint32_t CriticalEdgeMultiplier = 1000;
108 for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
109 TerminatorInst *TI = BB->getTerminator();
111 (BFI != nullptr ? BFI->getBlockFreq(&*BB).getFrequency() : 2);
113 if (int successors = TI->getNumSuccessors()) {
114 for (int i = 0; i != successors; ++i) {
115 BasicBlock *TargetBB = TI->getSuccessor(i);
116 bool Critical = isCriticalEdge(TI, i);
117 uint64_t scaleFactor = BBWeight;
119 if (scaleFactor < UINT64_MAX / CriticalEdgeMultiplier)
120 scaleFactor *= CriticalEdgeMultiplier;
122 scaleFactor = UINT64_MAX;
125 Weight = BPI->getEdgeProbability(&*BB, TargetBB).scale(scaleFactor);
126 addEdge(&*BB, TargetBB, Weight).IsCritical = Critical;
127 DEBUG(dbgs() << " Edge: from " << BB->getName() << " to "
128 << TargetBB->getName() << " w=" << Weight << "\n");
131 addEdge(&*BB, nullptr, BBWeight);
132 DEBUG(dbgs() << " Edge: from " << BB->getName() << " to exit"
133 << " w = " << BBWeight << "\n");
138 // Sort CFG edges based on its weight.
139 void sortEdgesByWeight() {
140 std::stable_sort(AllEdges.begin(), AllEdges.end(),
141 [](const std::unique_ptr<Edge> &Edge1,
142 const std::unique_ptr<Edge> &Edge2) {
143 return Edge1->Weight > Edge2->Weight;
147 // Traverse all the edges and compute the Minimum Weight Spanning Tree
148 // using union-find algorithm.
149 void computeMinimumSpanningTree() {
150 // First, put all the critical edge with landing-pad as the Dest to MST.
151 // This works around the insufficient support of critical edges split
152 // when destination BB is a landing pad.
153 for (auto &Ei : AllEdges) {
156 if (Ei->IsCritical) {
157 if (Ei->DestBB && Ei->DestBB->isLandingPad()) {
158 if (unionGroups(Ei->SrcBB, Ei->DestBB))
164 for (auto &Ei : AllEdges) {
167 if (unionGroups(Ei->SrcBB, Ei->DestBB))
172 // Dump the Debug information about the instrumentation.
173 void dumpEdges(raw_ostream &OS, const Twine &Message) const {
174 if (!Message.str().empty())
175 OS << Message << "\n";
176 OS << " Number of Basic Blocks: " << BBInfos.size() << "\n";
177 for (auto &BI : BBInfos) {
178 const BasicBlock *BB = BI.first;
179 OS << " BB: " << (BB == nullptr ? "FakeNode" : BB->getName()) << " "
180 << BI.second->infoString() << "\n";
183 OS << " Number of Edges: " << AllEdges.size()
184 << " (*: Instrument, C: CriticalEdge, -: Removed)\n";
186 for (auto &EI : AllEdges)
187 OS << " Edge " << Count++ << ": " << getBBInfo(EI->SrcBB).Index << "-->"
188 << getBBInfo(EI->DestBB).Index << EI->infoString() << "\n";
191 // Add an edge to AllEdges with weight W.
192 Edge &addEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W) {
193 uint32_t Index = BBInfos.size();
194 auto Iter = BBInfos.end();
196 std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Src, nullptr));
198 // Newly inserted, update the real info.
199 Iter->second = std::move(llvm::make_unique<BBInfo>(Index));
202 std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Dest, nullptr));
204 // Newly inserted, update the real info.
205 Iter->second = std::move(llvm::make_unique<BBInfo>(Index));
206 AllEdges.emplace_back(new Edge(Src, Dest, W));
207 return *AllEdges.back();
210 BranchProbabilityInfo *BPI;
211 BlockFrequencyInfo *BFI;
214 CFGMST(Function &Func, BranchProbabilityInfo *BPI_ = nullptr,
215 BlockFrequencyInfo *BFI_ = nullptr)
216 : F(Func), BPI(BPI_), BFI(BFI_) {
219 computeMinimumSpanningTree();
223 #undef DEBUG_TYPE // "cfgmst"
224 } // end namespace llvm