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