[llvm] [Coroutines] properly update CallGraph in CoroSplit (PR #107935)

Yuxuan Chen via llvm-commits llvm-commits at lists.llvm.org
Thu Sep 12 10:02:33 PDT 2024


https://github.com/yuxuanchen1997 updated https://github.com/llvm/llvm-project/pull/107935

>From 00da9849000624ae83ceee8a5629b6b5e2b71114 Mon Sep 17 00:00:00 2001
From: Yuxuan Chen <ych at meta.com>
Date: Mon, 9 Sep 2024 16:14:15 -0700
Subject: [PATCH 1/2] [Coroutines] Make CoroSplit properly update CallGraph

---
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp  | 18 ++++++---
 .../Coroutines/gh107139-split-in-scc.ll       | 38 +++++++++++++++++++
 2 files changed, 50 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index dc3829d7f28eb1..8ea460badaad5d 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -2080,12 +2080,13 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   return Shape;
 }
 
-static void updateCallGraphAfterCoroutineSplit(
+static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
     LazyCallGraph::Node &N, const coro::Shape &Shape,
     const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
     LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
     FunctionAnalysisManager &FAM) {
 
+  auto *CurrentSCC = &C;
   if (!Clones.empty()) {
     switch (Shape.ABI) {
     case coro::ABI::Switch:
@@ -2105,13 +2106,16 @@ static void updateCallGraphAfterCoroutineSplit(
     }
 
     // Let the CGSCC infra handle the changes to the original function.
-    updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR, FAM);
+    CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,
+                                                         UR, FAM);
   }
 
   // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges
   // to the split functions.
   postSplitCleanup(N.getFunction());
-  updateCGAndAnalysisManagerForFunctionPass(CG, C, N, AM, UR, FAM);
+  CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,
+                                                          AM, UR, FAM);
+  return *CurrentSCC;
 }
 
 /// Replace a call to llvm.coro.prepare.retcon.
@@ -2200,6 +2204,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
   if (Coroutines.empty() && PrepareFns.empty())
     return PreservedAnalyses::all();
 
+  auto *CurrentSCC = &C;
   // Split all the coroutines.
   for (LazyCallGraph::Node *N : Coroutines) {
     Function &F = N->getFunction();
@@ -2211,7 +2216,8 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     coro::Shape Shape =
         splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
-    updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
+    CurrentSCC = &updateCallGraphAfterCoroutineSplit(
+        *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);
 
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     ORE.emit([&]() {
@@ -2223,14 +2229,14 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
 
     if (!Shape.CoroSuspends.empty()) {
       // Run the CGSCC pipeline on the original and newly split functions.
-      UR.CWorklist.insert(&C);
+      UR.CWorklist.insert(CurrentSCC);
       for (Function *Clone : Clones)
         UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));
     }
   }
 
   for (auto *PrepareFn : PrepareFns) {
-    replaceAllPrepares(PrepareFn, CG, C);
+    replaceAllPrepares(PrepareFn, CG, *CurrentSCC);
   }
 
   return PreservedAnalyses::none();
diff --git a/llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll b/llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll
new file mode 100644
index 00000000000000..1fae4d2c21e80f
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll
@@ -0,0 +1,38 @@
+; Verify that we don't crash on mutually recursive coroutines
+; RUN: opt < %s -passes='cgscc(coro-split)' -S | FileCheck %s
+
+target triple = "x86_64-redhat-linux-gnu"
+
+; CHECK-LABEL: define void @foo
+define void @foo() presplitcoroutine personality ptr null {
+entry:
+
+  %0 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %1 = call ptr @llvm.coro.begin(token %0, ptr null)
+  %2 = call token @llvm.coro.save(ptr null)
+  %3 = call i8 @llvm.coro.suspend(token none, i1 false)
+  %4 = call token @llvm.coro.save(ptr null)
+  ; CHECK: call void @bar(ptr null, ptr null)
+  call void @llvm.coro.await.suspend.void(ptr null, ptr null, ptr @bar)
+  ret void
+}
+
+; CHECK-LABEL: define void @bar({{.*}})
+define void @bar(ptr %0, ptr %1) {
+entry:
+  ; CHECK: call void @foo()
+  call void @foo()
+  ret void
+}
+
+; CHECK-LABEL: @foo.resume({{.*}})
+; CHECK-LABEL: @foo.destroy({{.*}})
+; CHECK-LABEL: @foo.cleanup({{.*}})
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) #0
+declare ptr @llvm.coro.begin(token, ptr writeonly) nounwind
+declare token @llvm.coro.save(ptr) nomerge nounwind
+declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
+declare i8 @llvm.coro.suspend(token, i1) nounwind
+
+attributes #0 = { nocallback nofree nosync nounwind willreturn memory(argmem: read) }

>From a6220677fc57b50e4d12f89c323bca4f65686b04 Mon Sep 17 00:00:00 2001
From: Yuxuan Chen <ych at meta.com>
Date: Thu, 12 Sep 2024 10:02:01 -0700
Subject: [PATCH 2/2] use reviewer suggested style

---
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 8ea460badaad5d..5cce80b601dc60 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -2080,13 +2080,12 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   return Shape;
 }
 
-static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
+static void updateCallGraphAfterCoroutineSplit(
     LazyCallGraph::Node &N, const coro::Shape &Shape,
-    const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
+    const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC *&C,
     LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
     FunctionAnalysisManager &FAM) {
 
-  auto *CurrentSCC = &C;
   if (!Clones.empty()) {
     switch (Shape.ABI) {
     case coro::ABI::Switch:
@@ -2106,16 +2105,13 @@ static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
     }
 
     // Let the CGSCC infra handle the changes to the original function.
-    CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,
-                                                         UR, FAM);
+    C = &updateCGAndAnalysisManagerForCGSCCPass(CG, *C, N, AM, UR, FAM);
   }
 
   // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges
   // to the split functions.
   postSplitCleanup(N.getFunction());
-  CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,
-                                                          AM, UR, FAM);
-  return *CurrentSCC;
+  C = &updateCGAndAnalysisManagerForFunctionPass(CG, *C, N, AM, UR, FAM);
 }
 
 /// Replace a call to llvm.coro.prepare.retcon.
@@ -2216,8 +2212,8 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     coro::Shape Shape =
         splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
-    CurrentSCC = &updateCallGraphAfterCoroutineSplit(
-        *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);
+    updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, CurrentSCC, CG, AM,
+                                       UR, FAM);
 
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     ORE.emit([&]() {



More information about the llvm-commits mailing list