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

[Hans Wennborg via llvm-commits]

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

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

>From f9bd3480761edc6ba667e694c1639e6eaf701cd1 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Wed, 26 Mar 2025 13:09:42 +0100
Subject: [PATCH] [Coro] Allow spilling @llvm.coro.suspend() to the coro frame

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
---
 llvm/lib/Transforms/Coroutines/SpillUtils.cpp | 11 ++--
 .../Coroutines/coro-spill-suspend.ll          | 59 +++++++++++++++++++
 2 files changed, 64 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/Transforms/Coroutines/coro-spill-suspend.ll

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)