[llvm] [MemProf] Add dot graph dumping immediately after stack node update (PR #143025)

[Teresa Johnson via llvm-commits]

https://github.com/teresajohnson created https://github.com/llvm/llvm-project/pull/143025

To aid in debugging, (optionally) dump the dot graph immediately after
the stack update phase (which matches nodes to interior callsites) and
before we cleanup mismatched callee edges (either via tail call fixup,
indirect call fixup, or nulling otherwise).


>From 5b4c20365ae7297625ef2dca5b5b901fcf0bfa1d Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Thu, 5 Jun 2025 12:52:08 -0700
Subject: [PATCH] [MemProf] Add dot graph dumping immediately after stack node
 update

To aid in debugging, (optionally) dump the dot graph immediately after
the stack update phase (which matches nodes to interior callsites) and
before we cleanup mismatched callee edges (either via tail call fixup,
indirect call fixup, or nulling otherwise).
---
 llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp | 6 ++++++
 .../duplicate-context-ids.ll                             | 8 ++++++--
 .../MemProfContextDisambiguation/tailcall-nonunique.ll   | 9 +++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index cff38a8e68c6a..f28fe51fb6a5a 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -2183,6 +2183,9 @@ ModuleCallsiteContextGraph::ModuleCallsiteContextGraph(
 
   updateStackNodes();
 
+  if (ExportToDot)
+    exportToDot("poststackupdate");
+
   handleCallsitesWithMultipleTargets();
 
   markBackedges();
@@ -2285,6 +2288,9 @@ IndexCallsiteContextGraph::IndexCallsiteContextGraph(
 
   updateStackNodes();
 
+  if (ExportToDot)
+    exportToDot("poststackupdate");
+
   handleCallsitesWithMultipleTargets();
 
   markBackedges();
diff --git a/llvm/test/Transforms/MemProfContextDisambiguation/duplicate-context-ids.ll b/llvm/test/Transforms/MemProfContextDisambiguation/duplicate-context-ids.ll
index 93557e963752e..18521c013d555 100644
--- a/llvm/test/Transforms/MemProfContextDisambiguation/duplicate-context-ids.ll
+++ b/llvm/test/Transforms/MemProfContextDisambiguation/duplicate-context-ids.ll
@@ -64,6 +64,10 @@
 ; RUN:  --check-prefix=STATS --check-prefix=REMARKS
 
 ; RUN:  cat %t.ccg.prestackupdate.dot | FileCheck %s --check-prefix=DOTPRE
+;; There are no changes for this test case in between from immediately after
+;; the stack update through the end of graph finalization, so the
+;; corresponding dot graphs should be identical.
+; RUN:  cat %t.ccg.poststackupdate.dot | FileCheck %s --check-prefix=DOTPOST
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOTPOST
 ;; We should clone D once for the cold allocations via C.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
@@ -313,8 +317,8 @@ attributes #6 = { builtin }
 ; DOTPRE: }
 
 
-; DOTPOST:digraph "postbuild" {
-; DOTPOST:	label="postbuild";
+; DOTPOST:digraph "post
+; DOTPOST:	label="post
 ; DOTPOST:	Node[[D:0x[a-z0-9]+]] [shape=record,tooltip="N[[D]] ContextIds: 1 2 3 4",fillcolor="mediumorchid1",style="filled",label="{OrigId: Alloc0\n_Z1Dv -\> _Znam}"];
 ; DOTPOST:	Node[[F:0x[a-z0-9]+]] [shape=record,tooltip="N[[F]] ContextIds: 2",fillcolor="brown1",style="filled",label="{OrigId: 13543580133643026784\n_Z1Fv -\> _Z1Dv}"];
 ; DOTPOST:	Node[[F]] -> Node[[D]][tooltip="ContextIds: 2",fillcolor="brown1"
diff --git a/llvm/test/Transforms/MemProfContextDisambiguation/tailcall-nonunique.ll b/llvm/test/Transforms/MemProfContextDisambiguation/tailcall-nonunique.ll
index 985c381ad42ff..97e433ccefa99 100644
--- a/llvm/test/Transforms/MemProfContextDisambiguation/tailcall-nonunique.ll
+++ b/llvm/test/Transforms/MemProfContextDisambiguation/tailcall-nonunique.ll
@@ -6,11 +6,20 @@
 ; REQUIRES: asserts
 
 ; RUN: opt -passes=memprof-context-disambiguation -supports-hot-cold-new \
+; RUN:	-memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
 ; RUN:  -stats -debug %s -S 2>&1 | FileCheck %s --check-prefix=STATS \
 ; RUN:  --check-prefix=IR --check-prefix=DEBUG
 
 ; DEBUG: Not found through unique tail call chain: _Z3barv from main that actually called xyz (found multiple possible chains)
 
+;; Graph fixup should have nulled the call in the node that could not be
+;; resolved via tail call fixup. That happens in between the initial
+;; stack update and the end of graph building.
+; RUN:  cat %t.ccg.poststackupdate.dot | FileCheck %s --check-prefix=DOTPOSTSTACKUPDATE
+; DOTPOSTSTACKUPDATE: {OrigId: 15025054523792398438\nmain -\> xyz}
+; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOTPOSTBUILD
+; DOTPOSTBUILD: {OrigId: 15025054523792398438\nnull call (external)}
+
 ;; Check that all calls in the IR are to the original functions, leading to a
 ;; non-cold operator new call.