[Coroutine] Fix a bug where Coroutine incorrectly spills phi and invoke defs before CoroBegin

[lxfind]

Summary

When a spill definition is before CoroBegin, we cannot spill it to the frame immediately after the definition. We have to spill it after the frame is ready.
The current implementation handles it properly for any other kinds of instructions except for PhINode and InvokeInst, which could also be defined before CoroBegin.
This patch fixes it by moving the CoroBegin dominance check earlier, so that it covers all cases.
Added a test.

[lxfind]

fix lint

[junparser]

LGTM, Thanks for the fix.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

if (CurrentValue != E.def()) {
           // coro.begin from i8* to %f.frame*.
           InsertPt = FramePtr->getNextNode();
 
           // If we're spilling an Argument, make sure we clear 'nocapture'
           // from the coroutine function.
           Arg->getParent()->removeParamAttr(Arg->getArgNo(),
                                             Attribute::NoCapture);
 
-        } else if (auto *II = dyn_cast<InvokeInst>(CurrentValue)) {
-          // If we are spilling the result of the invoke instruction, split the
-          // normal edge and insert the spill in the new block.
-          auto NewBB = SplitEdge(II->getParent(), II->getNormalDest());
+        } else if (auto CSI = dyn_cast<AnyCoroSuspendInst>(CurrentValue)) {
+          // Don't spill immediately after a suspend; splitting assumes
+          // that the suspend will be followed by a branch.
+          InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHI();
+        } else {
+          auto *I = cast<Instruction>(CurrentValue);
+          if (!DT.dominates(CB, I)) {
+            // If it is not dominated by CoroBegin, then spill should be
+            // inserted immediately after CoroFrame is computed.
+            InsertPt = FramePtr->getNextNode();
+          } else if (auto *II = dyn_cast<InvokeInst>(I)) {
+            // If we are spilling the result of the invoke instruction, split
+            // the normal edge and insert the spill in the new block.
+            auto *NewBB = SplitEdge(II->getParent(), II->getNormalDest());
             InsertPt = NewBB->getTerminator();
-        } else if (isa<PHINode>(CurrentValue)) {
+          } else if (isa<PHINode>(I)) {
             // Skip the PHINodes and EH pads instructions.
-          BasicBlock *DefBlock = cast<Instruction>(E.def())->getParent();
-          if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
+            BasicBlock *DefBlock = I->getParent();
+            if (auto *CSI =
+                    dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
               InsertPt = splitBeforeCatchSwitch(CSI);
             else
               InsertPt = &*DefBlock->getFirstInsertionPt();
-        } else if (auto CSI = dyn_cast<AnyCoroSuspendInst>(CurrentValue)) {
-          // Don't spill immediately after a suspend; splitting assumes
-          // that the suspend will be followed by a branch.
-          InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHI();
           } else {
-          auto *I = cast<Instruction>(E.def());
             assert(!I->isTerminator() && "unexpected terminator");
             // For all other values, the spill is placed immediately after
             // the definition.
-          if (DT.dominates(CB, I)) {
             InsertPt = I->getNextNode();
-          } else {
-            // Unless, it is not dominated by CoroBegin, then it will be
-            // inserted immediately after CoroFrame is computed.
-            InsertPt = FramePtr->getNextNode();
           }
         }
 
         Builder.SetInsertPoint(InsertPt);
         auto *G = Builder.CreateConstInBoundsGEP2_32(
             FrameTy, FramePtr, 0, Index,
             CurrentValue->getName() + Twine(".spill.addr"));
         Builder.CreateStore(CurrentValue, G);

llvm/test/Transforms/Coroutines/coro-spill-defs-before-corobegin.ll

+; Verifies that phi and invoke definitions before CoroBegin are spilled properly.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+define i8* @f(i1 %n) "coroutine.presplit"="1" personality i32 0 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %flag = call i1 @check(i8* %alloc)
+  br i1 %flag, label %flag_true, label %flag_false
+
+flag_true:
+  br label %merge
+
+flag_false:
+  br label %merge
+
+merge:
+  %value_phi = phi i32 [ 0, %flag_true ], [ 1, %flag_false ]
+  %value_invoke = invoke i32 @calc() to label %normal unwind label %lpad
+
+normal:
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  call i32 @print(i32 %value_phi)
+  call i32 @print(i32 %value_invoke)
+  %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp1, label %suspend [i8 0, label %resume
+                                  i8 1, label %cleanup]
+resume:
+  call i32 @print(i32 %value_phi)
+  call i32 @print(i32 %value_invoke)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+
+lpad:
+  %lpval = landingpad { i8*, i32 }
+     cleanup
+
+  resume { i8*, i32 } %lpval
+}
+
+; Verifies that the both value_phi and value_invoke are stored correctly in the coroutine frame
+; CHECK: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i32, i32, i1 }
+; CHECK-LABEL: @f(
+; CHECK:       %alloc = call i8* @malloc(i32 32)
+; CHECK-NEXT:  %flag = call i1 @check(i8* %alloc)
+; CHECK-NEXT:  %value_phi = select i1 %flag, i32 0, i32 1
+; CHECK-NEXT:  %value_invoke = call i32 @calc()
+; CHECK-NEXT:  %hdl = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
+
+; CHECK:       store void (%f.Frame*)* @f.destroy, void (%f.Frame*)** %destroy.addr
+; CHECK-NEXT:  %value_invoke.spill.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 3
+; CHECK-NEXT:  store i32 %value_invoke, i32* %value_invoke.spill.addr
+; CHECK-NEXT:  %value_phi.spill.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 2
+; CHECK-NEXT:  store i32 %value_phi, i32* %value_phi.spill.addr
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare i32 @print(i32)
+declare i1 @check(i8*)
+declare i32 @calc()
+declare void @free(i8*)