Extending CFGPrinter and CallPrinter with Heat Colors

[rcorcs]

Summary

Extends the CFGPrinter and CallPrinter with heat colors based on heuristics or profiling information.
The colors can be enabled for CFGPrinter by using the option -cfg-heat-colors for both -dot-cfg[-only] and -view-cfg[-only].
Similarly, the colors can be enabled for CallPrinter by using the option -callgraph-heat-colors for both -dot-callgraph and -view-callgraph.

[lebedev.ri]

Hi, interesting stuff, some thoughts on the differential itself:

1. you did not set the repo. rLLVM should have been specified.
2. you did not put llvm-commits as subscriber. there are talks about automating this, but if the repo is not specified, it could not even be auto-detected.
3. please upload the diff with full context (-U999999).

Thus, you probably want to close this differential, and open a new one with those ^ three issues not present

[rcorcs]

Closed due to incorrect configuration.
Another revision will be opened with the correct configuration.

include/llvm/Analysis/CFGPrinter.h

 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/GraphWriter.h"
 
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/HeatUtils.h"
+
+#include <sstream>
+
 namespace llvm {
 class CFGViewerPass
     : public PassInfoMixin<CFGViewerPass> {

   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
+class CFGDOTInfo {
+private:
+   const BlockFrequencyInfo *BFI;
+   const Function *F;
+   uint64_t MaxFreq;
+   bool ShowHeat;
+   bool Heuristic;
+   bool EdgeWeights;
+   bool RawWeights;
+public:
+   CFGDOTInfo(const Function *F) : CFGDOTInfo(F,nullptr,0){}
+
+   CFGDOTInfo(const Function *F, const BlockFrequencyInfo *BFI, uint64_t MaxFreq){
+      this->BFI = BFI;
+      this->F = F;
+      this->MaxFreq = MaxFreq;
+      this->ShowHeat = true;
+      this->Heuristic = true;
+      this->EdgeWeights = true;
+      this->RawWeights = false;
+   }
+
+   const BlockFrequencyInfo *getBFI(){ return BFI; }
+ 
+   
+   const Function *getF(){ return this->F; }
+
+   uint64_t getMaxFreq() { return MaxFreq; }
+
+   uint64_t getFreq(const BasicBlock *BB){
+      return getBlockFreq(BB,BFI,Heuristic);
+   }
+
+   void setHeatColors(bool ShowHeat){
+      this->ShowHeat = ShowHeat;
+   }
+
+   bool showHeatColors(){
+      return ShowHeat;
+   }
+
+   void setHeuristic(bool Heuristic){
+      this->Heuristic = Heuristic;
+   }
+
+   bool useHeuristic(){
+      return Heuristic;
+   }
+ 
+   void setRawEdgeWeights(bool RawWeights){
+      this->RawWeights = RawWeights;
+   }
+
+   bool useRawEdgeWeights(){
+      return RawWeights;
+   }
+
+   void setEdgeWeights(bool EdgeWeights){
+      this->EdgeWeights = EdgeWeights;
+   }
+
+   bool showEdgeWeights(){
+      return EdgeWeights;
+   }
+};
+
+
+template <> struct GraphTraits<CFGDOTInfo *> :
+  public GraphTraits<const BasicBlock*> {
+  static NodeRef getEntryNode(CFGDOTInfo *CFGInfo) {
+    return &(CFGInfo->getF()->getEntryBlock());
+  }
+
+  // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
+  using nodes_iterator = pointer_iterator<Function::const_iterator>;
+
+  static nodes_iterator nodes_begin(CFGDOTInfo *CFGInfo) {
+    return nodes_iterator(CFGInfo->getF()->begin());
+  }
+
+  static nodes_iterator nodes_end(CFGDOTInfo *CFGInfo) {
+    return nodes_iterator(CFGInfo->getF()->end());
+  }
+
+  static size_t size(CFGDOTInfo *CFGInfo) { return CFGInfo->getF()->size(); }
+};
+
 template<>
-struct DOTGraphTraits<const Function*> : public DefaultDOTGraphTraits {
+struct DOTGraphTraits<CFGDOTInfo *> : public DefaultDOTGraphTraits {
 
   DOTGraphTraits (bool isSimple=false) : DefaultDOTGraphTraits(isSimple) {}
 
-  static std::string getGraphName(const Function *F) {
-    return "CFG for '" + F->getName().str() + "' function";
+  static std::string getGraphName(CFGDOTInfo *CFGInfo) {
+    return "CFG for '" + CFGInfo->getF()->getName().str() + "' function";
   }
 
   static std::string getSimpleNodeLabel(const BasicBlock *Node,
-                                        const Function *) {
+                                        CFGDOTInfo *) {
     if (!Node->getName().empty())
       return Node->getName().str();
 

   }
 
   static std::string getCompleteNodeLabel(const BasicBlock *Node,
-                                          const Function *) {
+                                          CFGDOTInfo *) {
     enum { MaxColumns = 80 };
     std::string Str;
     raw_string_ostream OS(Str);

     return OutStr;
   }
 
-  std::string getNodeLabel(const BasicBlock *Node,
-                           const Function *Graph) {
+  std::string getNodeLabel(const BasicBlock *Node, CFGDOTInfo *CFGInfo) {
     if (isSimple())
-      return getSimpleNodeLabel(Node, Graph);
+      return getSimpleNodeLabel(Node, CFGInfo);
     else
-      return getCompleteNodeLabel(Node, Graph);
+      return getCompleteNodeLabel(Node, CFGInfo);
   }
 
   static std::string getEdgeSourceLabel(const BasicBlock *Node,

 
   /// Display the raw branch weights from PGO.
   std::string getEdgeAttributes(const BasicBlock *Node, succ_const_iterator I,
-                                const Function *F) {
+                                CFGDOTInfo *CFGInfo) {
+
+    if (!CFGInfo->showEdgeWeights())
+      return "";
+
     const TerminatorInst *TI = Node->getTerminator();
     if (TI->getNumSuccessors() == 1)
       return "";
 
-    MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
-    if (!WeightsNode)
-      return "";
+    std::string Attrs = "";
 
-    MDString *MDName = cast<MDString>(WeightsNode->getOperand(0));
-    if (MDName->getString() != "branch_weights")
-      return "";
+    if (CFGInfo->useRawEdgeWeights()) {
+       MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
+       if (!WeightsNode)
+         return "";
 
-    unsigned OpNo = I.getSuccessorIndex() + 1;
-    if (OpNo >= WeightsNode->getNumOperands())
+       MDString *MDName = cast<MDString>(WeightsNode->getOperand(0));
+       if (MDName->getString() != "branch_weights")
          return "";
-    ConstantInt *Weight =
-        mdconst::dyn_extract<ConstantInt>(WeightsNode->getOperand(OpNo));
-    if (!Weight)
-      return "";
 
-    // Prepend a 'W' to indicate that this is a weight rather than the actual
-    // profile count (due to scaling).
-    Twine Attrs = "label=\"W:" + Twine(Weight->getZExtValue()) + "\"";
-    return Attrs.str();
+       unsigned OpNo = I.getSuccessorIndex() + 1;
+       if (OpNo >= WeightsNode->getNumOperands())
+         return "";
+       ConstantInt *Weight =
+           mdconst::dyn_extract<ConstantInt>(WeightsNode->getOperand(OpNo));
+       if (!Weight)
+         return "";
+
+       // Prepend a 'W' to indicate that this is a weight rather than the actual
+       // profile count (due to scaling).
+       Attrs = "label=\"W:" + std::to_string(Weight->getZExtValue()) + "\"";
+    } else {
+       uint64_t total = 0;
+       for (unsigned i = 0; i<TI->getNumSuccessors(); i++){
+          total += CFGInfo->getFreq(TI->getSuccessor(i));
+       }
+
+       unsigned OpNo = I.getSuccessorIndex();
+
+       if (OpNo >= TI->getNumSuccessors())
+         return "";
+
+       BasicBlock *SuccBB = TI->getSuccessor(OpNo);
+
+       double val = 0.0;
+       if (CFGInfo->getFreq(SuccBB)>0) {
+         double freq = CFGInfo->getFreq(SuccBB);
+         val = (int(round((freq/double(total))*10000)))/100.0;
+       }
+
+       std::stringstream ss;
+       ss.precision(2);
+       ss << std::fixed << val;
+       Attrs = "label=\"" + ss.str() + "%\"";
+    }
+    return Attrs;
   }
+
+  std::string getNodeAttributes(const BasicBlock *Node, CFGDOTInfo *CFGInfo) {
+    std::string attrs = "";
+    if(CFGInfo->showHeatColors()){
+       uint64_t freq = CFGInfo->getFreq(Node);
+       std::string color = getHeatColor(freq, CFGInfo->getMaxFreq());
+       std::string edgeColor = (freq<=(CFGInfo->getMaxFreq()/2))?
+                            (getHeatColor(0)):(getHeatColor(1));
+
+       attrs = "color=\"" + edgeColor +
+                        "ff\", style=filled, fillcolor=\"" + color + "70\"";
+
+    }
+    return attrs;
+  }  
 };
+
 } // End llvm namespace
 
 namespace llvm {

include/llvm/Analysis/HeatUtils.h

+//===-- HeatUtils.h - Utility for printing heat colors ----------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Utility for printing heat colors based on heuristics or profiling
+// information.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_HEATUTILS_H
+#define LLVM_ANALYSIS_HEATUTILS_H
+
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+
+#include <string>
+
+using namespace llvm;
+
+namespace llvm {
+
+bool hasProfiling(Module &M);
+
+uint64_t getBlockFreq(const BasicBlock *BB, const BlockFrequencyInfo *BFI,
+                      bool useHeuristic=true);
+
+uint64_t getNumOfCalls(Function &callerFunction, Function &calledFunction,
+                       function_ref<BlockFrequencyInfo *(Function &)> LookupBFI,
+                       bool useHeuristic=true);
+
+uint64_t getMaxFreq(const Function &F, const BlockFrequencyInfo *BFI,
+                    bool useHeuristic=true);
+
+uint64_t getMaxFreq(Module &M,
+                    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI,
+                    bool useHeuristic=true);
+
+std::string getHeatColor(uint64_t freq, uint64_t maxFreq);
+
+std::string getHeatColor(double percent);
+
+}
+
+#endif

lib/Analysis/CFGPrinter.cpp

 #include "llvm/Analysis/CFGPrinter.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/FileSystem.h"
+
+
 using namespace llvm;
 
+static cl::opt<bool>
+ShowHeatColors("cfg-heat-colors", cl::init(false), cl::Hidden,
+                   cl::desc("Show heat colors in CFG"));
+
+static cl::opt<bool>
+UseRawEdgeWeight("cfg-raw-weights", cl::init(false), cl::Hidden,
+                   cl::desc("Use raw profiling weights"));
+
+static cl::opt<bool>
+ShowEdgeWeight("cfg-weights", cl::init(false), cl::Hidden,
+                   cl::desc("Show edge labels with weights"));
+
+
+static void writeHeatCFGToDotFile(Function &F, BlockFrequencyInfo *BFI,
+                           uint64_t MaxFreq, bool UseHeuristic, bool isSimple) {
+  std::string Filename = ("cfg." + F.getName() + ".dot").str();
+  errs() << "Writing '" << Filename << "'...";
+
+  std::error_code EC;
+  raw_fd_ostream File(Filename, EC, sys::fs::F_Text);
+
+  CFGDOTInfo CFGInfo(&F,BFI,MaxFreq);
+  CFGInfo.setHeuristic(UseHeuristic);
+  CFGInfo.setHeatColors(ShowHeatColors);
+  CFGInfo.setEdgeWeights(ShowEdgeWeight);
+  CFGInfo.setRawEdgeWeights(UseRawEdgeWeight);
+
+  if (!EC)
+     WriteGraph(File, &CFGInfo, isSimple);
+  else
+     errs() << "  error opening file for writing!";
+  errs() << "\n";
+}
+
+static void viewHeatCFGToDotFile(Function &F, BlockFrequencyInfo *BFI,
+                           uint64_t MaxFreq, bool UseHeuristic, bool isSimple) {
+  CFGDOTInfo CFGInfo(&F,BFI,MaxFreq);
+  CFGInfo.setHeuristic(UseHeuristic);
+  CFGInfo.setHeatColors(ShowHeatColors);
+  CFGInfo.setEdgeWeights(ShowEdgeWeight);
+  CFGInfo.setRawEdgeWeights(UseRawEdgeWeight);
+
+  ViewGraph(&CFGInfo, "cfg." + F.getName(), isSimple);
+}
+
 namespace {
+
   struct CFGViewerLegacyPass : public FunctionPass {
     static char ID; // Pass identifcation, replacement for typeid
     CFGViewerLegacyPass() : FunctionPass(ID) {

     }
 
     bool runOnFunction(Function &F) override {
-      F.viewCFG();
+      auto *BFI = &getAnalysis<BlockFrequencyInfoWrapperPass>().getBFI();
+      bool UseHeuristic = true;
+      uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+      viewHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/false);
+      //F.viewCFG();
       return false;
     }
 

     void print(raw_ostream &OS, const Module* = nullptr) const override {}
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<BlockFrequencyInfoWrapperPass>();
       AU.setPreservesAll();
     }
   };

 
 PreservedAnalyses CFGViewerPass::run(Function &F,
                                      FunctionAnalysisManager &AM) {
-  F.viewCFG();
+  auto *BFI = &AM.getResult<BlockFrequencyAnalysis>(F);
+  bool UseHeuristic = true;
+  uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+  viewHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/false);
+  //F.viewCFG();
   return PreservedAnalyses::all();
 }
 

     }
 
     bool runOnFunction(Function &F) override {
-      F.viewCFGOnly();
+      auto *BFI = &getAnalysis<BlockFrequencyInfoWrapperPass>().getBFI();
+      bool UseHeuristic = true;
+      uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+      viewHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/true);
+      //F.viewCFGOnly();
       return false;
     }
 

     void print(raw_ostream &OS, const Module* = nullptr) const override {}
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<BlockFrequencyInfoWrapperPass>();
       AU.setPreservesAll();
     }
   };

 
 PreservedAnalyses CFGOnlyViewerPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
-  F.viewCFGOnly();
+
+  auto *BFI = &AM.getResult<BlockFrequencyAnalysis>(F);
+  bool UseHeuristic = true;
+  uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+  viewHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/true);
+  //F.viewCFGOnly();
   return PreservedAnalyses::all();
 }
 
-static void writeCFGToDotFile(Function &F) {
-  std::string Filename = ("cfg." + F.getName() + ".dot").str();
-  errs() << "Writing '" << Filename << "'...";
-
-  std::error_code EC;
-  raw_fd_ostream File(Filename, EC, sys::fs::F_Text);
-
-  if (!EC)
-    WriteGraph(File, (const Function*)&F);
-  else
-    errs() << "  error opening file for writing!";
-  errs() << "\n";
-}
-
 namespace {
   struct CFGPrinterLegacyPass : public FunctionPass {
     static char ID; // Pass identification, replacement for typeid

     }
 
     bool runOnFunction(Function &F) override {
-      writeCFGToDotFile(F);
+      auto *BFI = &getAnalysis<BlockFrequencyInfoWrapperPass>().getBFI();
+      bool UseHeuristic = true;
+      uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+      writeHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/false);
       return false;
     }
 

     void print(raw_ostream &OS, const Module* = nullptr) const override {}
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<BlockFrequencyInfoWrapperPass>();
       AU.setPreservesAll();
     }
   };

 
 PreservedAnalyses CFGPrinterPass::run(Function &F,
                                       FunctionAnalysisManager &AM) {
-  writeCFGToDotFile(F);
+  auto *BFI = &AM.getResult<BlockFrequencyAnalysis>(F);
+  bool UseHeuristic = true;
+  uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+  writeHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/false);
   return PreservedAnalyses::all();
 }
 

     }
 
     bool runOnFunction(Function &F) override {
-      writeCFGToDotFile(F);
+      auto *BFI = &this->getAnalysis<BlockFrequencyInfoWrapperPass>().getBFI();
+
+      bool UseHeuristic = true;
+      uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+      writeHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/true);
+
       return false;
     }
     void print(raw_ostream &OS, const Module* = nullptr) const override {}
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<BlockFrequencyInfoWrapperPass>();
       AU.setPreservesAll();
     }
   };

 INITIALIZE_PASS(CFGOnlyPrinterLegacyPass, "dot-cfg-only",
    "Print CFG of function to 'dot' file (with no function bodies)",
    false, true)
+/*
+INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
+INITIALIZE_PASS_END(CFGOnlyPrinterLegacyPass, "dot-cfg-only",
+   "Print CFG of function to 'dot' file (with no function bodies)",
+   false, true)
+*/
 
 PreservedAnalyses CFGOnlyPrinterPass::run(Function &F,
                                           FunctionAnalysisManager &AM) {
-  writeCFGToDotFile(F);
+  auto *BFI = &AM.getResult<BlockFrequencyAnalysis>(F);
+  bool UseHeuristic = true;
+  uint64_t MaxFreq = getMaxFreq(F,BFI,UseHeuristic);
+  writeHeatCFGToDotFile(F,BFI,MaxFreq,UseHeuristic,/*isSimple=*/true);
   return PreservedAnalyses::all();
 }
 

 /// being a 'dot' and 'gv' program in your path.
 ///
 void Function::viewCFG() const {
-  ViewGraph(this, "cfg" + getName());
+
+  CFGDOTInfo CFGInfo(this,nullptr,0);
+  CFGInfo.setHeuristic(false);
+  CFGInfo.setHeatColors(false);
+  CFGInfo.setEdgeWeights(true);
+  CFGInfo.setRawEdgeWeights(true);
+
+  ViewGraph(&CFGInfo, "cfg" + getName());
 }
 
 /// viewCFGOnly - This function is meant for use from the debugger.  It works

 /// this can make the graph smaller.
 ///
 void Function::viewCFGOnly() const {
-  ViewGraph(this, "cfg" + getName(), true);
+
+  CFGDOTInfo CFGInfo(this,nullptr,0);
+  CFGInfo.setHeuristic(false);
+  CFGInfo.setHeatColors(false);
+  CFGInfo.setEdgeWeights(true);
+  CFGInfo.setRawEdgeWeights(true);
+
+  ViewGraph(&CFGInfo, "cfg" + getName(), true);
 }
 
 FunctionPass *llvm::createCFGPrinterLegacyPassPass () {

lib/Analysis/CMakeLists.txt

   BlockFrequencyInfoImpl.cpp
   BranchProbabilityInfo.cpp
   CFG.cpp
+  HeatUtils.cpp
   CFGPrinter.cpp
   CFLAndersAliasAnalysis.cpp
   CFLSteensAliasAnalysis.cpp

lib/Analysis/CallPrinter.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/CallPrinter.h"
+
 #include "llvm/Analysis/CallGraph.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/DOTGraphTraitsPass.h"
+#include "llvm/Analysis/HeatUtils.h"
 
+#include <map>
+#include <set>
+
 using namespace llvm;
 
+
+static cl::opt<bool>
+ShowHeatColors("callgraph-heat-colors", cl::init(false), cl::Hidden,
+                   cl::desc("Show heat colors in call-graph"));
+
+static cl::opt<bool>
+EstimateEdgeWeight("callgraph-estimate-weight", cl::init(false),
+                   cl::Hidden, cl::desc("Estimate edge weights"));
+
+static cl::opt<bool>
+FullCallGraph("callgraph-full", cl::init(false), cl::Hidden,
+                   cl::desc("Print full call-graph (using external nodes)"));
+
+static cl::opt<bool>
+UseCallCounter("callgraph-call-count", cl::init(false), cl::Hidden,
+                   cl::desc("Use function's call counter as a heat metric"));
+
 namespace llvm {
 
-template <> struct DOTGraphTraits<CallGraph *> : public DefaultDOTGraphTraits {
-  DOTGraphTraits(bool isSimple = false) : DefaultDOTGraphTraits(isSimple) {}
 
-  static std::string getGraphName(CallGraph *Graph) { return "Call graph"; }
+class CallGraphDOTInfo {
+private:
+   CallGraph *CG;
+   Module *M;
+   std::map<const Function *, uint64_t> Freq;
+   uint64_t MaxFreq;
+public:
+   std::function<BlockFrequencyInfo *(Function &)> LookupBFI;
 
-  std::string getNodeLabel(CallGraphNode *Node, CallGraph *Graph) {
+   CallGraphDOTInfo(Module *M, CallGraph *CG,
+                     function_ref<BlockFrequencyInfo *(Function &)> LookupBFI){
+     this->M = M;
+     this->CG = CG;
+     MaxFreq = 0;
+
+     bool useHeuristic = !hasProfiling(*M);
+
+     for(Function &F : *M){
+       Freq[&F] = 0;
+       if(F.isDeclaration())
+         continue;
+       uint64_t localMaxFreq = 0;
+       if (UseCallCounter) {
+         Optional< uint64_t > EntryCount = F.getEntryCount();
+         if (EntryCount.hasValue())
+           localMaxFreq = EntryCount.getValue();       
+       } else {
+          localMaxFreq = llvm::getMaxFreq(F,LookupBFI(F),useHeuristic);
+       }
+       if(localMaxFreq>=MaxFreq) MaxFreq = localMaxFreq;
+       Freq[&F] = localMaxFreq;
+     }
+     this->LookupBFI = LookupBFI;
+     removeParallelEdges();
+   }
+
+   Module *getModule() const { return M; }
+   CallGraph *getCallGraph() const { return CG; }
+
+   uint64_t getFreq(const Function *F) { return Freq[F]; }
+
+   uint64_t getMaxFreq() { return MaxFreq; }
+
+private:
+   void removeParallelEdges(){
+      for (auto &I : (*CG)) {
+         CallGraphNode *Node = I.second.get();
+         
+         bool foundParallelEdge = true;
+         while (foundParallelEdge) {
+            std::set<Function *> visited;
+            foundParallelEdge = false;
+            for (std::vector<CallGraphNode::CallRecord>::iterator CI =
+                 Node->begin(); CI!=Node->end(); CI++) {
+               if (visited.find(CI->second->getFunction())==visited.end())
+                  visited.insert(CI->second->getFunction());
+               else {
+                  foundParallelEdge = true;
+                  Node->removeCallEdge(CI);
+                  break;
+               }
+            }
+         }
+      }
+   }
+};
+
+
+template <>
+struct GraphTraits<CallGraphDOTInfo *> : public GraphTraits<
+                                            const CallGraphNode *> {
+  static NodeRef getEntryNode(CallGraphDOTInfo *CGInfo) {
+    // Start at the external node!
+    return CGInfo->getCallGraph()->getExternalCallingNode();
+  }
+
+  typedef std::pair<const Function *const, std::unique_ptr<CallGraphNode>>
+      PairTy;
+  static const CallGraphNode *CGGetValuePtr(const PairTy &P) {
+    return P.second.get();
+  }
+
+  // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
+  typedef mapped_iterator<CallGraph::const_iterator, decltype(&CGGetValuePtr)>
+      nodes_iterator;
+
+  static nodes_iterator nodes_begin(CallGraphDOTInfo *CGInfo) {
+    return nodes_iterator(CGInfo->getCallGraph()->begin(), &CGGetValuePtr);
+  }
+  static nodes_iterator nodes_end(CallGraphDOTInfo *CGInfo) {
+    return nodes_iterator(CGInfo->getCallGraph()->end(), &CGGetValuePtr);
+  }
+};
+
+
+template<>
+struct DOTGraphTraits<CallGraphDOTInfo *> : public DefaultDOTGraphTraits {
+
+  DOTGraphTraits (bool isSimple=false) : DefaultDOTGraphTraits(isSimple) {}
+
+  static std::string getGraphName(CallGraphDOTInfo *CGInfo) {
+    return "Call graph: "+
+           std::string(CGInfo->getModule()->getModuleIdentifier());
+  }
+
+  static bool isNodeHidden(const CallGraphNode *Node) {
+    if (FullCallGraph)
+       return false;
+
+    if (Node->getFunction())
+       return false;
+
+    return true;
+  }
+
+  std::string getNodeLabel(const CallGraphNode *Node, CallGraphDOTInfo *CGInfo){
+
+    if (Node==CGInfo->getCallGraph()->getExternalCallingNode())
+       return "external caller";
+
+    if (Node==CGInfo->getCallGraph()->getCallsExternalNode())
+       return "external callee";
+
     if (Function *Func = Node->getFunction())
       return Func->getName();
 
     return "external node";
   }
-};
 
-struct AnalysisCallGraphWrapperPassTraits {
-  static CallGraph *getGraph(CallGraphWrapperPass *P) {
-    return &P->getCallGraph();
+  static const CallGraphNode *CGGetValuePtr(CallGraphNode::CallRecord P) {
+    return P.second;
   }
+
+  // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
+  typedef mapped_iterator<CallGraphNode::const_iterator,
+                          decltype(&CGGetValuePtr)>
+      nodes_iterator;
+
+
+  std::string getEdgeAttributes(const CallGraphNode *Node, nodes_iterator I,
+                                CallGraphDOTInfo *CGInfo) {
+    if (!EstimateEdgeWeight)
+       return "";
+
+    Function *F = Node->getFunction();
+    if (F==nullptr || F->isDeclaration())
+       return "";
+
+    Function *SuccFunction = (*I)->getFunction();
+    if (SuccFunction==nullptr)
+       return "";
+
+    uint64_t counter = getNumOfCalls(*F, *SuccFunction, CGInfo->LookupBFI);
+    std::string Attrs = "label=\"" + std::to_string(counter) + "\"";
+    return Attrs;
+  }
+
+  std::string getNodeAttributes(const CallGraphNode *Node,
+                                CallGraphDOTInfo *CGInfo) {
+    Function *F = Node->getFunction();
+    if (F==nullptr || F->isDeclaration())
+       return "";
+
+    std::string attrs = "";
+    if(ShowHeatColors){
+       uint64_t freq = CGInfo->getFreq(F);
+       std::string color = getHeatColor(freq, CGInfo->getMaxFreq());
+       std::string edgeColor = (freq<=(CGInfo->getMaxFreq()/2))?
+                            getHeatColor(0):getHeatColor(1);
+
+       attrs = "color=\"" + edgeColor +
+                        "ff\", style=filled, fillcolor=\"" + color + "80\"";
+    }
+    return attrs;
+  }
+
 };
 
 } // end llvm namespace

 
 namespace {
 
-struct CallGraphViewer
-    : public DOTGraphTraitsModuleViewer<CallGraphWrapperPass, true, CallGraph *,
-                                        AnalysisCallGraphWrapperPassTraits> {
+// Viewer
+
+class CallGraphViewer : public ModulePass {
+public:
   static char ID;
+  CallGraphViewer() : ModulePass(ID) {}
 
-  CallGraphViewer()
-      : DOTGraphTraitsModuleViewer<CallGraphWrapperPass, true, CallGraph *,
-                                   AnalysisCallGraphWrapperPassTraits>(
-            "callgraph", ID) {
-    initializeCallGraphViewerPass(*PassRegistry::getPassRegistry());
-  }
+  void getAnalysisUsage(AnalysisUsage &AU) const;
+  bool runOnModule(Module &M) override;
 };
 
-struct CallGraphDOTPrinter : public DOTGraphTraitsModulePrinter<
-                              CallGraphWrapperPass, true, CallGraph *,
-                              AnalysisCallGraphWrapperPassTraits> {
+
+void CallGraphViewer::getAnalysisUsage(AnalysisUsage &AU) const {
+  ModulePass::getAnalysisUsage(AU);
+  AU.addRequired<BlockFrequencyInfoWrapperPass>();
+  AU.setPreservesAll();
+}
+
+bool CallGraphViewer::runOnModule(Module &M) {
+  auto LookupBFI = [this](Function &F) {
+    return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
+  };
+
+  CallGraph CG(M);
+  CallGraphDOTInfo CFGInfo(&M,&CG,LookupBFI);
+
+  std::string Title = DOTGraphTraits<CallGraphDOTInfo*>::getGraphName(&CFGInfo);
+  ViewGraph(&CFGInfo, "callgraph", true, Title);
+
+  return false;
+}
+
+
+// DOT Printer
+
+class CallGraphDOTPrinter : public ModulePass {
+public:
   static char ID;
+  CallGraphDOTPrinter() : ModulePass(ID) {}
 
-  CallGraphDOTPrinter()
-      : DOTGraphTraitsModulePrinter<CallGraphWrapperPass, true, CallGraph *,
-                                    AnalysisCallGraphWrapperPassTraits>(
-            "callgraph", ID) {
-    initializeCallGraphDOTPrinterPass(*PassRegistry::getPassRegistry());
-  }
+  void getAnalysisUsage(AnalysisUsage &AU) const;
+  bool runOnModule(Module &M) override;
 };
 
+
+void CallGraphDOTPrinter::getAnalysisUsage(AnalysisUsage &AU) const {
+  ModulePass::getAnalysisUsage(AU);
+  AU.addRequired<BlockFrequencyInfoWrapperPass>();
+  AU.setPreservesAll();
+}
+
+bool CallGraphDOTPrinter::runOnModule(Module &M) {
+  auto LookupBFI = [this](Function &F) {
+    return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
+  };
+
+  std::string Filename = (std::string(M.getModuleIdentifier())+
+                         ".callgraph.dot");
+  errs() << "Writing '" << Filename << "'...";
+
+  std::error_code EC;
+  raw_fd_ostream File(Filename, EC, sys::fs::F_Text);
+
+  CallGraph CG(M);
+  CallGraphDOTInfo CFGInfo(&M,&CG,LookupBFI);
+
+  if(!EC)
+     WriteGraph(File, &CFGInfo);
+  else
+     errs() << "  error opening file for writing!";
+  errs() << "\n";
+
+  return false;
+}
+
 } // end anonymous namespace
 
 char CallGraphViewer::ID = 0;

lib/Analysis/DomPrinter.cpp

     if (!BB)
       return "Post dominance root node";
 
-
     if (isSimple())
-      return DOTGraphTraits<const Function*>
-        ::getSimpleNodeLabel(BB, BB->getParent());
+      return DOTGraphTraits<CFGDOTInfo*>
+        ::getSimpleNodeLabel(BB, nullptr);
     else
-      return DOTGraphTraits<const Function*>
-        ::getCompleteNodeLabel(BB, BB->getParent());
+      return DOTGraphTraits<CFGDOTInfo*>
+        ::getCompleteNodeLabel(BB, nullptr);
   }
 };
 

lib/Analysis/HeatUtils.cpp

+//===-- HeatUtils.cpp - Utility for printing heat colors --------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Utility for printing heat colors based on heuristics or profiling
+// information.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/HeatUtils.h"
+#include "llvm/IR/Instructions.h"
+
+namespace llvm {
+
+static const std::string heatPalette[100] = {"#3d50c3", "#4055c8", "#4358cb", "#465ecf", "#4961d2", "#4c66d6", "#4f69d9", "#536edd", "#5572df", "#5977e3", "#5b7ae5", "#5f7fe8", "#6282ea", "#6687ed", "#6a8bef", "#6c8ff1", "#7093f3", "#7396f5", "#779af7", "#7a9df8", "#7ea1fa", "#81a4fb", "#85a8fc", "#88abfd", "#8caffe", "#8fb1fe", "#93b5fe", "#96b7ff", "#9abbff", "#9ebeff", "#a1c0ff", "#a5c3fe", "#a7c5fe", "#abc8fd", "#aec9fc", "#b2ccfb", "#b5cdfa", "#b9d0f9", "#bbd1f8", "#bfd3f6", "#c1d4f4", "#c5d6f2", "#c7d7f0", "#cbd8ee", "#cedaeb", "#d1dae9", "#d4dbe6", "#d6dce4", "#d9dce1", "#dbdcde", "#dedcdb", "#e0dbd8", "#e3d9d3", "#e5d8d1", "#e8d6cc", "#ead5c9", "#ecd3c5", "#eed0c0", "#efcebd", "#f1ccb8", "#f2cab5", "#f3c7b1", "#f4c5ad", "#f5c1a9", "#f6bfa6", "#f7bca1", "#f7b99e", "#f7b599", "#f7b396", "#f7af91", "#f7ac8e", "#f7a889", "#f6a385", "#f5a081", "#f59c7d", "#f4987a", "#f39475", "#f29072", "#f08b6e", "#ef886b", "#ed8366", "#ec7f63", "#e97a5f", "#e8765c", "#e57058", "#e36c55", "#e16751", "#de614d", "#dc5d4a", "#d85646", "#d65244", "#d24b40", "#d0473d", "#cc403a", "#ca3b37", "#c53334", "#c32e31", "#be242e", "#bb1b2c", "#b70d28"};
+
+static const unsigned heatSize = 100;
+
+
+bool hasProfiling(Module &M){
+  for (Function &F : M) {
+    for (BasicBlock &BB : F) {
+      Instruction *TI = BB.getTerminator();
+      if (TI==nullptr)
+        continue;
+      if (TI->getMetadata(llvm::LLVMContext::MD_prof)!=nullptr)
+        return true;
+    }
+  }
+  return false;
+}
+
+uint64_t getBlockFreq(const BasicBlock *BB, const BlockFrequencyInfo *BFI,
+                      bool useHeuristic){
+  uint64_t freqVal = 0;
+  if (!useHeuristic) {
+    Optional< uint64_t > freq = BFI->getBlockProfileCount(BB);
+    if (freq.hasValue())
+      freqVal = freq.getValue();
+  } else {
+    freqVal = BFI->getBlockFreq(BB).getFrequency();
+  }
+  return freqVal;
+}
+
+uint64_t getNumOfCalls(Function &callerFunction, Function &calledFunction,
+                      function_ref<BlockFrequencyInfo *(Function &)> LookupBFI,
+                      bool useHeuristic){
+  auto *BFI = LookupBFI(callerFunction);
+  uint64_t counter = 0;
+  for (BasicBlock &BB : callerFunction) {
+     uint64_t freq = getBlockFreq(&BB,BFI,useHeuristic);
+     for (Instruction &I : BB) {
+        if (CallInst *Call = dyn_cast<CallInst>(&I)) {
+           if (Call->getCalledFunction()==(&calledFunction))
+              counter += freq;
+        }
+     }
+  }
+  return counter;
+}
+
+uint64_t getMaxFreq(const Function &F, const BlockFrequencyInfo *BFI, bool useHeuristic){
+  uint64_t maxFreq = 0;
+  for (const BasicBlock &BB : F) {
+     uint64_t freqVal = getBlockFreq(&BB,BFI,useHeuristic);
+     if (freqVal>=maxFreq)
+        maxFreq = freqVal;
+  }
+  return maxFreq;
+}
+
+
+uint64_t getMaxFreq(Module &M,
+                    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI,
+                    bool useHeuristic){
+  uint64_t maxFreq = 0;
+  for (Function &F : M) {
+    if (F.isDeclaration())
+      continue;
+    uint64_t localMaxFreq = getMaxFreq(F,LookupBFI(F),useHeuristic);
+    if (localMaxFreq>=maxFreq)
+       maxFreq = localMaxFreq;
+  }
+  return maxFreq;
+}
+
+std::string getHeatColor(uint64_t freq, uint64_t maxFreq){
+  if (freq>maxFreq) freq = maxFreq;
+  unsigned colorId = unsigned( round((double(freq)/maxFreq)*(heatSize-1.0)) );
+  return heatPalette[colorId];
+}
+
+std::string getHeatColor(double percent){
+  if (percent>1.0)
+     percent = 1.0;
+  if (percent<0.0)
+     percent = 0.0;
+  unsigned colorId = unsigned( round(percent*(heatSize-1.0)) );
+  return heatPalette[colorId];
+}
+
+}

lib/Analysis/RegionPrinter.cpp

       BasicBlock *BB = Node->getNodeAs<BasicBlock>();
 
       if (isSimple())
-        return DOTGraphTraits<const Function*>
-          ::getSimpleNodeLabel(BB, BB->getParent());
+        return DOTGraphTraits<CFGDOTInfo*>
+          ::getSimpleNodeLabel(BB, nullptr);
       else
-        return DOTGraphTraits<const Function*>
-          ::getCompleteNodeLabel(BB, BB->getParent());
+        return DOTGraphTraits<CFGDOTInfo*>
+          ::getCompleteNodeLabel(BB, nullptr);
     }
 
     return "Not implemented";

lib/Transforms/Scalar/NewGVN.cpp

 
 #ifndef NDEBUG
 static std::string getBlockName(const BasicBlock *B) {
-  return DOTGraphTraits<const Function *>::getSimpleNodeLabel(B, nullptr);
+  return DOTGraphTraits<CFGDOTInfo *>::getSimpleNodeLabel(B, nullptr);
 }
 #endif