[Coroutines] suspend metadata preservation

[yonillasky]

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.

Discourse post: https://discourse.llvm.org/t/metadata-preservation-in-coroutines/87494

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[Dinistro]

Dropped a few nit

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-coroutines

Author: None (yonillasky)

Changes

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.

Discourse post: https://discourse.llvm.org/t/metadata-preservation-in-coroutines/87494

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Coroutines/CoroSplit.cpp (+24)
• (added) llvm/test/Transforms/Coroutines/coro-split-md.ll (+69)

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 64b33e46404f0..6f0a43ca7cd80 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -79,6 +79,18 @@ using namespace llvm;
 
 #define DEBUG_TYPE "coro-split"
 
+/// If set, ensures that all metadata from CoroSuspendInst's is preserved in the
+/// containing function.
+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";
+
 // 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 +1515,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 +1602,10 @@ struct SwitchCoroutineSplitter {
       ++SuspendIndex;
     }
 
+    if (CoroSplitPreservesSuspendMD)
+      if (!SuspendMdEntries.empty())
+        F.setMetadata(CoroSuspendMDTableName, MDNode::get(C, SuspendMdEntries));
+
     Builder.SetInsertPoint(UnreachBB);
     Builder.CreateUnreachable();
     DBuilder.finalize();
diff --git a/llvm/test/Transforms/Coroutines/coro-split-md.ll b/llvm/test/Transforms/Coroutines/coro-split-md.ll
new file mode 100644
index 0000000000000..edf0a43e11685
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-split-md.ll
@@ -0,0 +1,69 @@
+; Tests that coro-split pass preserves metadata on suspend points.
+; RUN: opt < %s -coro-split-preserves-suspend-md -passes='cgscc(coro-split)' -S | FileCheck %s
+
+define ptr @f() presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %need.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.alloc, label %dyn.alloc, label %begin
+
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  br label %begin
+
+begin:
+  %phi = phi ptr [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %phi)
+  call void @print(i32 0)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false), !llvm.coro.suspend_md !0
+  switch i8 %0, label %suspend [i8 0, label %resume1
+                                i8 1, label %cleanup]
+resume1:
+  call void @print(i32 1)
+  %1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %1, label %suspend [i8 0, label %resume2
+                                i8 1, label %cleanup]
+resume2:
+  call void @print(i32 2)
+  %2 = call i8 @llvm.coro.suspend(token none, i1 true), !llvm.coro.suspend_md !1
+  switch i8 %2, label %suspend [i8 0, label %trap
+                                i8 1, label %cleanup]
+trap:
+  call void @llvm.trap()
+  unreachable
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret ptr %hdl
+}
+
+; CHECK: define ptr @f() !llvm.coro.suspend_md_table ![[MD_TABLE:[0-9]+]]
+
+; CHECK-DAG: ![[MD_TABLE]] = !{![[S0_ROW:[0-9]+]], ![[S1_ROW:[0-9]+]]}
+; CHECK-DAG: ![[S0_ROW]] = !{i2 0, ![[S0_MD:[0-9]+]]}
+; CHECK-DAG: ![[S0_MD]] = !{!"waiting"}
+; CHECK-DAG: ![[S1_ROW]] = !{i2 -2, ![[S1_MD:[0-9]+]]}
+; CHECK-DAG: ![[S1_MD]] = !{!"done"}
+
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(ptr)
+declare void @llvm.coro.destroy(ptr)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.end(ptr, i1, token)
+
+declare noalias ptr @malloc(i32) allockind("alloc,uninitialized") "alloc-family"="malloc"
+declare void @print(i32)
+declare void @free(ptr) willreturn allockind("free") "alloc-family"="malloc"
+
+!0 = !{!"waiting"}
+!1 = !{!"done"}
+attributes #1 = { coro_elide_safe }

[yonillasky]

A downstream solution exists, will close this PR.