[llvm] [Coro] Allow spilling @llvm.coro.suspend() to the coro frame (PR #133088)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-coroutines

Author: Hans Wennborg (zmodem)

<details>
<summary>Changes</summary>

It was excluded from spilling in a263a60, possibly by accident.

Since the return value is not constant throughout the coroutine and may get used across suspend points, it must be eligible for spilling like other values, as demonstrated by the bug.

Fixes: #<!-- -->130326

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Coroutines/SpillUtils.cpp (+5-6) 
- (added) llvm/test/Transforms/Coroutines/coro-spill-suspend.ll (+59) 


``````````diff
diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
index 5062ee97a665d..b3e5b7fa6e0b5 100644
--- a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
+++ b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
@@ -24,11 +24,10 @@ namespace {
 
 typedef SmallPtrSet<BasicBlock *, 8> VisitedBlocksSet;
 
-// Check for structural coroutine intrinsics that should not be spilled into
-// the coroutine frame.
-static bool isCoroutineStructureIntrinsic(Instruction &I) {
-  return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I) ||
-         isa<CoroSuspendInst>(&I);
+static bool isNonSpilledIntrinsic(Instruction &I) {
+  // Structural coroutine intrinsics that should not be spilled into the
+  // coroutine frame.
+  return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I);
 }
 
 /// Does control flow starting at the given block ever reach a suspend
@@ -467,7 +466,7 @@ void collectSpillsAndAllocasFromInsts(
   for (Instruction &I : instructions(F)) {
     // Values returned from coroutine structure intrinsics should not be part
     // of the Coroutine Frame.
-    if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
+    if (isNonSpilledIntrinsic(I) || &I == Shape.CoroBegin)
       continue;
 
     // Handle alloca.alloc specially here.
diff --git a/llvm/test/Transforms/Coroutines/coro-spill-suspend.ll b/llvm/test/Transforms/Coroutines/coro-spill-suspend.ll
new file mode 100644
index 0000000000000..8de02c8b7de23
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-spill-suspend.ll
@@ -0,0 +1,59 @@
+; Check that the return value of @llvm.coro.suspend gets spilled to the frame
+; if it may be used across suspend points.
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+
+; %sp1 should be part of the frame (the i8 value).
+; CHECK: %f.Frame = type { ptr, ptr, i1, i8 }
+
+; If the coro resumes, %sp1 is set to 0.
+; CHECK-LABEL: define{{.*}} void @f.resume
+; CHECK: AfterCoroSuspend:
+; CHECK: %sp1.spill.addr = getelementptr inbounds %f.Frame
+; CHECK: store i8 0, ptr %sp1.spill.addr
+
+; In the coro destroy function, %sp1 is reloaded from the frame. Its value
+; depends on whether the coroutine was resumed or not.
+; CHECK-LABEL: define{{.*}} void @f.destroy
+; CHECK: cleanup:
+; CHECK: %sp1.reload.addr = getelementptr inbounds %f.Frame
+; CHECK: %sp1.reload = load i8, ptr %sp1.reload.addr
+; CHECK: call void @print(i8 %sp1.reload)
+
+
+define ptr @f(i32 %n) presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+
+  %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp1, label %suspend [i8 0, label %resume1
+                                  i8 1, label %cleanup]
+
+resume1:
+  %sp2 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp2, label %suspend [i8 0, label %resume2
+                                  i8 1, label %cleanup]
+
+resume2:
+  br label %cleanup
+
+cleanup:
+  ; This use of %sp1 may cross a suspend point (%sp2).
+  call void @print(i8 %sp1)
+
+  %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
+}
+
+
+declare noalias ptr @malloc(i32)
+declare void @print(i8)
+declare void @free(ptr)

``````````

</details>


https://github.com/llvm/llvm-project/pull/133088