[Coroutines] Allocas used after @coro.end should not go in the coro frame

Because the coro frame may not exist anymore at that point. See the bug
for a motivating example.

Fixes #127499

Author: zmodem

llvm/lib/Transforms/Coroutines/SpillUtils.cpp

@@ -162,6 +162,16 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
 
   void visit(Instruction &I) {
     Users.insert(&I);
+
+    for (const Instruction *CoroEnd : CoroShape.CoroEnds) {
+      if (DT.dominates(CoroEnd, &I)) {
+        // Data accessed after coro.end (such as the return object and callee
+        // destroyed parameters) should always go on the stack.
+        ShouldLiveOnFrame = false;
+        return;
+      }
+    }
+
     Base::visit(I);
     // If the pointer is escaped prior to CoroBegin, we have to assume it would
     // be written into before CoroBegin as well.

llvm/test/Transforms/Coroutines/callee-destructed-param.ll

@@ -0,0 +1,94 @@
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+; Inspired by coro-param-copy.ll and the following:
+;
+;   $ clang++ -std=c++20 -S -emit-llvm -g0 -fno-exceptions \
+;         -mllvm -disable-llvm-optzns -fno-discard-value-names a.cc -o -
+;
+;   #include <coroutine>
+;
+;   class BasicCoroutine {
+;   public:
+;     struct Promise {
+;       BasicCoroutine get_return_object();
+;       void unhandled_exception() noexcept;
+;       void return_void() noexcept;
+;       std::suspend_never initial_suspend() noexcept;
+;       std::suspend_never final_suspend() noexcept;
+;     };
+;     using promise_type = Promise;
+;   };
+;
+;   struct [[clang::trivial_abi]] Trivial {
+;     Trivial(int x) : x(x) {}
+;     ~Trivial();
+;     int x;
+;   };
+;
+;   BasicCoroutine coro(Trivial t) {
+;     co_return;
+;   }
+;
+;
+; Check that even though %x.local may escape via use() in the beginning of @f,
+; it is not put in the corountine frame, since %x.local is used after
+; @llvm.coro.end, at which point the coroutine frame may have been deallocated.
+;
+; In the program above, a move constructor (or just memcpy) is invoked to copy
+; t to a coroutine-local alloca. At the end of the function, t's destructor is
+; called because of trivial_abi. At that point, t must not have been stored in
+; the coro frame.
+
+; The frame should not contain an i64.
+; CHECK: %f.Frame = type { ptr, ptr, i1 }
+
+; Check that we have both uses of %x.local (and they're not using the frame).
+; CHECK-LABEL: define ptr @f(i64 %x)
+; CHECK: call void @use(ptr %x.local)
+; CHECK: call void @use(ptr %x.local)
+
+
+define ptr @f(i64 %x) presplitcoroutine {
+entry:
+  %x.local = alloca i64
+  store i64 %x, ptr %x.local
+  br label %coro.alloc
+
+coro.alloc:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @myAlloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  call void @use(ptr %x.local)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %cleanup]
+resume:
+  br label %cleanup
+
+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)
+  call void @use(ptr %x.local)  ; It better not be on the frame, that's gone.
+  ret ptr %hdl
+}
+
+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 @myAlloc(i32)
+declare void @use(ptr)
+declare void @free(ptr)

llvm/test/Transforms/Coroutines/coro-debug.ll

@@ -64,7 +64,7 @@ indirect.dest:
   br label %coro_Cleanup
 
 coro_Cleanup:                                     ; preds = %sw.epilog, %sw.bb1
-  %5 = load ptr, ptr %coro_hdl, align 8, !dbg !24
+  %5 = load ptr, ptr %coro_hdl, align 8, !dbg !24  ; XXX: Here %coro_hdl is used between a coro.suspend and coro.end, which seems wrong?
   %6 = call ptr @llvm.coro.free(token %0, ptr %5), !dbg !24
   call void @free(ptr %6), !dbg !24
   call void @llvm.dbg.declare(metadata i32 %asm_res, metadata !32, metadata !13), !dbg !16
@@ -172,14 +172,14 @@ attributes #7 = { noduplicate }
 !32 = !DILocalVariable(name: "inline_asm", scope: !6, file: !7, line: 55, type: !11)
 
 ; CHECK: define ptr @f(i32 %x) #0 personality i32 0 !dbg ![[ORIG:[0-9]+]]
-; CHECK: define internal fastcc void @f.resume(ptr noundef nonnull align 8 dereferenceable(40) %0) #0 personality i32 0 !dbg ![[RESUME:[0-9]+]]
+; CHECK: define internal fastcc void @f.resume(ptr noundef nonnull align 8 dereferenceable(32) %0) #0 personality i32 0 !dbg ![[RESUME:[0-9]+]]
 ; CHECK: entry.resume:
 ; CHECK: %[[DBG_PTR:.*]] = alloca ptr
-; CHECK: #dbg_declare(ptr %[[DBG_PTR]], ![[RESUME_COROHDL:[0-9]+]], !DIExpression(DW_OP_deref, DW_OP_plus_uconst,
 ; CHECK: #dbg_declare(ptr %[[DBG_PTR]], ![[RESUME_X:[0-9]+]], !DIExpression(DW_OP_deref, DW_OP_plus_uconst, [[EXPR_TAIL:.*]])
 ; CHECK: #dbg_declare(ptr %[[DBG_PTR]], ![[RESUME_DIRECT:[0-9]+]], !DIExpression(DW_OP_deref, DW_OP_plus_uconst, [[EXPR_TAIL]])
 ; CHECK: store ptr {{.*}}, ptr %[[DBG_PTR]]
-; CHECK-NOT: alloca ptr
+; CHECK: alloca ptr
+; CHECK: #dbg_declare(ptr %[[CORO_HDL:.*]], ![[RESUME_COROHDL:[0-9]+]], !DIExpression()
 ; CHECK: #dbg_declare(i8 0, ![[RESUME_CONST:[0-9]+]], !DIExpression(DW_OP_LLVM_convert, 8, DW_ATE_signed, DW_OP_LLVM_convert, 32, DW_ATE_signed),
 ; Note that keeping the undef value here could be acceptable, too.
 ; CHECK-NOT: #dbg_declare(ptr undef, !{{[0-9]+}}, !DIExpression(),
@@ -195,15 +195,15 @@ attributes #7 = { noduplicate }
 ; CHECK: [[DEFAULT_DEST]]:
 ; CHECK-NOT: {{.*}}:
 ; CHECK: #dbg_declare(i32 %[[CALLBR_RES]]
-; CHECK: define internal fastcc void @f.destroy(ptr noundef nonnull align 8 dereferenceable(40) %0) #0 personality i32 0 !dbg ![[DESTROY:[0-9]+]]
-; CHECK: define internal fastcc void @f.cleanup(ptr noundef nonnull align 8 dereferenceable(40) %0) #0 personality i32 0 !dbg ![[CLEANUP:[0-9]+]]
+; CHECK: define internal fastcc void @f.destroy(ptr noundef nonnull align 8 dereferenceable(32) %0) #0 personality i32 0 !dbg ![[DESTROY:[0-9]+]]
+; CHECK: define internal fastcc void @f.cleanup(ptr noundef nonnull align 8 dereferenceable(32) %0) #0 personality i32 0 !dbg ![[CLEANUP:[0-9]+]]
 
 ; CHECK: ![[ORIG]] = distinct !DISubprogram(name: "f", linkageName: "flink"
 
 ; CHECK: ![[RESUME]] = distinct !DISubprogram(name: "f", linkageName: "flink"
-; CHECK: ![[RESUME_COROHDL]] = !DILocalVariable(name: "coro_hdl", scope: ![[RESUME]]
 ; CHECK: ![[RESUME_X]] = !DILocalVariable(name: "x", arg: 1, scope: ![[RESUME]]
 ; CHECK: ![[RESUME_DIRECT]] = !DILocalVariable(name: "direct_mem", scope: ![[RESUME]]
+; CHECK: ![[RESUME_COROHDL]] = !DILocalVariable(name: "coro_hdl", scope: ![[RESUME]]
 ; CHECK: ![[RESUME_CONST]] = !DILocalVariable(name: "direct_const", scope: ![[RESUME]]
 ; CHECK: ![[RESUME_DIRECT_VALUE]] = !DILocalVariable(name: "direct_value", scope: ![[RESUME]]
 

llvm/test/Transforms/Coroutines/coro-lifetime-end.ll

@@ -50,16 +50,18 @@ exit:
 define void @LifetimeEndAfterCoroEnd() presplitcoroutine {
 ; CHECK-LABEL: define void @LifetimeEndAfterCoroEnd() {
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TESTVAL:%.*]] = alloca
 ; CHECK-NEXT:    [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @LifetimeEndAfterCoroEnd.resumers)
 ; CHECK-NEXT:    [[ALLOC:%.*]] = call ptr @malloc(i64 16)
 ; CHECK-NEXT:    [[VFRAME:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
 ; CHECK-NEXT:    store ptr @LifetimeEndAfterCoroEnd.resume, ptr [[VFRAME]], align 8
 ; CHECK-NEXT:    [[DESTROY_ADDR:%.*]] = getelementptr inbounds nuw [[LIFETIMEENDAFTERCOROEND_FRAME:%.*]], ptr [[VFRAME]], i32 0, i32 1
 ; CHECK-NEXT:    store ptr @LifetimeEndAfterCoroEnd.destroy, ptr [[DESTROY_ADDR]], align 8
-; CHECK-NEXT:    [[INDEX_ADDR1:%.*]] = getelementptr inbounds [[LIFETIMEENDAFTERCOROEND_FRAME]], ptr [[VFRAME]], i32 0, i32 2
-; CHECK-NEXT:    call void @consume.i8.array(ptr [[INDEX_ADDR1]])
-; CHECK-NEXT:    [[INDEX_ADDR2:%.*]] = getelementptr inbounds nuw [[LIFETIMEENDAFTERCOROEND_FRAME]], ptr [[VFRAME]], i32 0, i32 3
-; CHECK-NEXT:    store i1 false, ptr [[INDEX_ADDR2]], align 1
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 100, ptr [[TESTVAL]])
+; CHECK-NEXT:    call void @consume.i8.array(ptr [[TESTVAL]])
+; CHECK-NEXT:    [[INDEX_ADDR1:%.*]] = getelementptr inbounds nuw [[LIFETIMEENDAFTERCOROEND_FRAME]], ptr [[VFRAME]], i32 0, i32 2
+; CHECK-NEXT:    store i1 false, ptr [[INDEX_ADDR1]], align 1
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 100, ptr [[TESTVAL]])
 ; CHECK-NEXT:    ret void
 ;
 entry: