[llvm] [mlgo][coro] Assign coro split-ed functions a `FunctionLevel` (PR #68263)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/68263

>From e6a05bcb9fa5258325834cc0b999f76d766bc735 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Wed, 4 Oct 2023 14:37:16 -0700
Subject: [PATCH 1/2] [mlgo][coro] Fix Issue #62616

The coroutine handle methods needed to be added to the `FunctionLevels`. We already had logic for their appearing between calls into the `Inliner` pass.
---
 llvm/lib/Analysis/MLInlineAdvisor.cpp         | 15 ++++++-
 .../Inline/ML/coro-split-func-levels.ll       | 42 +++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll

diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 0660a9993b6dd7b..75eb8ece2e447e2 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -192,7 +192,9 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
   // - in addition, if new Nodes were created by a pass (e.g. CoroSplit),
   // they'd be adjacent to Nodes in the last SCC. So we just need to check the
   // boundary of Nodes in NodesInLastSCC for Nodes we haven't seen. We don't
-  // care about the nature of the Edge (call or ref).
+  // care about the nature of the Edge (call or ref). `FunctionLevels`-wise, we
+  // record them at the same level as the original node (this is a choice, may
+  // need revisiting).
   NodeCount -= static_cast<int64_t>(NodesInLastSCC.size());
   while (!NodesInLastSCC.empty()) {
     const auto *N = *NodesInLastSCC.begin();
@@ -204,12 +206,15 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
     }
     ++NodeCount;
     EdgeCount += getLocalCalls(N->getFunction());
+    const auto NLevel = FunctionLevels.at(N);
     for (const auto &E : *(*N)) {
       const auto *AdjNode = &E.getNode();
       assert(!AdjNode->isDead() && !AdjNode->getFunction().isDeclaration());
       auto I = AllNodes.insert(AdjNode);
-      if (I.second)
+      if (I.second) {
         NodesInLastSCC.insert(AdjNode);
+        FunctionLevels[AdjNode] = NLevel;
+      }
     }
   }
 
@@ -461,6 +466,12 @@ void MLInlineAdvisor::print(raw_ostream &OS) const {
     OS << "\n";
   }
   OS << "\n";
+  OS << "[MLInlineAdvisor] FuncLevels:\n";
+  for (auto I : FunctionLevels)
+    OS << (I.first->isDead() ? "<deleted>" : I.first->getFunction().getName())
+       << " : " << I.second << "\n";
+
+  OS << "\n";
 }
 
 MLInlineAdvice::MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB,
diff --git a/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll b/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
new file mode 100644
index 000000000000000..b07536285dcec9b
--- /dev/null
+++ b/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
@@ -0,0 +1,42 @@
+
+; REQUIRES: llvm_inliner_model_autogenerated
+; RUN: opt -S -passes='coro-early,scc-oz-module-inliner,print<inline-advisor>' \
+; RUN:  -enable-ml-inliner=release -keep-inline-advisor-for-printing < %s
+
+define void @_Z5get_sv() presplitcoroutine {
+  %1 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %2 = call ptr @llvm.coro.begin(token %1, ptr null)
+  %3 = call token @llvm.coro.save(ptr null)
+  %4 = call i8 @llvm.coro.suspend(token none, i1 false)
+  call void @_ZN1S12promise_typeD2Ev()
+  ret void
+}
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
+declare ptr @llvm.coro.begin(token, ptr writeonly)
+declare token @llvm.coro.save(ptr)
+declare i8 @llvm.coro.suspend(token, i1)
+
+declare void @__clang_call_terminate()
+
+define void @_ZN1S12promise_typeD2Ev() personality ptr null {
+  invoke void @_Z4funcv()
+          to label %1 unwind label %2
+
+1:                                                ; preds = %0
+  ret void
+
+2:                                                ; preds = %0
+  %3 = landingpad { ptr, i32 }
+          catch ptr null
+  call void @__clang_call_terminate()
+  unreachable
+}
+declare void @_Z4funcv()
+
+; CHECK:      [MLInlineAdvisor] FuncLevels:
+; CHECK-NEXT: _Z5get_sv : 1
+; CHECK-NEXT: _ZN1S12promise_typeD2Ev : 0
+; CHECK-NEXT: _Z5get_sv.resume : 1
+; CHECK-NEXT: _Z5get_sv.destroy : 1
+; CHECK-NEXT: _Z5get_sv.cleanup : 1
\ No newline at end of file

>From 4d14c3af5f2239c1309b367390bba7c155ebc02c Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Wed, 4 Oct 2023 15:07:20 -0700
Subject: [PATCH 2/2] Deleted trailing line / added eof line.

---
 llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll b/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
index b07536285dcec9b..79e1ebeec17b929 100644
--- a/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
+++ b/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
@@ -1,4 +1,3 @@
-
 ; REQUIRES: llvm_inliner_model_autogenerated
 ; RUN: opt -S -passes='coro-early,scc-oz-module-inliner,print<inline-advisor>' \
 ; RUN:  -enable-ml-inliner=release -keep-inline-advisor-for-printing < %s
@@ -39,4 +38,4 @@ declare void @_Z4funcv()
 ; CHECK-NEXT: _ZN1S12promise_typeD2Ev : 0
 ; CHECK-NEXT: _Z5get_sv.resume : 1
 ; CHECK-NEXT: _Z5get_sv.destroy : 1
-; CHECK-NEXT: _Z5get_sv.cleanup : 1
\ No newline at end of file
+; CHECK-NEXT: _Z5get_sv.cleanup : 1