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

[via llvm-commits]

Author: Hans Wennborg
Date: 2025-03-27T10:23:30+01:00
New Revision: 00c527abab3dfea5f3122f8151d69e77655eb033

URL: https://github.com/llvm/llvm-project/commit/00c527abab3dfea5f3122f8151d69e77655eb033
DIFF: https://github.com/llvm/llvm-project/commit/00c527abab3dfea5f3122f8151d69e77655eb033.diff

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

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

In the linked bug, we hit a situation like this:

```
  %s = call @llvm.coro.suspend(...)
             |
       switch (%s)
case v1: /       \ case v2:
        ...      ...
         |       suspend point
         |       ...
         \       /
 %x = phi [v1] [v2]
             |
            ...
             |
         use(%x)
```

Instcombine will notice that %x correlates exactly with %s, and so
use(%x) becomes use(%s).

However, corosplit would substitute different values for %s when
splitting the function, so even though %s had a particular value when
control actually passed through the switch, it could have a *different*
value when reaching use(%s).

This illustrates that while IR from the frontend typically does not use
these suspend return values across suspend points, mid-level
optimizations on the presplit coroutine may introduce new uses of
suspend values, so they must be considered eligible to spill to the
coroutine frame.

Fixes: #130326

Added: 
    llvm/test/Transforms/Coroutines/coro-spill-suspend.ll

Modified: 
    llvm/lib/Transforms/Coroutines/SpillUtils.cpp

Removed: 
    


################################################################################
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)