[Coroutines] Part13: Handle single edge PHINodes across suspends

[GorNishanov]

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

[majnemer]

lib/Transforms/Coroutines/CoroEarly.cpp

169–174

Can't the FE do this?

[GorNishanov]

lib/Transforms/Coroutines/CoroEarly.cpp

169–174

Coroutine await C++ FE, sure.

This particular helper allows construct coroutines in plain C using __builtins mapping to llvm.coro.* intrinsics.
coro.id creates trouble, since it returns tokens and those do not have a good mapping in the builtin type system.

The workaround is to allow __builtin_coro_free(mem) output llvm.coro.free(token none, %mem) and then fixup in CoroEarly to point to an appropriate coro.id.

[GorNishanov]

Take out coro.free fixup code. We can do it in the frontend. Alphabetize the list of interesting coro intrinsics in coro.early

[GorNishanov]

lib/Transforms/Coroutines/CoroEarly.cpp

175–180

After I took out the fixup code, I am having second thoughts on it and want to put it back.

CoroEarly purpose is to cleanup after frontend. The fixup for coro.frees fits nicely there. Having nearly identical code that looks like LLVM pass at the end of CodeGenFunction::FinishFunction seems awkward. Having it in CoroEarly puts all of the cleanup logic in one place.

[majnemer]

lib/Transforms/Coroutines/CoroEarly.cpp

175–180

Ok, let's bring it back.

[GorNishanov]

Put coro.free fixup code back

[majnemer]

LGTM

lib/Transforms/Coroutines/CoroCleanup.cpp

for (auto IB = inst_begin(F), E = inst_end(F); IB != E;) {
     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
       switch (II->getIntrinsicID()) {
       default:
         continue;
       case Intrinsic::coro_begin:
         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));
         break;
       case Intrinsic::coro_id:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
       }

lib/Transforms/Coroutines/CoroEarly.cpp

 static void setCannotDuplicate(CoroIdInst *CoroId) {
   for (User *U : CoroId->users())
     if (auto *CB = dyn_cast<CoroBeginInst>(U))
       CB->setCannotDuplicate();
 }
 
 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.
         if (cast<CoroSuspendInst>(&I)->isFinal())
           CS.setCannotDuplicate();
         break;
       case Intrinsic::coro_end:
         // Make sure that fallthrough coro.end is not duplicated as CoroSplit
         // pass expects that there is at most one fallthrough coro.end.
         if (cast<CoroEndInst>(&I)->isFallthrough())
           CS.setCannotDuplicate();
         break;
       case Intrinsic::coro_id:
         // Mark a function that comes out of the frontend that has a coro.id
         // with a coroutine attribute.
         if (auto *CII = cast<CoroIdInst>(&I)) {
           if (CII->getInfo().isPreSplit()) {
             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
             setCannotDuplicate(CII);
             CII->setCoroutineSelf();
+            CoroId = cast<CoroIdInst>(&I);
           }
         }
         break;
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
         break;
       case Intrinsic::coro_promise:
         lowerCoroPromise(cast<CoroPromiseInst>(&I));
         break;
       case Intrinsic::coro_done:
         lowerCoroDone(cast<IntrinsicInst>(&I));
         break;
       }
       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;
 }
 
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 struct CoroEarly : public FunctionPass {
   static char ID; // Pass identification, replacement for typeid.
   CoroEarly() : FunctionPass(ID) {}
 
   std::unique_ptr<Lowerer> L;

lib/Transforms/Coroutines/CoroFrame.cpp

for (auto const &E : Spills) {
     }
 
     // If we have not seen the use block, generate a reload in it.
     if (CurrentBlock != E.userBlock()) {
       CurrentBlock = E.userBlock();
       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);
   }
 
   BasicBlock *FramePtrBB = FramePtr->getParent();
   Shape.AllocaSpillBlock =
       FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
   Shape.AllocaSpillBlock->splitBasicBlock(&Shape.AllocaSpillBlock->front(),

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

+; 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*)