[MemProf] Handle recursion during stack node update

[teresajohnson]

If we are replacing a sequence of stack nodes with a single node
representing inlined IR, and the stack id sequence contains recursion,
we may have already removed some edges. Handle this case correctly by
skipping the now removed edge.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Teresa Johnson (teresajohnson)

Changes

If we are replacing a sequence of stack nodes with a single node
representing inlined IR, and the stack id sequence contains recursion,
we may have already removed some edges. Handle this case correctly by
skipping the now removed edge.

---

Full diff: https://github.com/llvm/llvm-project/pull/135837.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp (+6-1)
• (modified) llvm/test/Transforms/MemProfContextDisambiguation/inlined2.ll (+7-4)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index a5e0251277d8f..561c01cb01f82 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -1711,7 +1711,12 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
       // edge from the prior node.
       if (PrevNode) {
         auto *PrevEdge = CurNode->findEdgeFromCallee(PrevNode);
-        assert(PrevEdge);
+        // If the sequence contained recursion, we might have already removed
+        // some edges during the connectNewNode calls above.
+        if (!PrevEdge) {
+          PrevNode = CurNode;
+          continue;
+        }
         set_subtract(PrevEdge->getContextIds(), SavedContextIds);
         if (PrevEdge->getContextIds().empty())
           removeEdgeFromGraph(PrevEdge);
diff --git a/llvm/test/Transforms/MemProfContextDisambiguation/inlined2.ll b/llvm/test/Transforms/MemProfContextDisambiguation/inlined2.ll
index cc97b5290e25a..2cc655e927d12 100644
--- a/llvm/test/Transforms/MemProfContextDisambiguation/inlined2.ll
+++ b/llvm/test/Transforms/MemProfContextDisambiguation/inlined2.ll
@@ -40,6 +40,9 @@
 ;; in the input IR to ensure that the MIB call chain is matched to the longer
 ;; inline sequences from main.
 ;;
+;; Update: the inlined sequence of callsite ids was manually modified to
+;; include some recursion, which reproduced an error before it was fixed.
+;;
 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
 ; RUN: opt -passes=memprof-context-disambiguation -supports-hot-cold-new \
@@ -96,13 +99,13 @@ attributes #7 = { builtin }
 !6 = !{i32 7, !"frame-pointer", i32 2}
 !7 = !{!8, !10}
 !8 = !{!9, !"notcold"}
-!9 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 8632435727821051414}
+!9 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 -5964873800580613432, i64 8632435727821051414}
 !10 = !{!11, !"cold"}
-!11 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 -3421689549917153178}
+!11 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 -5964873800580613432, i64 -3421689549917153178}
 !12 = !{i64 9086428284934609951}
 !13 = !DIBasicType(name: "char", size: 8, encoding: DW_ATE_signed_char)
-!14 = !{i64 -5964873800580613432, i64 2732490490862098848, i64 8632435727821051414}
-!15 = !{i64 -5964873800580613432, i64 2732490490862098848, i64 -3421689549917153178}
+!14 = !{i64 -5964873800580613432, i64 2732490490862098848, i64 -5964873800580613432, i64 8632435727821051414}
+!15 = !{i64 -5964873800580613432, i64 2732490490862098848, i64 -5964873800580613432, i64 -3421689549917153178}
 
 
 ; DUMP: CCG before cloning:

[snehasish]

lgtm