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 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 SuspendMdEntries; size_t SuspendIndex = 0; for (auto *AnyS : Shape.CoroSuspends) { auto *S = cast(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 }