[Coroutines] Part13: Handle single edge PHINodes across suspends

Summary:
If one of the uses of the value is a single edge PHINode, handle it.

Original:

    %val = something
    <suspend>
    %p = PHINode [%val]

After Spill + Part13:

    %val = something
    %slot = gep val.spill.slot
    store %val, %slot
    <suspend>
    %p = load %slot

Plus tiny fixes/changes:
   * use correct index for coro.free in CoroCleanup
   * fixup id parameter in coro.free to allow authoring coroutine in plain C with __builtins

Reviewers: majnemer

Subscribers: mehdi_amini, llvm-commits

Differential Revision: https://reviews.llvm.org/D24242

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@281020 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Transforms/Coroutines/CoroCleanup.cpp

@@ -49,7 +49,7 @@ bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
         II->replaceAllUsesWith(II->getArgOperand(1));
         break;
       case Intrinsic::coro_free:
-        II->replaceAllUsesWith(II->getArgOperand(0));
+        II->replaceAllUsesWith(II->getArgOperand(1));
         break;
       case Intrinsic::coro_alloc:
         II->replaceAllUsesWith(ConstantInt::getTrue(Context));

lib/Transforms/Coroutines/CoroEarly.cpp

@@ -115,12 +115,17 @@ static void setCannotDuplicate(CoroIdInst *CoroId) {
 
 bool Lowerer::lowerEarlyIntrinsics(Function &F) {
   bool Changed = false;
+  CoroIdInst *CoroId = nullptr;
+  SmallVector<CoroFreeInst *, 4> CoroFrees;
   for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
     Instruction &I = *IB++;
     if (auto CS = CallSite(&I)) {
       switch (CS.getIntrinsicID()) {
       default:
         continue;
+      case Intrinsic::coro_free:
+        CoroFrees.push_back(cast<CoroFreeInst>(&I));
+        break;
       case Intrinsic::coro_suspend:
         // Make sure that final suspend point is not duplicated as CoroSplit
         // pass expects that there is at most one final suspend point.
@@ -141,6 +146,7 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) {
             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
             setCannotDuplicate(CII);
             CII->setCoroutineSelf();
+            CoroId = cast<CoroIdInst>(&I);
           }
         }
         break;
@@ -160,6 +166,12 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) {
       Changed = true;
     }
   }
+  // Make sure that all CoroFree reference the coro.id intrinsic.
+  // Token type is not exposed through coroutine C/C++ builtins to plain C, so
+  // we allow specifying none and fixing it up here.
+  if (CoroId)
+    for (CoroFreeInst *CF : CoroFrees)
+      CF->setArgOperand(0, CoroId);
   return Changed;
 }
 
@@ -178,9 +190,10 @@ struct CoroEarly : public FunctionPass {
   // This pass has work to do only if we find intrinsics we are going to lower
   // in the module.
   bool doInitialization(Module &M) override {
-    if (coro::declaresIntrinsics(M, {"llvm.coro.begin", "llvm.coro.end",
-                                     "llvm.coro.resume", "llvm.coro.destroy",
-                                     "llvm.coro.done", "llvm.coro.suspend"}))
+    if (coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.destroy",
+                                     "llvm.coro.done", "llvm.coro.end",
+                                     "llvm.coro.free", "llvm.coro.promise",
+                                     "llvm.coro.resume", "llvm.coro.suspend"}))
       L = llvm::make_unique<Lowerer>(M);
     return false;
   }

lib/Transforms/Coroutines/CoroFrame.cpp

@@ -438,6 +438,17 @@ static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
       CurrentReload = CreateReload(&*CurrentBlock->getFirstInsertionPt());
     }
 
+    // If we have a single edge PHINode, remove it and replace it with a reload
+    // from the coroutine frame. (We already took care of multi edge PHINodes
+    // by rewriting them in the rewritePHIs function).
+    if (auto *PN = dyn_cast<PHINode>(E.user())) {
+      assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
+                                                "values in the PHINode");
+      PN->replaceAllUsesWith(CurrentReload);
+      PN->eraseFromParent();
+      continue;
+    }
+
     // Replace all uses of CurrentValue in the current instruction with reload.
     E.user()->replaceUsesOfWith(CurrentValue, CurrentReload);
   }

test/Transforms/Coroutines/phi-coro-end.ll

@@ -0,0 +1,48 @@
+; Verify that we correctly handle suspend when the coro.end block contains phi
+; RUN: opt < %s -O2 -enable-coroutines -S | FileCheck %s
+
+define i8* @f(i32 %n) {
+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)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %cleanup i8 1, label %cleanup]
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+
+suspend:
+  %r = phi i32 [%n, %entry], [1, %cleanup]
+  call void @llvm.coro.end(i8* %hdl, i1 false)  
+  call void @print(i32 %r)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @main
+define i32 @main() {
+entry:
+  %hdl = call i8* @f(i32 4)
+  call void @llvm.coro.resume(i8* %hdl)
+  ret i32 0
+;CHECK: call void @print(i32 4)
+;CHECK: ret i32 0
+}
+
+declare i8* @llvm.coro.alloc()
+declare i32 @llvm.coro.size.i32()
+declare i8* @llvm.coro.free(token, i8*)
+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 i8* @llvm.coro.begin(token, i8*)
+declare void @llvm.coro.end(i8*, i1) 
+
+declare noalias i8* @malloc(i32)
+declare void @print(i32)
+declare void @free(i8*)