[llvm] [Analysis] Fix null ptr dereference when using WriteGraph without branch probability info (PR #104102)

[Tyler Nowicki via llvm-commits]

https://github.com/TylerNowicki updated https://github.com/llvm/llvm-project/pull/104102

>From 8fd9fe5e2eb68f3b5d69e3afe9350064e54ab792 Mon Sep 17 00:00:00 2001
From: tnowicki <tnowicki.nowicki at amd.com>
Date: Wed, 14 Aug 2024 13:49:33 -0400
Subject: [PATCH 1/2] [Analysis] Fix null ptr dereference when using WriteGraph
 without branch probability info

---
 llvm/include/llvm/Analysis/CFGPrinter.h | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Analysis/CFGPrinter.h b/llvm/include/llvm/Analysis/CFGPrinter.h
index 7d4d3a574b9e81..cd785331d1f146 100644
--- a/llvm/include/llvm/Analysis/CFGPrinter.h
+++ b/llvm/include/llvm/Analysis/CFGPrinter.h
@@ -265,18 +265,19 @@ struct DOTGraphTraits<DOTFuncInfo *> : public DefaultDOTGraphTraits {
   /// Display the raw branch weights from PGO.
   std::string getEdgeAttributes(const BasicBlock *Node, const_succ_iterator I,
                                 DOTFuncInfo *CFGInfo) {
+    // If BPI is not provided do not display any edge attributes
+    if (!CFGInfo->showEdgeWeights())
+      return "";
+
     unsigned OpNo = I.getSuccessorIndex();
     const Instruction *TI = Node->getTerminator();
     BasicBlock *SuccBB = TI->getSuccessor(OpNo);
     auto BranchProb = CFGInfo->getBPI()->getEdgeProbability(Node, SuccBB);
     double WeightPercent = ((double)BranchProb.getNumerator()) /
                            ((double)BranchProb.getDenominator());
-
     std::string TTAttr =
         formatv("tooltip=\"{0} -> {1}\\nProbability {2:P}\" ", getBBName(Node),
                 getBBName(SuccBB), WeightPercent);
-    if (!CFGInfo->showEdgeWeights())
-      return TTAttr;
 
     if (TI->getNumSuccessors() == 1)
       return TTAttr + "penwidth=2";

>From f7f06a9e9d10934dd61c80db6d87c1c7d2529ff6 Mon Sep 17 00:00:00 2001
From: tnowicki <tnowicki.nowicki at amd.com>
Date: Fri, 16 Aug 2024 13:12:10 -0400
Subject: [PATCH 2/2] [Analysis] unittest for WriteGraph without BPI

---
 llvm/unittests/Analysis/CMakeLists.txt      |  1 +
 llvm/unittests/Analysis/GraphWriterTest.cpp | 77 +++++++++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 llvm/unittests/Analysis/GraphWriterTest.cpp

diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt
index 3cba630867a83b..d9eb81faac42ad 100644
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -25,6 +25,7 @@ set(ANALYSIS_TEST_SOURCES
   DDGTest.cpp
   DomTreeUpdaterTest.cpp
   DXILResourceTest.cpp
+  GraphWriterTest.cpp
   GlobalsModRefTest.cpp
   FunctionPropertiesAnalysisTest.cpp
   InlineCostTest.cpp
diff --git a/llvm/unittests/Analysis/GraphWriterTest.cpp b/llvm/unittests/Analysis/GraphWriterTest.cpp
new file mode 100644
index 00000000000000..a723c92d157618
--- /dev/null
+++ b/llvm/unittests/Analysis/GraphWriterTest.cpp
@@ -0,0 +1,77 @@
+//===- GraphWriterTest.cpp - GraphWriter unit tests -----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/GraphWriter.h"
+#include "llvm/Analysis/CFGPrinter.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+#include "gtest/gtest.h"
+#include <string>
+
+#define ASSERT_NO_ERROR(x)                                                     \
+  if (std::error_code ASSERT_NO_ERROR_ec = x) {                                \
+    SmallString<128> MessageStorage;                                           \
+    raw_svector_ostream Message(MessageStorage);                               \
+    Message << #x ": did not return errc::success.\n"                          \
+            << "error number: " << ASSERT_NO_ERROR_ec.value() << "\n"          \
+            << "error message: " << ASSERT_NO_ERROR_ec.message() << "\n";      \
+    GTEST_FATAL_FAILURE_(MessageStorage.c_str());                              \
+  } else {                                                                     \
+  }
+
+namespace llvm {
+namespace {
+
+class GraphWriterTest : public testing::Test {
+protected:
+  LLVMContext C;
+
+  std::unique_ptr<Module> makeLLVMModule() {
+    const char *ModuleStrig = "define i32 @f(i32 %x) {\n"
+                              "bb0:\n"
+                              "  %y1 = icmp eq i32 %x, 0 \n"
+                              "  br i1 %y1, label %bb1, label %bb2 \n"
+                              "bb1:\n"
+                              "  br label %bb3\n"
+                              "bb2:\n"
+                              "  br label %bb3\n"
+                              "bb3:\n"
+                              "  %y2 = phi i32 [0, %bb1], [1, %bb2] \n"
+                              "  ret i32 %y2\n"
+                              "}\n";
+    SMDiagnostic Err;
+    return parseAssemblyString(ModuleStrig, Err, C);
+  }
+};
+
+static void writeCFGToDotFile(Function &F, std::string Name,
+                              bool CFGOnly = false) {
+  std::error_code EC;
+  raw_fd_ostream File(Name + ".dot", EC, sys::fs::OpenFlags::OF_Text);
+
+  DOTFuncInfo CFGInfo(&F);
+
+  ASSERT_NO_ERROR(EC);
+  // Test intentionally does not pass BPI, WriteGraph should work without it.
+  WriteGraph(File, &CFGInfo, CFGOnly);
+}
+
+TEST_F(GraphWriterTest, WriteCFGDotFileTest) {
+  auto M = makeLLVMModule();
+  Function *F = M->getFunction("f");
+
+  writeCFGToDotFile(*F, "test-full");
+  writeCFGToDotFile(*F, "test-cfg-only", true);
+}
+
+} // end anonymous namespace
+} // end namespace llvm