[Coroutines] Handle the writes to promise alloca prior to

llvm.coro.begin

Previously we've taken care of the writes to allocas prior to llvm.coro.begin.
However, since the promise alloca is special so that we never handled it
before. For the long time, since the programmers can't access the
promise_type due to the c++ language specification, we still failed to
recognize the problem until a recent report:
https://github.com/llvm/llvm-project/issues/57861

And we've tested many codes that the problem gone away after we handle
the writes to the promise alloca prior to @llvm.coro.begin()
prope until a recent report:
https://github.com/llvm/llvm-project/issues/57861

And we've tested many codes that the problem gone away after we handle
the writes to the promise alloca prior to @llvm.coro.begin() properly.

Closes https://github.com/llvm/llvm-project/issues/57861

2 files changed, 105 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index b8f5c9d..4973ecd 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1817,6 +1817,47 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
           AliasPtrTyped, [&](Use &U) { return DT.dominates(CB, U); });
     }
   }
+
+  // PromiseAlloca is not collected in FrameData.Allocas. So we don't handle
+  // the case that the PromiseAlloca may have writes before CoroBegin in the
+  // above codes. And it may be problematic in edge cases. See
+  // https://github.com/llvm/llvm-project/issues/57861 for an example.
+  if (Shape.ABI == coro::ABI::Switch && Shape.SwitchLowering.PromiseAlloca) {
+    AllocaInst *PA = Shape.SwitchLowering.PromiseAlloca;
+    // If there is memory accessing to promise alloca before CoroBegin;
+    bool HasAccessingPromiseBeforeCB = llvm::any_of(PA->uses(), [&](Use &U) {
+      auto *Inst = dyn_cast<Instruction>(U.getUser());
+      if (!Inst || DT.dominates(CB, Inst))
+        return false;
+
+      if (auto *CI = dyn_cast<CallInst>(Inst)) {
+        // It is fine if the call wouldn't write to the Promise.
+        // This is possible for @llvm.coro.id intrinsics, which
+        // would take the promise as the second argument as a
+        // marker.
+        if (CI->onlyReadsMemory() ||
+            CI->onlyReadsMemory(CI->getArgOperandNo(&U)))
+          return false;
+        return true;
+      }
+
+      return isa<StoreInst>(Inst) ||
+             // It may take too much time to track the uses.
+             // Be conservative about the case the use may escape.
+             isa<GetElementPtrInst>(Inst) ||
+             // There would always be a bitcast for the promise alloca
+             // before we enabled Opaque pointers. And now given
+             // opaque pointers are enabled by default. This should be
+             // fine.
+             isa<BitCastInst>(Inst);
+    });
+    if (HasAccessingPromiseBeforeCB) {
+      Builder.SetInsertPoint(Shape.getInsertPtAfterFramePtr());
+      auto *G = GetFramePointer(PA);
+      auto *Value = Builder.CreateLoad(PA->getAllocatedType(), PA);
+      Builder.CreateStore(Value, G);
+    }
+  }
 }
 
 // Moves the values in the PHIs in SuccBB that correspong to PredBB into a new
diff --git a/llvm/test/Transforms/Coroutines/coro-spill-promise-02.ll b/llvm/test/Transforms/Coroutines/coro-spill-promise-02.ll
new file mode 100644
index 0000000..42d16b9
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-spill-promise-02.ll
@@ -0,0 +1,64 @@
+; Check that we would take care of the value written to promise before @llvm.coro.begin.
+; RUN: opt < %s -opaque-pointers -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+%"class.task::promise_type" = type { [64 x i8] }
+
+declare void @consume(i32*)
+declare void @consume2(%"class.task::promise_type"*)
+
+define ptr @f() presplitcoroutine {
+entry:
+  %data = alloca i32, align 4
+  %__promise = alloca %"class.task::promise_type", align 64
+  %id = call token @llvm.coro.id(i32 0, ptr %__promise, ptr null, ptr null)
+  call void @consume2(%"class.task::promise_type"* %__promise)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+  call void @consume(i32* %data)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %cleanup]
+resume:
+  call void @consume(i32* %data)
+  call void @consume2(%"class.task::promise_type"* %__promise)
+  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)
+  ret ptr %hdl
+}
+
+; CHECK-LABEL: %f.Frame = type { ptr, ptr, i32, i1, [43 x i8], %"class.task::promise_type" }
+
+; CHECK-LABEL: @f(
+; CHECK: %__promise = alloca %"class.task::promise_type"
+; CHECK: call void @consume2(ptr %__promise)
+; CHECK: call{{.*}}@llvm.coro.begin
+; CHECK: %[[PROMISE_ADDR:.+]] = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 5
+; CHECK: %[[PROMISE_VALUE:.+]] = load %"class.task::promise_type", ptr %__promise
+; CHECK: store %"class.task::promise_type" %[[PROMISE_VALUE]], ptr %[[PROMISE_ADDR]]
+
+; CHECK-LABEL: @f.resume(
+; CHECK: %[[DATA:.+]] = getelementptr inbounds %f.Frame, ptr %hdl, i32 0, i32 5
+; CHECK: call void @consume2(ptr %[[DATA]])
+; CHECK: ret void
+
+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)
+
+declare noalias ptr @malloc(i32)
+declare double @print(double)
+declare void @free(ptr)