contrib/llvm/lib/Transforms/Instrumentation/CFGMST.h

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