Introduce -dot-cfg-mssa option which creates dot-cfg style file with mssa comments included in source

Summary:
Expand the print-memoryssa and print<memoryssa> passes with a new hidden
option -cfg-dot-mssa that names a file. When set, a dot-cfg style file
will be generated into the named file with the memoryssa comments retained
and those blocks containing them shown in light pink. The option does
nothing in isolation.

Author: Jamie Schmeiser <schmeise@ca.ibm.com>

Reviewed By: asbirlea (Alina Sbirlea), dblaikie (David Blaikie)

Differential Revision: https://reviews.llvm.org/D90638

llvm/include/llvm/Analysis/CFGPrinter.h

@@ -18,6 +18,7 @@
 #ifndef LLVM_ANALYSIS_CFGPRINTER_H
 #define LLVM_ANALYSIS_CFGPRINTER_H
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/HeatUtils.h"
@@ -141,8 +142,18 @@ struct DOTGraphTraits<DOTFuncInfo *> : public DefaultDOTGraphTraits {
     return OS.str();
   }
 
-  static std::string getCompleteNodeLabel(const BasicBlock *Node,
-                                          DOTFuncInfo *) {
+  static void eraseComment(std::string &OutStr, unsigned &I, unsigned Idx) {
+    OutStr.erase(OutStr.begin() + I, OutStr.begin() + Idx);
+    --I;
+  }
+
+  static std::string getCompleteNodeLabel(
+      const BasicBlock *Node, DOTFuncInfo *,
+      llvm::function_ref<void(raw_string_ostream &, const BasicBlock &)>
+          HandleBasicBlock = [](raw_string_ostream &OS,
+                                const BasicBlock &Node) -> void { OS << Node; },
+      llvm::function_ref<void(std::string &, unsigned &, unsigned)>
+          HandleComment = eraseComment) {
     enum { MaxColumns = 80 };
     std::string Str;
     raw_string_ostream OS(Str);
@@ -152,7 +163,7 @@ struct DOTGraphTraits<DOTFuncInfo *> : public DefaultDOTGraphTraits {
       OS << ":";
     }
 
-    OS << *Node;
+    HandleBasicBlock(OS, *Node);
     std::string OutStr = OS.str();
     if (OutStr[0] == '\n')
       OutStr.erase(OutStr.begin());
@@ -168,8 +179,7 @@ struct DOTGraphTraits<DOTFuncInfo *> : public DefaultDOTGraphTraits {
         LastSpace = 0;
       } else if (OutStr[i] == ';') {             // Delete comments!
         unsigned Idx = OutStr.find('\n', i + 1); // Find end of line
-        OutStr.erase(OutStr.begin() + i, OutStr.begin() + Idx);
-        --i;
+        HandleComment(OutStr, i, Idx);
       } else if (ColNum == MaxColumns) { // Wrap lines.
         // Wrap very long names even though we can't find a space.
         if (!LastSpace)

llvm/lib/Analysis/MemorySSA.cpp

@@ -24,6 +24,7 @@
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/CFGPrinter.h"
 #include "llvm/Analysis/IteratedDominanceFrontier.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Config/llvm-config.h"
@@ -59,6 +60,11 @@ using namespace llvm;
 
 #define DEBUG_TYPE "memoryssa"
 
+static cl::opt<std::string>
+    DotCFGMSSA("dot-cfg-mssa",
+               cl::value_desc("file name for generated dot file"),
+               cl::desc("file name for generated dot file"), cl::init(""));
+
 INITIALIZE_PASS_BEGIN(MemorySSAWrapperPass, "memoryssa", "Memory SSA", false,
                       true)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
@@ -2256,9 +2262,94 @@ void MemorySSAPrinterLegacyPass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<MemorySSAWrapperPass>();
 }
 
+class DOTFuncMSSAInfo {
+private:
+  const Function &F;
+  MemorySSAAnnotatedWriter MSSAWriter;
+
+public:
+  DOTFuncMSSAInfo(const Function &F, MemorySSA &MSSA)
+      : F(F), MSSAWriter(&MSSA) {}
+
+  const Function *getFunction() { return &F; }
+  MemorySSAAnnotatedWriter &getWriter() { return MSSAWriter; }
+};
+
+template <>
+struct GraphTraits<DOTFuncMSSAInfo *> : public GraphTraits<const BasicBlock *> {
+  static NodeRef getEntryNode(DOTFuncMSSAInfo *CFGInfo) {
+    return &(CFGInfo->getFunction()->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(DOTFuncMSSAInfo *CFGInfo) {
+    return nodes_iterator(CFGInfo->getFunction()->begin());
+  }
+
+  static nodes_iterator nodes_end(DOTFuncMSSAInfo *CFGInfo) {
+    return nodes_iterator(CFGInfo->getFunction()->end());
+  }
+
+  static size_t size(DOTFuncMSSAInfo *CFGInfo) {
+    return CFGInfo->getFunction()->size();
+  }
+};
+
+template <>
+struct DOTGraphTraits<DOTFuncMSSAInfo *> : public DefaultDOTGraphTraits {
+
+  DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {}
+
+  static std::string getGraphName(DOTFuncMSSAInfo *CFGInfo) {
+    return "MSSA CFG for '" + CFGInfo->getFunction()->getName().str() +
+           "' function";
+  }
+
+  std::string getNodeLabel(const BasicBlock *Node, DOTFuncMSSAInfo *CFGInfo) {
+    return DOTGraphTraits<DOTFuncInfo *>::getCompleteNodeLabel(
+        Node, nullptr,
+        [CFGInfo](raw_string_ostream &OS, const BasicBlock &BB) -> void {
+          BB.print(OS, &CFGInfo->getWriter(), true, true);
+        },
+        [](std::string &S, unsigned &I, unsigned Idx) -> void {
+          std::string Str = S.substr(I, Idx - I);
+          StringRef SR = Str;
+          if (SR.count(" = MemoryDef(") || SR.count(" = MemoryPhi(") ||
+              SR.count("MemoryUse("))
+            return;
+          DOTGraphTraits<DOTFuncInfo *>::eraseComment(S, I, Idx);
+        });
+  }
+
+  static std::string getEdgeSourceLabel(const BasicBlock *Node,
+                                        const_succ_iterator I) {
+    return DOTGraphTraits<DOTFuncInfo *>::getEdgeSourceLabel(Node, I);
+  }
+
+  /// Display the raw branch weights from PGO.
+  std::string getEdgeAttributes(const BasicBlock *Node, const_succ_iterator I,
+                                DOTFuncMSSAInfo *CFGInfo) {
+    return "";
+  }
+
+  std::string getNodeAttributes(const BasicBlock *Node,
+                                DOTFuncMSSAInfo *CFGInfo) {
+    return getNodeLabel(Node, CFGInfo).find(';') != std::string::npos
+               ? "style=filled, fillcolor=lightpink"
+               : "";
+  }
+};
+
 bool MemorySSAPrinterLegacyPass::runOnFunction(Function &F) {
   auto &MSSA = getAnalysis<MemorySSAWrapperPass>().getMSSA();
-  MSSA.print(dbgs());
+  if (DotCFGMSSA != "") {
+    DOTFuncMSSAInfo CFGInfo(F, MSSA);
+    WriteGraph(&CFGInfo, "", false, "MSSA", DotCFGMSSA);
+  } else
+    MSSA.print(dbgs());
+
   if (VerifyMemorySSA)
     MSSA.verifyMemorySSA();
   return false;
@@ -2284,8 +2375,14 @@ bool MemorySSAAnalysis::Result::invalidate(
 
 PreservedAnalyses MemorySSAPrinterPass::run(Function &F,
                                             FunctionAnalysisManager &AM) {
-  OS << "MemorySSA for function: " << F.getName() << "\n";
-  AM.getResult<MemorySSAAnalysis>(F).getMSSA().print(OS);
+  auto &MSSA = AM.getResult<MemorySSAAnalysis>(F).getMSSA();
+  if (DotCFGMSSA != "") {
+    DOTFuncMSSAInfo CFGInfo(F, MSSA);
+    WriteGraph(&CFGInfo, "", false, "MSSA", DotCFGMSSA);
+  } else {
+    OS << "MemorySSA for function: " << F.getName() << "\n";
+    MSSA.print(OS);
+  }
 
   return PreservedAnalyses::all();
 }

llvm/test/Analysis/MemorySSA/print-dot-cfg-mssa.ll

@@ -0,0 +1,97 @@
+; RUN: opt -basic-aa -print-memoryssa -dot-cfg-mssa=out.dot -enable-new-pm=0 -analyze < %s 2>&1 > /dev/null
+;RUN: FileCheck %s -input-file=out.dot
+; RUN: opt -aa-pipeline=basic-aa -passes='print<memoryssa>' -dot-cfg-mssa=out.dot < %s 2>&1 > /dev/null
+;RUN: FileCheck %s -input-file=out.dot
+
+; Test -dot-cfg-mssa option for -print-memoryssa.
+; Test is based on following C code with some forwarding basic blocks
+; added to show that only those blocks with memory ssa comments
+; are colourized.
+
+;void g();
+
+;int f(int *p, int *q, int *r) {
+;  int i = 0;
+;  if (*r)
+;    i = 1;
+;  else
+;    g();
+;  *p = *q + 1;
+;  if (i)
+;    ++i;
+;  return *q;
+;}
+
+define signext i32 @f(i32* %p, i32* %q, i32* %r) {
+entry:
+  br label %bb1
+
+bb1:
+  %p.addr = alloca i32*, align 8
+  %q.addr = alloca i32*, align 8
+  %r.addr = alloca i32*, align 8
+  %i = alloca i32, align 4
+  store i32* %p, i32** %p.addr, align 8
+  store i32* %q, i32** %q.addr, align 8
+  store i32* %r, i32** %r.addr, align 8
+  %0 = bitcast i32* %i to i8*
+  store i32 0, i32* %i, align 4
+  %1 = load i32*, i32** %r.addr, align 8
+  %2 = load i32, i32* %1, align 4
+  %tobool = icmp ne i32 %2, 0
+  br i1 %tobool, label %if.then, label %if.else
+
+if.then:
+  store i32 1, i32* %i, align 4
+  br label %bb2
+
+bb2:
+  br label %if.end
+
+if.else:
+  call void bitcast (void (...)* @g to void ()*)()
+  br label %if.end
+
+if.end:
+  %3 = load i32*, i32** %q.addr, align 8
+  %4 = load i32, i32* %3, align 4
+  %add = add nsw i32 %4, 1
+  %5 = load i32*, i32** %p.addr, align 8
+  store i32 %add, i32* %5, align 4
+  %6 = load i32, i32* %i, align 4
+  %tobool1 = icmp ne i32 %6, 0
+  br i1 %tobool1, label %if.then2, label %if.end3
+
+if.then2:
+  %7 = load i32, i32* %i, align 4
+  %inc = add nsw i32 %7, 1
+  br label %bb3
+
+bb3:
+  store i32 %inc, i32* %i, align 4
+  br label %if.end3
+
+if.end3:
+  br label %bb4
+
+bb4:
+  %8 = load i32*, i32** %q.addr, align 8
+  %9 = load i32, i32* %8, align 4
+  %10 = bitcast i32* %i to i8*
+  ret i32 %9
+}
+
+declare void @g(...)
+
+; CHECK: digraph "MSSA"
+; CHECK-NEXT: label="MSSA";
+; CHECK: {{Node0x.* [shape=record,label="{entry:.*}"]}}
+; CHECK: {{[shape=record,style=filled, fillcolor=lightpink,label="{bb1:.*1 = MemoryDef(liveOnEntry).*2 = MemoryDef(1).*3 = MemoryDef(2).*4 = MemoryDef(3).*MemoryUse(3).*MemoryUse(liveOnEntry).*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{if.then:.*5 = MemoryDef(4).*}"]}}
+; CHECK {{[shape=record,label="{bb2:.*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{if.else:.*6 = MemoryDef(4).*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{if.end:.*10 = MemoryPhi({bb2,5},{if.else,6})/*MemoryUse(2).*MemoryUse(10).*MemoryUse(1).*7 = MemoryDef(10).*MemoryUse(10).*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{if.then2:.*MemoryUse(10).*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{bb3:.*8 = MemoryDef(7).*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{if.end3:.*9 = MemoryPhi({if.end,7},{bb3,8}).*}"]}}
+; CHECK {{[shape=record,style=filled, fillcolor=lightpink,label="{bb4:.*MemoryUse(2).*MemoryUse(7).*}"]}}