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
28namespace 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.
36template <class Edge, class BBInfo> class CFGMST {
37public:
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
206public:
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