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