[mlgo][coro] Assign coro split-ed functions a FunctionLevel

[mtrofin]

The coroutine handle methods needed to be added to the FunctionLevels. We already had logic for their appearing between calls into the Inliner pass.

Fixes Issue #62616

[llvmbot]

@llvm/pr-subscribers-llvm-analysis
@llvm/pr-subscribers-mlgo

@llvm/pr-subscribers-llvm-transforms

Changes

The coroutine handle methods needed to be added to the FunctionLevels. We already had logic for their appearing between calls into the Inliner pass.

---

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

2 Files Affected:

• (modified) llvm/lib/Analysis/MLInlineAdvisor.cpp (+13-2)
• (added) llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll (+42)

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

[aeubanks]

I'd make the title more descriptive and put the bug number in the description