[llvm] [DRAFT] coros: suspend metadata preservation (PR #150077)

[via llvm-commits]

https://github.com/yonillasky created https://github.com/llvm/llvm-project/pull/150077

This change allows earlier compiler passes to bind metadata to individual `llvm.coro.suspend` intrinsics, and then have that metadata survive lowering, in the form of a table mapping SuspendPointIdx to the original suspend MD, bound to the coroutine function itself, so that it can be utilized by later compiler passes.


>From 136f9fb172eb91390dbb4bf277d0f1ec0dddc98c Mon Sep 17 00:00:00 2001
From: Yoni Lavi <yoni.lavi at nextsilicon.com>
Date: Tue, 22 Jul 2025 21:11:28 +0300
Subject: [PATCH] coros: suspend metadata preservation

---
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 26 ++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 64b33e46404f0..c0d116cd81e91 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -79,6 +79,20 @@ using namespace llvm;
 
 #define DEBUG_TYPE "coro-split"
 
+/// If set, ensures that all metadata from CoroSuspendInst's is preserved in the
+/// containing function.
+namespace llvm {
+static cl::opt<bool> CoroSplitPreservesSuspendMD(
+  "coro-split-preserves-suspend-md", cl::Hidden,
+  cl::desc("llvm.coro.suspend_md metadata from all suspend point instructions "
+          "will be preserved inside llvm.coro.suspend_md_table metadata on the "
+          "containing coroutine"));
+
+static StringRef CoroSuspendMDName = "llvm.coro.suspend_md";
+static StringRef CoroSuspendMDTableName = "llvm.coro.suspend_md_table";
+
+} // namespace llvm
+
 // FIXME:
 // Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape
 // and it is known that other transformations, for example, sanitizers
@@ -1503,10 +1517,18 @@ struct SwitchCoroutineSplitter {
     Shape.SwitchLowering.ResumeSwitch = Switch;
 
     // Split all coro.suspend calls
+    SmallVector<Metadata *> suspendMDEntries;
     size_t SuspendIndex = 0;
     for (auto *AnyS : Shape.CoroSuspends) {
       auto *S = cast<CoroSuspendInst>(AnyS);
       ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);
+      if (CoroSplitPreservesSuspendMD) {
+        MDNode *SuspendMD = S->getMetadata(CoroSuspendMDName);
+        if (SuspendMD) {
+          Metadata *KeyMD = ConstantAsMetadata::get(IndexVal);
+          suspendMDEntries.push_back(MDNode::get(C, {KeyMD, SuspendMD}));
+        }
+      }
 
       // Replace CoroSave with a store to Index:
       //    %index.addr = getelementptr %f.frame... (index field number)
@@ -1582,6 +1604,10 @@ struct SwitchCoroutineSplitter {
       ++SuspendIndex;
     }
 
+    if (CoroSplitPreservesSuspendMD)
+      if (!suspendMDEntries.empty())
+        F.setMetadata(CoroSuspendMDTableName, MDNode::get(C, suspendMDEntries));
+
     Builder.SetInsertPoint(UnreachBB);
     Builder.CreateUnreachable();
     DBuilder.finalize();